#!/usr/bin/env python3
from pathlib import Path
from pprint import pformat
from re import compile as regex
from typing import Optional, Dict, List, Set, Any, Union

import click
from rich import print
from rich.text import Text
from rich.tree import Tree
from yaml import safe_load


def print_help() -> None:
    ctx = click.get_current_context()
    click.echo(ctx.get_help())
    ctx.exit()


class Role:
    def __init__(self, role_dir: Path, root_dir: Path) -> None:
        self.role_name = str(role_dir.relative_to(root_dir))
        self.role_dir = role_dir.resolve()
        self.dependencies: Dict[str, List[str]] = self.scan_dependencies()
        # TODO: find tags

    @property
    def dependency_file(self) -> Path:
        return self.role_dir / "meta" / "main.yaml"

    def scan_dependencies(self) -> List[str]:
        result = {}
        if self.dependency_file.is_file():
            meta = safe_load(self.dependency_file.open())
            if (
                meta is not None
                and "dependencies" in meta
                and meta["dependencies"] is not None
            ):
                for role in meta["dependencies"]:
                    if type(role) == str:
                        result[role] = []
                    else:
                        result[role["role"]] = role["tags"] if "tags" in role else []
        return result

    def tree(
        self, roles: Dict[str, "Role"], print_tags: bool = False, tags: List[str] = []
    ):
        root = Tree(
            f'{self.role_name}{" " + str(tags) if tags and print_tags else ""}',
            style="green" if not tags else "bold blue",
        )
        for dependency, tags in self.dependencies.items():
            if dependency not in roles:
                root.add(
                    Text(
                        f'{dependency}{tags if tags and print_tags else ""}',
                        style="bold red",
                    )
                )
            else:
                root.add(
                    roles[dependency].tree(roles, print_tags=print_tags, tags=tags)
                )
        return root


def recurse(
    element: Union[Dict[str, Any], List[Any]],
    collection: Optional[Dict[str, List[str]]] = None,
) -> None:
    collection = collection if collection is not None else {}
    if type(element) is dict:
        for key, value in element.items():
            if key == "include_role":
                collection[value["name"]] = element["tags"] if "tags" in element else []
            if key == "roles":
                for role in value:
                    if type(role) == str:
                        collection[role] = []
                    else:
                        collection[role["role"]] = (
                            role["tags"] if "tags" in role else []
                        )
            else:
                recurse(value, collection)
    elif type(element) is list:
        for value in element:
            recurse(value, collection)
    return collection


def get_roles_from_playbook(playbook: Path) -> Dict[str, List[str]]:
    play = safe_load(playbook.open())
    return recurse(play)


class DependencyTree:
    def __init__(self, root_dir: Path, playbook: Optional[Path] = None) -> None:
        self.root_dir: Path = root_dir
        self.playbook = playbook
        self.playbook_roles = (
            get_roles_from_playbook(playbook) if self.playbook else None
        )
        self.roles: Dict[str, Role] = {}
        self.scan()

    @staticmethod
    def _is_role(candidate: Path) -> bool:
        if not candidate.is_dir():
            return False
        yaml_files = set(str(c) for c in candidate.rglob("main.yaml"))
        return (
            str(candidate / "tasks" / "main.yaml") in yaml_files
            or str(candidate / "meta" / "main.yaml") in yaml_files
        )

    def scan(self) -> None:
        for path in self.root_dir.glob("**/*"):
            if self._is_role(path):
                self.roles[str(path.relative_to(self.root_dir))] = Role(
                    path, self.root_dir
                )

    def print(
        self,
        root_text: str = ".",
        filter_roles: Optional[Set[str]] = None,
        tags: bool = False,
    ) -> None:
        root = Tree(Text(root_text))
        for name, role in self.roles.items():
            if (not filter_roles or role.role_name in filter_roles) and (
                not self.playbook or role.role_name in self.playbook_roles
            ):
                root.add(
                    role.tree(
                        self.roles,
                        print_tags=tags,
                        tags=self.playbook_roles[name]
                        if self.playbook_roles and name in self.playbook_roles
                        else [],
                    )
                )
        print(root)


def generate_roles_tree(
    root_dir: Path, playbook: Optional[Path] = None
) -> DependencyTree:
    return DependencyTree(root_dir, playbook)


ROLES_PATH_PAT = regex(r"^roles_path\s*\=\s*(.+)$")


def get_roles_path_from_ansible_cfg(ansible_cfg: Path) -> Path:
    for line in ansible_cfg.open().readlines():
        if match := ROLES_PATH_PAT.match(line):
            return Path(match.group(1))


@click.group()
def cli():
    pass


@cli.command()
@click.option(
    "-r",
    "--role",
    default=None,
    help="Name of role for which to generate a dependency tree for",
)
@click.option(
    "-d",
    "--role-dir",
    default=None,
    help="Directory to search for roles in (generates dependency tree for all roles if no playbook [-p] or specific role [-r] is provided)",
)
@click.option(
    "-p",
    "--playbook",
    default=None,
    help="Path to a playbook YAML file from which a dependency tree will be generated",
)
@click.option(
    "-f",
    "--find",
    default=None,
    help="Find a particular role in the dependency tree and highlight the path to it (not implemented)",
)
@click.option("--tags", is_flag=True, help="Display tags in the tree as well")
def roles(
    role: Optional[str],
    role_dir: Optional[str],
    playbook: Optional[str],
    find: Optional[str],
    tags: bool = False,
) -> None:
    # TODO: find
    if role_dir is None:
        if (Path.cwd() / "ansible.cfg").is_file():
            role_dir = get_roles_path_from_ansible_cfg(Path.cwd() / "ansible.cfg")
        elif (Path.cwd() / "roles").is_dir():
            role_dir = Path.cwd() / "roles"
        else:
            role_dir = Path.cwd()
    elif Path(role_dir).is_dir():
        role_dir = Path(role_dir)
    else:
        print(f"Role directory: {role_dir} is not accessible")
        sys.exit(1)
    if playbook is not None:
        if not Path(playbook).is_file():
            print(f"Playbook file: {playbook} is not accessible.")
            sys.exit(1)
        tree = generate_roles_tree(role_dir, Path(playbook))
    else:
        tree = generate_roles_tree(role_dir)
    tree.print(
        root_text=str(role_dir), filter_roles={role} if role else None, tags=tags
    )


if __name__ == "__main__":
    cli()