import enum
import time
from collections import defaultdict
from contextlib import contextmanager
from typing import Callable
from typing import Dict
from typing import List
from typing import Optional
from typing import TYPE_CHECKING, Iterable
from typing import Tuple
from typing import Union

from pkm_buildsys.packages.package import Package
from pkm_main.mixology import resolve_version
from pkm_main.mixology.failure import SolveFailure
from pkm_main.packages import DependencyPackage
from pkm_main.repositories import Repository
from pkm_main.utils.env import Env

from .exceptions import OverrideNeeded
from .exceptions import SolverProblemError
from .provider import Provider
from ..console import console, Printer

if TYPE_CHECKING:
    from pkm_buildsys.packages.dependency import Dependency
    from pkm_buildsys.packages.directory_dependency import DirectoryDependency
    from pkm_buildsys.packages.file_dependency import FileDependency
    from pkm_buildsys.packages.url_dependency import URLDependency
    from pkm_buildsys.packages.vcs_dependency import VCSDependency
    from pkm_main.managed_project import ManagedProject

    from .transaction import Transaction


class Solver:
    def __init__(
            self,
            project: "ManagedProject",
            installed: Repository,
            locked: Repository,
            package: Optional[Package] = None,
            provider: Optional[Provider] = None,
            printer: Optional[Printer] = None
    ):
        self._package = package or project.package
        self._pool = project.pool
        self._installed = installed
        self._locked = locked
        self._printer = printer or console

        if provider is None:
            provider = Provider(project, package=self._package)

        self._provider = provider
        self._overrides = []

    @property
    def provider(self) -> Provider:
        return self._provider

    @contextmanager
    def use_environment(self, env: Env) -> None:
        with self.provider.use_environment(env):
            yield

    def solve(self, use_latest: List[str] = None) -> "Transaction":
        from .transaction import Transaction

        self._printer.println("<info>Resolving dependencies...</info>")
        start = time.time()
        packages, depths = self._solve(use_latest=use_latest)
        end = time.time()

        self._printer.println(f"<success>Dependency Resolution done within {end - start:.2f} seconds")

        if len(self._overrides) > 1:
            self._provider.debug(
                "Complete version solving took {:.3f} seconds with {} overrides".format(
                    end - start, len(self._overrides)
                )
            )
            self._provider.debug(
                f"Resolved with overrides: {', '.join(f'({b})' for b in self._overrides)}"
            )

        return Transaction(
            self._locked.packages,
            list(zip(packages, depths)),
            installed_packages=self._installed.packages,
            root_package=self._package,
        )

    def solve_in_compatibility_mode(
            self, overrides: Tuple[Dict], use_latest: List[str] = None
    ) -> Tuple[List["Package"], List[int]]:
        locked = {}
        for package in self._locked.packages:
            locked[package.name] = DependencyPackage(package.to_dependency(), package)

        packages = []
        depths = []
        for override in overrides:
            self._provider.debug(
                "<comment>Retrying dependency resolution "
                f"with the following overrides ({override}).</comment>"
            )
            self._provider.set_overrides(override)
            _packages, _depths = self._solve(use_latest=use_latest)
            for index, package in enumerate(_packages):
                if package not in packages:
                    packages.append(package)
                    depths.append(_depths[index])
                    continue
                else:
                    idx = packages.index(package)
                    pkg = packages[idx]
                    depths[idx] = max(depths[idx], _depths[index])

                    for dep in package.requires:
                        if dep not in pkg.requires:
                            pkg.add_dependency(dep)

        return packages, depths

    def _solve(self, use_latest: List[str] = None) -> Tuple[List[Package], List[int]]:
        if self._provider._overrides:
            self._overrides.append(self._provider._overrides)

        locked = {}
        for package in self._locked.packages:
            try:
                locked[package.name] = DependencyPackage(package.to_dependency(), package)
            except FileNotFoundError:
                # if in lock there is a dependency inside the filesystem that cannot be resolved,
                # it just means that we need to re-resolve its version
                pass

        try:
            result = resolve_version(
                self._package, self._provider, locked=locked, use_latest=use_latest
            )

            packages = result.packages
        except OverrideNeeded as e:
            return self.solve_in_compatibility_mode(e.overrides, use_latest=use_latest)
        except SolveFailure as e:
            raise SolverProblemError(e)

        # NOTE passing explicit empty array for seen to reset between invocations during update + install cycle
        results = dict(
            depth_first_search(
                PackageNode(self._package, packages, seen=[]), aggregate_package_nodes
            )
        )

        # Merging feature packages with base packages
        final_packages = []
        depths = []
        for package in packages:
            if package.features:
                for _package in packages:
                    if (
                            _package.name == package.name
                            and not _package.is_same_package_as(package)
                            and _package.version == package.version
                    ):
                        for dep in package.requires:
                            if dep.is_same_package_as(_package):
                                continue

                            if dep not in _package.requires:
                                _package.add_dependency(dep)

                continue

            final_packages.append(package)
            depths.append(results[package])

        # Return the packages in their original order with associated depths
        return final_packages, depths


_NODE_ID_T = Union[str, Tuple[str, bool]]


class DFSNode:
    def __init__(
            self, id: _NODE_ID_T, name: str, base_name: str
    ) -> None:
        self.id = id
        self.name = name
        self.base_name = base_name

    def reachable(self) -> List:
        return []

    def visit(self, parents: List["PackageNode"]) -> None:
        pass

    def __str__(self) -> str:
        return str(self.id)


class VisitedState(enum.Enum):
    Unvisited = 0
    PartiallyVisited = 1
    Visited = 2


def depth_first_search(
        source: "PackageNode", aggregator: Callable
) -> List[Tuple[Package, int]]:
    back_edges = defaultdict(list)
    visited = {}
    topo_sorted_nodes = []

    dfs_visit(source, back_edges, visited, topo_sorted_nodes)

    # Combine the nodes by name
    combined_nodes = defaultdict(list)
    name_children = defaultdict(list)
    for node in topo_sorted_nodes:
        node.visit(back_edges[node.id])
        name_children[node.name].extend(node.reachable())
        combined_nodes[node.name].append(node)

    combined_topo_sorted_nodes = []
    for node in topo_sorted_nodes:
        if node.name in combined_nodes:
            combined_topo_sorted_nodes.append(combined_nodes.pop(node.name))

    results = [
        aggregator(nodes, name_children[nodes[0].name])
        for nodes in combined_topo_sorted_nodes
    ]
    return results


def dfs_visit(
        node: "PackageNode",
        back_edges: Dict[_NODE_ID_T, List["PackageNode"]],
        visited: Dict[_NODE_ID_T, VisitedState],
        sorted_nodes: List["PackageNode"],
) -> bool:
    if visited.get(node.id, VisitedState.Unvisited) == VisitedState.Visited:
        return True
    if visited.get(node.id, VisitedState.Unvisited) == VisitedState.PartiallyVisited:
        # We have a circular dependency.
        # Since the dependencies are resolved we can
        # simply skip it because we already have it
        return True

    visited[node.id] = VisitedState.PartiallyVisited
    for neighbor in node.reachable():
        back_edges[neighbor.id].append(node)
        if not dfs_visit(neighbor, back_edges, visited, sorted_nodes):
            return False
    visited[node.id] = VisitedState.Visited
    sorted_nodes.insert(0, node)
    return True


class PackageNode(DFSNode):
    def __init__(
            self,
            package: Package,
            packages: List[Package],
            seen: List[Package],
            previous: Optional["PackageNode"] = None,
            previous_dep: Optional[
                Union[
                    "DirectoryDependency",
                    "FileDependency",
                    "URLDependency",
                    "VCSDependency",
                    "Dependency",
                ]
            ] = None,
            dep: Optional[
                Union[
                    "DirectoryDependency",
                    "FileDependency",
                    "URLDependency",
                    "VCSDependency",
                    "Dependency",
                ]
            ] = None,
    ) -> None:
        self.package = package
        self.packages = packages
        self.seen = seen

        self.previous = previous
        self.previous_dep = previous_dep
        self.dep = dep
        self.depth = -1

        if not previous:
            self.optional = True
        else:
            self.optional = dep.is_optional()

        super().__init__(
            (package.complete_name, self.optional),
            package.complete_name,
            package.name,
        )

    def reachable(self) -> List["PackageNode"]:
        children: List[PackageNode] = []

        # skip already traversed packages
        if self.package in self.seen:
            return []
        else:
            self.seen.append(self.package)

        if (
                self.previous_dep
                and self.previous_dep is not self.dep
                and self.previous_dep.name == self.dep.name
        ):
            return []

        for dependency in self.package.all_requires:
            if self.previous and self.previous.name == dependency.name:
                # We have a circular dependency.
                # Since the dependencies are resolved we can
                # simply skip it because we already have it
                # N.B. this only catches cycles of length 2;
                # dependency cycles in general are handled by the DFS traversal
                continue

            for pkg in self.packages:
                if pkg.complete_name == dependency.complete_name and (
                        dependency.constraint.allows(pkg.version)
                        or dependency.allows_prereleases()
                        and pkg.version.is_unstable()
                        and dependency.constraint.allows(pkg.version.stable)
                ):
                    # If there is already a child with this name
                    # we merge the requirements
                    if any(
                            child.package.name == pkg.name
                            for child in children
                    ):
                        continue

                    children.append(
                        PackageNode(
                            pkg,
                            self.packages,
                            self.seen,
                            self,
                            dependency,
                            self.dep or dependency,
                        )
                    )

        return children

    def visit(self, parents: Iterable["PackageNode"]) -> None:
        # The root package, which has no parents, is defined as having depth -1
        # So that the root package's top-level dependencies have depth 0.
        self.depth = 1 + max(
            [
                parent.depth if parent.base_name != self.base_name else parent.depth - 1
                for parent in parents
            ]
            + [-2]
        )


def aggregate_package_nodes(
        nodes: List[PackageNode], children: List[PackageNode]
) -> Tuple[Package, int]:
    package = nodes[0].package
    depth = max(node.depth for node in nodes)

    optional = all(node.optional for node in children + nodes)
    for node in nodes:
        node.depth = depth
        node.optional = optional

    package.optional = optional

    return package, depth
