from importlib.metadata import distributions, PackageNotFoundError
from typing import Set, Optional, Dict, Tuple, List, Any, Union, Pattern
import re
import json
import yaml
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from functools import lru_cache
from enum import Enum
from dataclasses import dataclass
from pathlib import Path
import time


class DependencyType(Enum):
    REQUIRED = "required"
    OPTIONAL = "optional"
    DEVELOPMENT = "development"


class OutputFormat(Enum):
    TEXT = "text"
    JSON = "json"
    YAML = "yaml"
    DICT = "dict"


@dataclass
class PackageInfo:
    name: str
    version: str
    requirement: str = ""
    status: str = "installed"
    dependencies: List["PackageInfo"] = None
    dep_type: DependencyType = DependencyType.REQUIRED
    extras: List[str] = None

    def __post_init__(self):
        if self.dependencies is None:
            self.dependencies = []
        if self.extras is None:
            self.extras = []


class PackageCache:
    """Cache for package metadata to improve performance"""

    _instance = None
    _cache = {}
    _lock = threading.RLock()

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

    @classmethod
    def get_distributions(cls, pkg_name: str):
        with cls._lock:
            if pkg_name in cls._cache:
                return cls._cache[pkg_name]

            try:
                dist = distributions(pkg_name)
                cls._cache[pkg_name] = dist
                return dist
            except PackageNotFoundError:
                cls._cache[pkg_name] = None
                return None
            except Exception:
                cls._cache[pkg_name] = None
                return None

    @classmethod
    def clear_cache(cls):
        with cls._lock:
            cls._cache.clear()


def parse_requirement(req: str) -> Tuple[Optional[str], str, List[str], str]:
    """
    Parse a requirement string into (package name, version spec, extras, marker).

    Args:
        req: requirement string (like "package[extra1,extra2]>=1.0; python_version>'3.6'")

    Returns:
        Tuple: (name, version_spec, extras, marker) or (None, "", [], "") if invalid
    """
    try:
        # Pattern to match complex requirements with extras and markers
        pattern = r"""
            ^\s*
            ([a-zA-Z0-9][a-zA-Z0-9._-]*)      # Package name
            (?:\[([^\]]+)\])?                 # Extras in brackets
            \s*
            ([^;]*)                           # Version specifiers (before ;)
            (?:;\s*(.*))?                     # Environment marker (after ;)
            \s*$
        """

        match = re.match(pattern, req, re.VERBOSE)

        if match:
            name = match.group(1)
            extras_str = match.group(2) or ""
            version_spec = (match.group(3) or "").strip()
            marker = (match.group(4) or "").strip()

            # Parse extras
            extras = (
                [extra.strip() for extra in extras_str.split(",")] if extras_str else []
            )

            # Validate package name according to PEP 508
            if not re.match(
                r"^([A-Z0-9]|[A-Z0-9][A-Z0-9._-]*[A-Z0-9])$", name, re.IGNORECASE
            ):
                return None, "", [], ""

            return name, version_spec, extras, marker
        else:
            # Fallback to simple parsing
            simple_pattern = (
                r"^\s*([a-zA-Z0-9][a-zA-Z0-9._-]*)(?:\s*([^\s;]*))?(?:;\s*(.*))?\s*$"
            )
            simple_match = re.match(simple_pattern, req)
            if simple_match:
                name = simple_match.group(1)
                version_spec = (simple_match.group(2) or "").strip()
                marker = (simple_match.group(3) or "").strip()
                return name, version_spec, [], marker

            return None, "", [], ""

    except Exception:
        return None, "", [], ""


def should_include_requirement(
    version_spec: str, marker: str, skip_optional: bool = False
) -> bool:
    """
    Determine if a requirement should be included based on markers and optional flags.
    """
    if skip_optional and ("extra" in marker.lower() or "optional" in marker.lower()):
        return False

    # TODO: Implement environment marker evaluation
    # For now, we include all requirements and let the user filter
    return True


def print_dep_tree(
    pkg: str,
    indent: int = 0,
    seen: Optional[Set[str]] = None,
    max_depth: Optional[int] = 1,
    current_depth: int = 0,
    show_versions: bool = True,
    show_required: bool = True,
    show_installed: bool = True,
    verbose: bool = False,
    pattern_filter: Optional[Pattern] = None,
    skip_optional: bool = False,
    parallel_processing: int = 1,
    show_extras: bool = False,
    show_markers: bool = False,
) -> bool:
    """
    Print the dependency tree of a package with enhanced features.

    Args:
        pkg: package name to analyze
        indent: indentation for tree printing
        seen: set of already visited packages
        max_depth: maximum depth of the tree
        current_depth: current recursion depth
        show_versions: show installed versions
        show_required: show required version specifiers
        show_installed: show installed versions of dependencies
        verbose: show detailed errors
        pattern_filter: regex pattern to filter packages
        skip_optional: skip optional dependencies
        parallel_processing: number of parallel threads for processing
        show_extras: show extras in requirements
        show_markers: show environment markers

    Returns:
        bool: True if successful, False if there were errors
    """

    if max_depth is not None and current_depth > max_depth:
        return True

    if seen is None:
        seen = set()

    normalized_pkg = pkg.lower().replace("_", "-")

    # Apply pattern filter
    if pattern_filter and not pattern_filter.search(pkg):
        return True

    if indent:
        indent_lvl = " " * indent
    else:
        indent_lvl = ""

    # Prevent cycles with depth tracking
    cycle_key = f"{normalized_pkg}:{current_depth}"
    if cycle_key in seen:
        print(indent_lvl + f"{pkg} (cycle detected at depth {current_depth})")
        return True
    seen.add(cycle_key)

    try:
        dist = PackageCache.get_distributions(pkg)
        if dist is None:
            print(indent_lvl + f"{pkg} [not installed]")
            return False

        name = dist.metadata["Name"]
        installed_version = dist.version
    except Exception as e:
        if verbose:
            print(indent_lvl + f"{pkg} [error reading metadata: {str(e)}]")
        else:
            print(indent_lvl + f"{pkg} [error reading metadata]")
        return False

    version_info = f"=={installed_version}" if show_versions else ""
    print(f"{indent_lvl}{name}{version_info}")

    try:
        requires = dist.requires or []
    except Exception as e:
        if verbose:
            print(indent_lvl + f"  [error reading requirements: {str(e)}]")
        return True

    success = True
    requirements_to_process = []

    # First pass: parse all requirements
    for req in requires:
        child_name, req_spec, extras, marker = parse_requirement(req)
        if not child_name:
            if verbose:
                print(indent_lvl + f"  [invalid requirement: {req}]")
            continue

        if not should_include_requirement(req_spec, marker, skip_optional):
            continue

        requirements_to_process.append((child_name, req_spec, extras, marker))

    # Process requirements in parallel or sequentially
    if parallel_processing > 1 and len(requirements_to_process) > 1:
        results = []
        with ThreadPoolExecutor(max_workers=parallel_processing) as executor:
            future_to_req = {
                executor.submit(
                    process_single_requirement,
                    child_name,
                    req_spec,
                    extras,
                    marker,
                    indent_lvl,
                    i,
                    len(requirements_to_process),
                    show_required,
                    show_installed,
                    show_extras,
                    show_markers,
                    verbose,
                ): (child_name, req_spec, extras, marker)
                for i, (child_name, req_spec, extras, marker) in enumerate(
                    requirements_to_process
                )
            }

            for future in as_completed(future_to_req):
                result = future.result()
                results.append((future_to_req[future], result))

        # Sort results by original order
        results.sort(key=lambda x: requirements_to_process.index(x[0]))
        for (child_name, req_spec, extras, marker), child_success in results:
            if not process_child_dependency(
                child_name,
                req_spec,
                extras,
                marker,
                indent,
                indent_lvl,
                requirements_to_process.index((child_name, req_spec, extras, marker)),
                len(requirements_to_process),
                seen,
                max_depth,
                current_depth,
                show_versions,
                show_required,
                show_installed,
                verbose,
                pattern_filter,
                skip_optional,
                parallel_processing,
                show_extras,
                show_markers,
            ):
                success = False
    else:
        # Sequential processing
        for i, (child_name, req_spec, extras, marker) in enumerate(
            requirements_to_process
        ):
            if not process_child_dependency(
                child_name,
                req_spec,
                extras,
                marker,
                indent,
                indent_lvl,
                i,
                len(requirements_to_process),
                seen,
                max_depth,
                current_depth,
                show_versions,
                show_required,
                show_installed,
                verbose,
                pattern_filter,
                skip_optional,
                parallel_processing,
                show_extras,
                show_markers,
            ):
                success = False

    # Clean up cycle tracking
    seen.remove(cycle_key)
    return success


def process_single_requirement(
    child_name,
    req_spec,
    extras,
    marker,
    indent_lvl,
    index,
    total_reqs,
    show_required,
    show_installed,
    show_extras,
    show_markers,
    verbose,
):
    """Process a single requirement for parallel execution"""
    prefix = "├── " if index < total_reqs - 1 else "└── "

    installed_info = ""
    if show_installed:
        try:
            child_dist = PackageCache.get_distributions(child_name)
            if child_dist:
                installed_info = f", installed: {child_dist.version}"
            else:
                installed_info = ", not installed"
        except Exception:
            installed_info = ", unknown"

    # Build requirement info
    req_parts = []
    if show_required and req_spec:
        req_parts.append(f"required: {req_spec}")
    if show_extras and extras:
        req_parts.append(f"extras: [{', '.join(extras)}]")
    if show_markers and marker:
        req_parts.append(f"marker: {marker}")
    if show_installed and installed_info:
        req_parts.append(installed_info.lstrip(", "))

    req_info = f" [{', '.join(req_parts)}]" if req_parts else ""
    print(f"{indent_lvl}{prefix}{child_name}{req_info}")

    return True


def process_child_dependency(
    child_name,
    req_spec,
    extras,
    marker,
    indent,
    indent_lvl,
    index,
    total_reqs,
    seen,
    max_depth,
    current_depth,
    show_versions,
    show_required,
    show_installed,
    verbose,
    pattern_filter,
    skip_optional,
    parallel_processing,
    show_extras,
    show_markers,
):
    """Process a child dependency with proper indentation"""
    new_indent = indent + 4
    new_indent_lvl = indent_lvl + ("│   " if index < total_reqs - 1 else "    ")

    child_success = print_dep_tree(
        child_name,
        new_indent,
        seen,
        max_depth,
        current_depth + 1,
        show_versions,
        show_required,
        show_installed,
        verbose,
        pattern_filter,
        skip_optional,
        parallel_processing,
        show_extras,
        show_markers,
    )
    return child_success


def get_dep_tree(
    pkg: str,
    max_depth: Optional[int] = None,
    output_format: OutputFormat = OutputFormat.DICT,
    skip_optional: bool = False,
    include_stats: bool = False,
    group_by_type: bool = False,
) -> Union[Dict, str]:
    """
    Return the dependency tree in various formats with enhanced features.

    Args:
        pkg: package name
        max_depth: maximum recursion depth
        output_format: output format (TEXT, JSON, YAML, DICT)
        skip_optional: skip optional dependencies
        include_stats: include statistics about the tree
        group_by_type: group dependencies by type

    Returns:
        Dependency tree in the requested format
    """

    def build_tree(
        pkg_name, current_depth=0, seen=None, dep_type=DependencyType.REQUIRED
    ):
        if seen is None:
            seen = set()

        if max_depth is not None and current_depth > max_depth:
            return None

        normalized_pkg = pkg_name.lower().replace("_", "-")
        cycle_key = f"{normalized_pkg}:{current_depth}"
        if cycle_key in seen:
            return {
                "name": pkg_name,
                "status": "cycle_detected",
                "depth": current_depth,
                "type": dep_type.value,
            }

        seen.add(cycle_key)

        try:
            dist = PackageCache.get_distributions(pkg_name)
            if dist is None:
                return {
                    "name": pkg_name,
                    "status": "not_installed",
                    "type": dep_type.value,
                }

            node = {
                "name": dist.metadata["Name"],
                "version": dist.version,
                "status": "installed",
                "type": dep_type.value,
                "dependencies": [],
            }

            requires = dist.requires or []
            for req in requires:
                child_name, req_spec, extras, marker = parse_requirement(req)
                if not child_name:
                    continue

                if skip_optional and not should_include_requirement(
                    req_spec, marker, skip_optional
                ):
                    continue

                # Determine dependency type
                child_dep_type = DependencyType.REQUIRED
                if "extra" in marker.lower() or "optional" in marker.lower():
                    child_dep_type = DependencyType.OPTIONAL
                elif "dev" in marker.lower() or "test" in marker.lower():
                    child_dep_type = DependencyType.DEVELOPMENT

                child_node = build_tree(
                    child_name, current_depth + 1, seen, child_dep_type
                )
                if child_node:
                    child_node.update(
                        {"requirement": req_spec, "extras": extras, "marker": marker}
                    )
                    if group_by_type:
                        # Group by type instead of flat list
                        if "dependencies_by_type" not in node:
                            node["dependencies_by_type"] = {
                                t.value: [] for t in DependencyType
                            }
                        node["dependencies_by_type"][child_dep_type.value].append(
                            child_node
                        )
                    else:
                        node["dependencies"].append(child_node)

            # Add statistics if requested
            if include_stats:
                node["statistics"] = calculate_statistics(node)

            seen.remove(cycle_key)
            return node

        except Exception:
            return {"name": pkg_name, "status": "error", "type": dep_type.value}

    def calculate_statistics(node):
        """Calculate statistics for the dependency tree"""
        stats = {
            "total_dependencies": 0,
            "max_depth": 0,
            "by_status": {},
            "by_type": {t.value: 0 for t in DependencyType},
        }

        def traverse_for_stats(subnode, depth=0):
            stats["total_dependencies"] += 1
            stats["max_depth"] = max(stats["max_depth"], depth)
            stats["by_status"][subnode.get("status", "unknown")] = (
                stats["by_status"].get(subnode.get("status", "unknown"), 0) + 1
            )
            stats["by_type"][subnode.get("type", DependencyType.REQUIRED.value)] += 1

            for dep in subnode.get("dependencies", []):
                traverse_for_stats(dep, depth + 1)

            if "dependencies_by_type" in subnode:
                for deps_list in subnode["dependencies_by_type"].values():
                    for dep in deps_list:
                        traverse_for_stats(dep, depth + 1)

        traverse_for_stats(node)
        return stats

    tree = build_tree(pkg)

    if output_format == OutputFormat.DICT:
        return tree
    elif output_format == OutputFormat.JSON:
        return json.dumps(tree, indent=2, default=str)
    elif output_format == OutputFormat.YAML:
        try:
            return yaml.dump(tree, default_flow_style=False, allow_unicode=True)
        except ImportError:
            raise ImportError("PyYAML is required for YAML output")
    else:
        return format_tree_text(tree)


def format_tree_text(tree: Dict, indent: int = 0) -> str:
    """Format dependency tree as text"""
    if tree is None:
        return ""

    indent_str = " " * indent
    result = []

    name = tree.get("name", "")
    version = tree.get("version", "")
    status = tree.get("status", "")
    dep_type = tree.get("type", "")

    line = f"{indent_str}{name}"
    if version:
        line += f"=={version}"
    if status != "installed":
        line += f" [{status}]"
    if dep_type != DependencyType.REQUIRED.value:
        line += f" ({dep_type})"

    result.append(line)

    # Handle both flat and grouped dependencies
    dependencies = []
    if "dependencies" in tree:
        dependencies.extend(tree["dependencies"])
    if "dependencies_by_type" in tree:
        for deps_list in tree["dependencies_by_type"].values():
            dependencies.extend(deps_list)

    for i, dep in enumerate(dependencies):
        prefix = "├── " if i < len(dependencies) - 1 else "└── "
        child_text = format_tree_text(dep, indent + 4)
        if child_text:
            result.append(f"{indent_str}{prefix}{child_text}")

    return "\n".join(result)


def export_to_graphviz(tree: Dict, output_file: str = None) -> str:
    """
    Export dependency tree to Graphviz DOT format.

    Args:
        tree: Dependency tree dictionary
        output_file: Optional file to save the output

    Returns:
        Graphviz DOT string
    """
    dot_lines = ["digraph Dependencies {", "    rankdir=LR;", "    node [shape=box];"]
    node_ids = {}

    def add_nodes_and_edges(node, parent_id=None):
        if node is None:
            return

        node_name = node.get("name", "")
        node_version = node.get("version", "")
        node_status = node.get("status", "")

        # Create unique node ID
        if node_name not in node_ids:
            node_id = f"node{len(node_ids)}"
            node_ids[node_name] = node_id

            # Node label with optional version and status
            label = node_name
            if node_version:
                label += f"\\n{node_version}"
            if node_status != "installed":
                label += f"\\n[{node_status}]"

            # Node style based on status
            style = "filled"
            fillcolor = "white"
            if node_status == "not_installed":
                fillcolor = "lightgray"
            elif node_status == "error":
                fillcolor = "lightcoral"
            elif node_status == "cycle_detected":
                fillcolor = "lightyellow"

            dot_lines.append(
                f'    {node_id} [label="{label}", style="{style}", fillcolor="{fillcolor}"];'
            )
        else:
            node_id = node_ids[node_name]

        # Add edge from parent
        if parent_id:
            dot_lines.append(f"    {parent_id} -> {node_id};")

        # Process dependencies
        dependencies = []
        if "dependencies" in node:
            dependencies.extend(node["dependencies"])
        if "dependencies_by_type" in node:
            for deps_list in node["dependencies_by_type"].values():
                dependencies.extend(deps_list)

        for dep in dependencies:
            add_nodes_and_edges(dep, node_id)

    add_nodes_and_edges(tree)
    dot_lines.append("}")

    dot_output = "\n".join(dot_lines)

    if output_file:
        with open(output_file, "w") as f:
            f.write(dot_output)

    return dot_output


def export_to_mermaid(tree: Dict, output_file: str = None) -> str:
    """
    Export dependency tree to Mermaid format.

    Args:
        tree: Dependency tree dictionary
        output_file: Optional file to save the output

    Returns:
        Mermaid string
    """
    mermaid_lines = ["graph LR"]
    node_ids = {}

    def add_mermaid_nodes(node, parent_id=None):
        if node is None:
            return

        node_name = node.get("name", "")
        node_version = node.get("version", "")
        node_status = node.get("status", "")

        # Create unique node ID
        if node_name not in node_ids:
            node_id = f"n{len(node_ids)}"
            node_ids[node_name] = node_id

            # Node text with optional version
            node_text = node_name
            if node_version:
                node_text += f" ({node_version})"

            # Node style based on status
            style = ""
            if node_status == "not_installed":
                style = ":::not_installed"
            elif node_status == "error":
                style = ":::error"
            elif node_status == "cycle_detected":
                style = ":::cycle"

            mermaid_lines.append(f'    {node_id}["{node_text}"]{style}')

        # Add edge from parent
        if parent_id:
            mermaid_lines.append(f"    {parent_id} --> {node_ids[node_name]}")

        # Process dependencies
        dependencies = []
        if "dependencies" in node:
            dependencies.extend(node["dependencies"])
        if "dependencies_by_type" in node:
            for deps_list in node["dependencies_by_type"].values():
                dependencies.extend(deps_list)

        for dep in dependencies:
            add_mermaid_nodes(dep, node_ids[node_name])

    add_mermaid_nodes(tree)

    # Add CSS classes for styling
    mermaid_lines.extend(
        [
            "    classDef default fill:#fff,stroke:#333,stroke-width:2px;",
            "    classDef not_installed fill:#f0f0f0,stroke:#999,stroke-width:1px;",
            "    classDef error fill:#ffcccc,stroke:#ff0000,stroke-width:2px;",
            "    classDef cycle fill:#ffffcc,stroke:#ffcc00,stroke-width:2px;",
        ]
    )

    mermaid_output = "\n".join(mermaid_lines)

    if output_file:
        with open(output_file, "w") as f:
            f.write(mermaid_output)

    return mermaid_output


def conflicts(tree: Dict) -> Dict:
    """
    Analyze version conflicts in the dependency tree.

    Args:
        tree: Dependency tree dictionary

    Returns:
        Dictionary with conflict information
    """
    conflicts = {}
    package_requirements = {}

    def collect_requirements(node, path=""):
        if node is None:
            return

        node_name = node.get("name", "")
        node_version = node.get("version", "")
        requirement = node.get("requirement", "")

        if node_name and path:
            if node_name not in package_requirements:
                package_requirements[node_name] = []
            package_requirements[node_name].append(
                {
                    "required_by": path,
                    "requirement": requirement,
                    "installed_version": node_version,
                }
            )

        current_path = f"{path}->{node_name}" if path else node_name

        # Process dependencies
        dependencies = []
        if "dependencies" in node:
            dependencies.extend(node["dependencies"])
        if "dependencies_by_type" in node:
            for deps_list in node["dependencies_by_type"].values():
                dependencies.extend(deps_list)

        for dep in dependencies:
            collect_requirements(dep, current_path)

    collect_requirements(tree)

    # Analyze conflicts
    for package, reqs in package_requirements.items():
        if len(reqs) > 1:
            versions = set()
            for req in reqs:
                if req["installed_version"]:
                    versions.add(req["installed_version"])

            if len(versions) > 1:
                conflicts[package] = {
                    "different_versions": list(versions),
                    "requirements": reqs,
                }

    return conflicts


def print_tree(
    pkg: str,
    max_depth: Optional[int] = 1,
    output_format: OutputFormat = OutputFormat.TEXT,
    pattern_filter: str = None,
    skip_optional: bool = False,
    parallel_processing: int = 1,
    show_extras: bool = False,
    show_markers: bool = False,
) -> bool:
    """
    Enhanced helper to print the dependency tree with various options.
    """
    if output_format == OutputFormat.TEXT:
        pattern = re.compile(pattern_filter) if pattern_filter else None
        success = print_dep_tree(
            pkg,
            max_depth=max_depth,
            pattern_filter=pattern,
            skip_optional=skip_optional,
            parallel_processing=parallel_processing,
            show_extras=show_extras,
            show_markers=show_markers,
        )
        if not success:
            print("\nNote: some errors occurred while building the tree.")
        return success
    else:
        result = get_dep_tree(
            pkg,
            max_depth=max_depth,
            output_format=output_format,
            skip_optional=skip_optional,
        )
        print(result)
        return True


def tree_project(
    path: str,
    indent="",
    max_depth=None,
    show_files=True,
    show_dirs=True,
    sort_by="name",  # "name", "size", "mtime"
    full_path=False,
    ignore=None,  # list of suffixes to ignore, e.g. [".pyc", ".log"]
    current_depth=0,
) -> str:
    output = ""
    path = Path(path)

    if not path.exists():
        raise FileNotFoundError(
            f"path '{path}' not found in current path '{Path.cwd()}'"
        )

    if max_depth is not None and current_depth > max_depth:
        return ""

    items = [p for p in path.iterdir() if not (ignore and p.suffix in ignore)]

    if sort_by == "name":
        items.sort(key=lambda x: x.name.lower())
    elif sort_by == "size":
        items.sort(key=lambda x: x.stat().st_size)
    elif sort_by == "mtime":
        items.sort(key=lambda x: x.stat().st_mtime)

    output += f"{path.name}\n"
    for i, item in enumerate(items):
        if item.is_file() and not show_files:
            continue
        if item.is_dir() and not show_dirs:
            continue

        branch = "└── " if i == len(items) - 1 else "├── "
        display = str(item) if full_path else item.name
        output += indent + branch + display + "\n"

        if item.is_dir():
            extension = "    " if i == len(items) - 1 else "│   "
            output += tree_project(
                item,
                indent + extension,
                max_depth=max_depth,
                show_files=show_files,
                show_dirs=show_dirs,
                sort_by=sort_by,
                full_path=full_path,
                ignore=ignore,
                current_depth=current_depth + 1,
            )
    return output
