#!/usr/bin/env python

import traceback
from collections.abc import Callable
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Self

from tqdm import tqdm
from tree_sitter import Language, Parser, Query, QueryCursor
from tree_sitter_language_pack import SupportedLanguage, get_language

from .data import load_stdlibs

###############################################################################


class ExtractorType(Enum):
    """Enum for selecting which file types to extract from."""

    PYTHON = "python"
    R = "r"
    ALL = "all"


@dataclass
class ImportedLibraries:
    """Container for categorized dependencies."""

    stdlib: set[str]
    third_party: set[str]
    first_party: set[str]


@dataclass
class DirectoryExtractionResult:
    """Result of extracting imports from a directory."""

    extracted: dict[Path, ImportedLibraries] = field(default_factory=dict)
    failed: dict[Path, str] = field(default_factory=dict)


###############################################################################


def _collect_files_to_extract(
    directory: Path,
    extractor_type: ExtractorType,
    recursive: bool = False,
) -> list[tuple[Path, str]]:
    """Collect all source files in a directory based on extractor type."""
    files: list[tuple[Path, str]] = []
    glob_pattern = "**/*" if recursive else "*"

    if extractor_type in (ExtractorType.PYTHON, ExtractorType.ALL):
        files.extend((f, "python") for f in directory.glob(f"{glob_pattern}.py"))

    if extractor_type in (ExtractorType.R, ExtractorType.ALL):
        files.extend((f, "r") for f in directory.glob(f"{glob_pattern}.r"))
        files.extend((f, "r") for f in directory.glob(f"{glob_pattern}.R"))

    return files


###############################################################################


class Extractor:
    SUPPORTED_LANGUAGES = (
        "python",
        "r",
        # "go",
        # "rust",
        # "javascript",
        # "typescript",
    )

    _PYTHON_QUERY = """
        (import_statement
          name: (dotted_name) @import)

        (import_statement
          name: (aliased_import
            name: (dotted_name) @import))

        (import_from_statement
          module_name: (dotted_name) @import)

        (import_from_statement
          module_name: (relative_import) @relative_import)
        """

    _R_QUERY = """
        (call
          function: (identifier) @func_name
          arguments: (arguments
            (argument [(identifier) (string)] @package)))

        (namespace_operator
          lhs: (identifier) @package)

        (namespace_operator
          lhs: (string) @package)
        """

    def __init__(self) -> None:
        self.languages: dict[SupportedLanguage, Language] = {}
        self.parsers: dict[SupportedLanguage, Parser] = {}
        self.queries: dict[SupportedLanguage, Query] = {}
        self.stdlibs = load_stdlibs()

    def _load_language(self: Self, lang: SupportedLanguage) -> None:
        """Load a language parser and query if not already loaded."""
        if lang not in self.parsers:
            self.languages[lang] = get_language(lang)
            self.parsers[lang] = Parser(self.languages[lang])

            # Cache the query for this language
            query_map: dict[SupportedLanguage, str] = {
                "python": self._PYTHON_QUERY,
                "r": self._R_QUERY,
            }
            if lang in query_map:
                self.queries[lang] = Query(self.languages[lang], query_map[lang])

    def _categorize_libraries(
        self: Self,
        deps: set[str],
        stdlib_set: set[str],
        first_party: set[str] | None = None,
        stdlib_check_func: Callable[[str], bool] | None = None,
    ) -> ImportedLibraries:
        """Categorize imports into stdlib, third-party, and first-party."""
        stdlib: set[str] = set()
        third_party: set[str] = set()
        first_party_set = first_party or set()

        for dep in deps:
            # Skip if already identified as first-party
            if dep in first_party_set:
                continue

            if stdlib_check_func:
                if stdlib_check_func(dep):
                    stdlib.add(dep)
                else:
                    third_party.add(dep)
            elif dep in stdlib_set:
                stdlib.add(dep)
            else:
                third_party.add(dep)

        return ImportedLibraries(
            stdlib=stdlib, third_party=third_party, first_party=first_party_set
        )

    def extract_python_libraries(self: Self, code: str) -> ImportedLibraries:  # noqa: C901
        """Extract imported libraries from Python code."""
        self._load_language("python")
        tree = self.parsers["python"].parse(bytes(code, "utf8"))

        imported_libs: set[str] = set()
        first_party: set[str] = set()
        query_cursor = QueryCursor(self.queries["python"])
        captures = query_cursor.captures(tree.root_node)

        # Handle absolute imports
        for node in captures.get("import", []):
            # Check if this is part of a relative import by examining the parent
            parent = node.parent
            if parent and parent.type == "relative_import":
                # Skip - will be handled in relative imports section
                continue

            dep_name = code[node.start_byte : node.end_byte]
            top_level = dep_name.split(".")[0]
            imported_libs.add(top_level)

        # Handle relative imports (first-party)
        for node in captures.get("relative_import", []):
            # Get the full import statement to extract the module
            import_statement = node.parent
            if import_statement and import_statement.type == "import_from_statement":
                # First check if there's a dotted_name in the relative_import
                has_dotted = False
                for child in node.children:
                    if child.type == "dotted_name":
                        module_text = code[child.start_byte : child.end_byte]
                        top_level = module_text.split(".")[0]
                        first_party.add(top_level)
                        has_dotted = True
                        break

                # If no dotted_name in relative_import, check what's imported
                if not has_dotted:
                    # Look for the name being imported after "import"
                    for child in import_statement.children:
                        if child.type == "dotted_name":
                            # This is after the "import" keyword
                            module_text = code[child.start_byte : child.end_byte]
                            top_level = module_text.split(".")[0]
                            first_party.add(top_level)
                            break
                        elif child.type == "aliased_import":
                            # Handle "from . import utils as u"
                            for subchild in child.children:
                                if subchild.type == "dotted_name":
                                    module_text = code[subchild.start_byte : subchild.end_byte]
                                    top_level = module_text.split(".")[0]
                                    first_party.add(top_level)
                                    break

        return self._categorize_libraries(
            imported_libs, self.stdlibs["python"], first_party=first_party
        )

    def extract_r_libraries(self: Self, code: str) -> ImportedLibraries:  # noqa: C901
        """Extract imported libraries from R code."""
        self._load_language("r")
        tree = self.parsers["r"].parse(bytes(code, "utf8"))

        imported_libs: set[str] = set()
        first_party: set[str] = set()
        source_arg_positions: set[tuple[int, int]] = set()
        query_cursor = QueryCursor(self.queries["r"])
        captures = query_cursor.captures(tree.root_node)

        func_nodes = captures.get("func_name", [])
        package_nodes = captures.get("package", [])

        # Sort both lists by position to ensure consistent ordering
        func_nodes_sorted = sorted(func_nodes, key=lambda n: n.start_byte)
        package_nodes_sorted = sorted(package_nodes, key=lambda n: n.start_byte)

        # Handle library(), require(), and source() calls
        for func_node in func_nodes_sorted:
            func_name = code[func_node.start_byte : func_node.end_byte]

            # Find the closest package node that comes after this function node
            # and is within the same call expression (before the next function)
            closest_pkg = None
            min_distance = float("inf")

            for pkg_node in package_nodes_sorted:
                if pkg_node.start_byte > func_node.start_byte:
                    distance = pkg_node.start_byte - func_node.start_byte
                    if distance < min_distance:
                        # Check if there's another function between them
                        has_func_between = any(
                            func_node.start_byte < f.start_byte < pkg_node.start_byte
                            for f in func_nodes_sorted
                        )
                        if not has_func_between:
                            closest_pkg = pkg_node
                            min_distance = distance

            if closest_pkg:
                pkg_text = code[closest_pkg.start_byte : closest_pkg.end_byte]
                pkg_text = pkg_text.strip("\"'")

                if func_name in ("library", "require"):
                    imported_libs.add(pkg_text)
                elif func_name == "source":
                    # Mark this node's position so we skip it later
                    source_arg_positions.add((closest_pkg.start_byte, closest_pkg.end_byte))
                    # Extract the base name without extension
                    base_name = Path(pkg_text).stem
                    if base_name:
                        first_party.add(base_name)

        # Handle :: and ::: operators
        for node in package_nodes_sorted:
            # Skip if this node was part of a source() call
            if (node.start_byte, node.end_byte) in source_arg_positions:
                continue

            pkg_text = code[node.start_byte : node.end_byte]
            pkg_text = pkg_text.strip("\"'")
            # Skip if this looks like a file path (contains / or ends with .R)
            if "/" not in pkg_text and not pkg_text.endswith(".R"):
                imported_libs.add(pkg_text)

        return self._categorize_libraries(
            imported_libs, self.stdlibs["r"], first_party=first_party
        )

    def extract_from_file(self: Self, file_path: str | Path) -> ImportedLibraries:
        """Extract imported libraries from a file based on its extension."""
        path = Path(file_path)
        if not path.exists():
            raise FileNotFoundError(f"File not found: {file_path}")

        # Parse
        code = path.read_text(encoding="utf-8")

        # Map file extensions to extraction methods
        ext_map = {
            ".py": self.extract_python_libraries,
            ".r": self.extract_r_libraries,
            ".R": self.extract_r_libraries,
        }

        # Parse and return or handle unsupported extension
        if path.suffix in ext_map:
            return ext_map[path.suffix](code)

        supported_exts = ", ".join(sorted(set(ext_map.keys())))
        raise ValueError(
            f"Unsupported file extension: {path.suffix}. Supported: {supported_exts}"
        )

    def extract_from_directory(
        self: Self,
        directory: str | Path,
        extractor_type: ExtractorType = ExtractorType.ALL,
        recursive: bool = False,
        show_progress: bool = True,
        progress_leave: bool = True,
    ) -> DirectoryExtractionResult:
        """
        Extract imported libraries from all source files in a directory.

        Parameters
        ----------
        directory : str | Path
            Path to the directory containing source files.
        extractor_type : ExtractorType
            Which file types to extract from: PYTHON (only .py files),
            R (only .r/.R files), or ALL (default, extracts from all supported types).
        recursive : bool
            Whether to search recursively in subdirectories (default False).
        show_progress : bool
            Whether to display a progress bar (default True).
        progress_leave : bool
            Whether to leave the progress bar visible after completion (default True).
            Set to False to remove the progress bar when done.

        Returns
        -------
        DirectoryExtractionResult
            Dataclass containing successfully extracted files and failed extractions
            with their tracebacks.

        Raises
        ------
        NotADirectoryError
            If the provided path is not a directory.
        """
        directory = Path(directory).resolve()
        if not directory.is_dir():
            raise NotADirectoryError(f"{directory} is not a directory")

        files_to_extract = _collect_files_to_extract(directory, extractor_type, recursive)
        result = DirectoryExtractionResult()

        if not files_to_extract:
            return result

        progress_bar = tqdm(
            files_to_extract,
            desc="Extracting imports",
            leave=progress_leave,
            disable=not show_progress,
        )

        for file_path, _file_type in progress_bar:
            progress_bar.set_description(f"Extracting {file_path.name}")
            try:
                result.extracted[file_path] = self.extract_from_file(file_path)
            except Exception:
                result.failed[file_path] = traceback.format_exc()

        return result
