"""
indexer.py — Walks a codebase using Python's ast module and builds an index
of all defined functions, classes, importable names, and existing files.
"""

from __future__ import annotations

import ast
import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple


@dataclass
class FunctionDef:
    name: str
    file: str
    line: int
    is_method: bool = False
    parent_class: Optional[str] = None


@dataclass
class ClassDef:
    name: str
    file: str
    line: int
    methods: List[str] = field(default_factory=list)


@dataclass
class ImportedName:
    name: str
    module: str
    file: str
    line: int


@dataclass
class CodebaseIndex:
    """Complete index of a codebase."""

    # All function definitions: name -> list of FunctionDef (multiple files OK)
    functions: Dict[str, List[FunctionDef]] = field(default_factory=dict)

    # All class definitions: name -> list of ClassDef
    classes: Dict[str, List[ClassDef]] = field(default_factory=dict)

    # All methods across all classes: method_name -> list of (ClassDef, line)
    methods: Dict[str, List[Tuple[str, str, int]]] = field(default_factory=dict)

    # All existing file paths (relative to codebase root)
    files: Set[str] = field(default_factory=set)

    # All importable names: module_path -> set of exported names
    module_exports: Dict[str, Set[str]] = field(default_factory=dict)

    # All top-level variable/constant names: name -> (file, line)
    variables: Dict[str, List[Tuple[str, int]]] = field(default_factory=dict)

    # Stats
    python_files_count: int = 0
    total_files_count: int = 0
    parse_errors: List[str] = field(default_factory=list)


class CodebaseIndexer:
    """Indexes a codebase directory for all code symbols."""

    # Extensions to include when walking (non-Python files still get indexed
    # for file-existence checks)
    PYTHON_EXTENSIONS = {".py", ".pyw"}
    SKIP_DIRS = {
        ".git",
        ".hg",
        ".svn",
        "__pycache__",
        ".tox",
        ".venv",
        "venv",
        "env",
        ".env",
        "node_modules",
        ".mypy_cache",
        ".pytest_cache",
        "dist",
        "build",
        "*.egg-info",
        ".eggs",
        "site-packages",
    }

    def __init__(self, root: str) -> None:
        self.root = Path(root).resolve()
        self._index = CodebaseIndex()

    def build(self) -> CodebaseIndex:
        """Walk the codebase and populate the index."""
        self._walk_files()
        return self._index

    # ------------------------------------------------------------------
    # Private helpers
    # ------------------------------------------------------------------

    def _walk_files(self) -> None:
        for dirpath, dirnames, filenames in os.walk(self.root):
            # Prune directories we don't want to descend into
            dirnames[:] = [
                d
                for d in dirnames
                if d not in self.SKIP_DIRS and not d.endswith(".egg-info")
            ]

            for filename in filenames:
                full_path = Path(dirpath) / filename
                rel_path = str(full_path.relative_to(self.root))

                self._index.files.add(rel_path)
                self._index.total_files_count += 1

                if full_path.suffix in self.PYTHON_EXTENSIONS:
                    self._index.python_files_count += 1
                    self._index_python_file(full_path, rel_path)

    def _index_python_file(self, full_path: Path, rel_path: str) -> None:
        try:
            source = full_path.read_text(encoding="utf-8", errors="replace")
            tree = ast.parse(source, filename=str(full_path))
        except SyntaxError as exc:
            self._index.parse_errors.append(f"{rel_path}: SyntaxError: {exc}")
            return
        except Exception as exc:
            self._index.parse_errors.append(f"{rel_path}: {exc}")
            return

        # Derive module name from relative path
        module_name = self._path_to_module(rel_path)
        exported: Set[str] = set()

        for node in ast.walk(tree):
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                self._handle_function(node, rel_path, exported)
            elif isinstance(node, ast.ClassDef):
                self._handle_class(node, rel_path, exported)
            elif isinstance(node, ast.Assign):
                self._handle_assignment(node, rel_path, exported)
            elif isinstance(node, ast.AnnAssign):
                self._handle_ann_assignment(node, rel_path, exported)

        if module_name:
            self._index.module_exports[module_name] = exported

    def _handle_function(
        self,
        node: ast.FunctionDef,  # type: ignore[type-arg]
        rel_path: str,
        exported: Set[str],
    ) -> None:
        name = node.name
        fdef = FunctionDef(name=name, file=rel_path, line=node.lineno)
        self._index.functions.setdefault(name, []).append(fdef)
        exported.add(name)

    def _handle_class(
        self,
        node: ast.ClassDef,
        rel_path: str,
        exported: Set[str],
    ) -> None:
        class_name = node.name
        cdef = ClassDef(name=class_name, file=rel_path, line=node.lineno)

        # Collect methods
        for child in ast.walk(node):
            if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)):
                method_name = child.name
                cdef.methods.append(method_name)
                # Register method under its bare name
                self._index.methods.setdefault(method_name, []).append(
                    (class_name, rel_path, child.lineno)
                )
                # Also register as a FunctionDef so it shows up in function lookups
                mdef = FunctionDef(
                    name=method_name,
                    file=rel_path,
                    line=child.lineno,
                    is_method=True,
                    parent_class=class_name,
                )
                self._index.functions.setdefault(method_name, []).append(mdef)

        self._index.classes.setdefault(class_name, []).append(cdef)
        exported.add(class_name)

    def _handle_assignment(
        self,
        node: ast.Assign,
        rel_path: str,
        exported: Set[str],
    ) -> None:
        for target in node.targets:
            if isinstance(target, ast.Name):
                name = target.id
                self._index.variables.setdefault(name, []).append(
                    (rel_path, node.lineno)
                )
                exported.add(name)
            elif isinstance(target, ast.Tuple):
                for elt in target.elts:
                    if isinstance(elt, ast.Name):
                        name = elt.id
                        self._index.variables.setdefault(name, []).append(
                            (rel_path, node.lineno)
                        )
                        exported.add(name)

    def _handle_ann_assignment(
        self,
        node: ast.AnnAssign,
        rel_path: str,
        exported: Set[str],
    ) -> None:
        if isinstance(node.target, ast.Name):
            name = node.target.id
            self._index.variables.setdefault(name, []).append(
                (rel_path, node.lineno)
            )
            exported.add(name)

    @staticmethod
    def _path_to_module(rel_path: str) -> str:
        """Convert a relative file path to a dotted module name."""
        p = Path(rel_path)
        if p.suffix not in {".py", ".pyw"}:
            return ""
        parts = list(p.with_suffix("").parts)
        if parts and parts[-1] == "__init__":
            parts = parts[:-1]
        return ".".join(parts) if parts else ""
