"""Add missing imports to test code based on source module symbols."""

from __future__ import annotations

import builtins
import logging
from dataclasses import dataclass, field

import libcst as cst
import sentry_sdk
from libcst.codemod import CodemodContext
from libcst.codemod.visitors import AddImportsVisitor

from core.languages.python.cst_utils import (
    DefinitionRemover,
    DepthTrackingMixin,
    collect_imported_names_from_import,
    collect_imported_names_from_import_from,
    extract_imports_from_import,
    extract_imports_from_import_from,
    file_path_to_module_path,
    get_dotted_name,
)

logger = logging.getLogger(__name__)

BUILTIN_NAMES = frozenset(dir(builtins))


@dataclass
class SourceAnalysis:
    """Analysis results for a source module."""

    defined_symbols: set[str] = field(default_factory=set)
    public_symbols: set[str] = field(default_factory=set)
    imports: dict[str, tuple[str, str]] = field(default_factory=dict)
    referenced_names: set[str] = field(default_factory=set)
    symbol_to_module: dict[str, str] = field(default_factory=dict)


@dataclass
class TestAnalysis:
    """Analysis results for a test module."""

    local_definitions: set[str] = field(default_factory=set)
    imported_names: set[str] = field(default_factory=set)
    used_names: set[str] = field(default_factory=set)

    @property
    def undefined_names(self) -> set[str]:
        return self.used_names - self.local_definitions - self.imported_names - BUILTIN_NAMES


class SourceModuleAnalyzer(DepthTrackingMixin, cst.CSTVisitor):
    """Visitor to analyze a source module and collect symbol information."""

    def __init__(self, module_path: str) -> None:
        super().__init__()
        self.module_path = module_path
        self.defined_symbols: set[str] = set()
        self.public_symbols: set[str] = set()
        self.imports: dict[str, tuple[str, str]] = {}
        self.referenced_names: set[str] = set()
        self._in_assignment_target = False

    def visit_ClassDef(self, node: cst.ClassDef) -> bool | None:
        if self._is_top_level():
            name = node.name.value
            self.defined_symbols.add(name)
            if not name.startswith("_"):
                self.public_symbols.add(name)
        self._visit_class()
        return True

    def leave_ClassDef(self, original_node: cst.ClassDef) -> None:  # noqa: ARG002
        self._leave_class()

    def visit_FunctionDef(self, node: cst.FunctionDef) -> bool | None:
        if self._is_top_level():
            name = node.name.value
            self.defined_symbols.add(name)
            if not name.startswith("_"):
                self.public_symbols.add(name)
        self._visit_function()
        return True

    def leave_FunctionDef(self, original_node: cst.FunctionDef) -> None:  # noqa: ARG002
        self._leave_function()

    def visit_ImportFrom(self, node: cst.ImportFrom) -> bool | None:
        self.imports.update(extract_imports_from_import_from(node))
        return False

    def visit_Import(self, node: cst.Import) -> bool | None:
        self.imports.update(extract_imports_from_import(node))
        return False

    def visit_AssignTarget(self, node: cst.AssignTarget) -> bool | None:  # noqa: ARG002
        self._in_assignment_target = True
        return True

    def leave_AssignTarget(self, original_node: cst.AssignTarget) -> None:  # noqa: ARG002
        self._in_assignment_target = False

    def visit_AnnAssign(self, node: cst.AnnAssign) -> bool | None:
        # For annotated assignments like x: int = 5, skip the target
        # but visit the annotation and value
        if node.annotation:
            node.annotation.visit(self)
        if node.value:
            node.value.visit(self)
        return False

    def visit_Attribute(self, node: cst.Attribute) -> bool | None:
        # For ElementType.TEXT, we want to collect ElementType but not TEXT
        # Manually visit the value (base) and skip the attr
        node.value.visit(self)
        return False

    def visit_Name(self, node: cst.Name) -> bool | None:
        if self._in_assignment_target:
            return False

        name = node.value
        if self.is_importable_symbol(name):
            self.referenced_names.add(name)
        return False

    def is_importable_symbol(self, name: str) -> bool:
        """Check if a name looks like an importable symbol.

        Includes:
        - Capitalized names (classes like ElementType)
        - Private classes (_ChunkerSpec)
        - Constants (MAX_SIZE, but not single chars)

        Excludes:
        - Local variables (lowercase)
        - Magic methods (__init__)
        - Single character names
        """
        if not name:
            return False
        if name[0].isupper():
            return True
        if len(name) > 1 and name[0] == "_" and name[1].isupper():
            return True
        return len(name) > 1 and name.isupper() and name.isalpha()

    def get_analysis(self) -> SourceAnalysis:
        return SourceAnalysis(
            defined_symbols=self.defined_symbols,
            public_symbols=self.public_symbols,
            imports=self.imports,
            referenced_names=self.referenced_names,
            symbol_to_module=dict.fromkeys(self.defined_symbols, self.module_path),
        )


class TestModuleAnalyzer(DepthTrackingMixin, cst.CSTVisitor):
    """Visitor to analyze a test module and collect symbol information."""

    def __init__(self) -> None:
        super().__init__()
        self.local_definitions: set[str] = set()
        self.imported_names: set[str] = set()
        self.used_names: set[str] = set()
        self._assigned_names: set[str] = set()

    def visit_ClassDef(self, node: cst.ClassDef) -> bool | None:
        if self._is_top_level():
            self.local_definitions.add(node.name.value)
        self._visit_class()
        return True

    def leave_ClassDef(self, original_node: cst.ClassDef) -> None:  # noqa: ARG002
        self._leave_class()

    def visit_FunctionDef(self, node: cst.FunctionDef) -> bool | None:
        if self._is_top_level():
            self.local_definitions.add(node.name.value)
        self.collect_function_parameters(node.params)
        self._visit_function()
        return True

    def leave_FunctionDef(self, original_node: cst.FunctionDef) -> None:  # noqa: ARG002
        self._leave_function()

    def collect_function_parameters(self, params: cst.Parameters) -> None:
        """Collect all parameter names from a function signature."""
        for param in params.params:
            self._assigned_names.add(param.name.value)
        if params.star_arg and isinstance(params.star_arg, cst.Param):
            self._assigned_names.add(params.star_arg.name.value)
        if params.star_kwarg:
            self._assigned_names.add(params.star_kwarg.name.value)
        for param in params.kwonly_params:
            self._assigned_names.add(param.name.value)

    def visit_ImportFrom(self, node: cst.ImportFrom) -> bool | None:
        self.imported_names.update(collect_imported_names_from_import_from(node))
        return False

    def visit_Import(self, node: cst.Import) -> bool | None:
        self.imported_names.update(collect_imported_names_from_import(node))
        return False

    def visit_Assign(self, node: cst.Assign) -> bool | None:
        # Track assigned names
        for target in node.targets:
            self._collect_assigned_names(target.target)
        return True

    def visit_AnnAssign(self, node: cst.AnnAssign) -> bool | None:
        if node.target:
            self._collect_assigned_names(node.target)
        return True

    def visit_NamedExpr(self, node: cst.NamedExpr) -> bool | None:
        self._collect_assigned_names(node.target)
        return True

    def visit_For(self, node: cst.For) -> bool | None:
        self._collect_assigned_names(node.target)
        return True

    def visit_CompFor(self, node: cst.CompFor) -> bool | None:
        self._collect_assigned_names(node.target)
        return True

    def visit_AsName(self, node: cst.AsName) -> bool | None:
        # Handles 'with X as y' and 'except E as e'
        if isinstance(node.name, cst.Name):
            self._assigned_names.add(node.name.value)
        return True

    def _collect_assigned_names(self, target: cst.BaseExpression) -> None:
        if isinstance(target, cst.Name):
            self._assigned_names.add(target.value)
        elif isinstance(target, (cst.Tuple, cst.List)):
            for element in target.elements:
                if isinstance(element, (cst.Element, cst.StarredElement)):
                    self._collect_assigned_names(element.value)

    def visit_Arg(self, node: cst.Arg) -> bool | None:
        if node.value:
            node.value.visit(self)
        return False

    def visit_Name(self, node: cst.Name) -> bool | None:
        self.used_names.add(node.value)
        return False

    def get_analysis(self) -> TestAnalysis:
        return TestAnalysis(
            local_definitions=self.local_definitions | self._assigned_names,
            imported_names=self.imported_names,
            used_names=self.used_names,
        )


def analyze_source(source_cst: cst.Module, module_path: str) -> SourceAnalysis:
    analyzer = SourceModuleAnalyzer(module_path)
    source_cst.visit(analyzer)
    return analyzer.get_analysis()


def analyze_test(test_module: cst.Module) -> TestAnalysis:
    analyzer = TestModuleAnalyzer()
    test_module.visit(analyzer)
    return analyzer.get_analysis()


def merge_source_analyses(analyses: list[tuple[str, SourceAnalysis]]) -> SourceAnalysis:
    merged = SourceAnalysis()
    for module_path, analysis in analyses:
        merged.defined_symbols.update(analysis.defined_symbols)
        merged.public_symbols.update(analysis.public_symbols)
        merged.referenced_names.update(analysis.referenced_names)

        for name, value in analysis.imports.items():
            if name not in merged.imports:
                merged.imports[name] = value

        for symbol in analysis.symbol_to_module:
            if symbol not in merged.symbol_to_module:
                merged.symbol_to_module[symbol] = module_path

    return merged


def get_source_analysis(source_cst: cst.Module | dict[str, cst.Module], module_path: str) -> SourceAnalysis | None:
    """Get analysis of source modules, handling both single and multi-context sources."""
    if not isinstance(source_cst, dict):
        return analyze_source(source_cst, module_path)

    if not source_cst:
        return None

    analyses: list[tuple[str, SourceAnalysis]] = []
    for file_path, module in source_cst.items():
        mod_path = file_path_to_module_path(str(file_path))
        analysis = analyze_source(module, mod_path)  # type: ignore[arg-type]
        analyses.append((mod_path, analysis))

    return merge_source_analyses(analyses)


def collect_symbols_to_import(
    test_analysis: TestAnalysis, source_analysis: SourceAnalysis, definitions_to_remove: set[str]
) -> set[str]:
    symbols: set[str] = set()
    symbols.update(definitions_to_remove)
    symbols.update(test_analysis.undefined_names & source_analysis.defined_symbols)
    symbols.update(test_analysis.undefined_names & set(source_analysis.imports.keys()))

    remaining_undefined = (
        test_analysis.undefined_names - source_analysis.defined_symbols - set(source_analysis.imports.keys())
    )
    symbols.update(remaining_undefined & source_analysis.referenced_names)

    return symbols


def apply_imports(
    module: cst.Module, symbols: set[str], source_analysis: SourceAnalysis, module_path: str
) -> cst.Module:
    context = CodemodContext()

    for symbol in symbols:
        if symbol in source_analysis.imports:
            import_module, original_name = source_analysis.imports[symbol]
            if original_name == symbol:
                AddImportsVisitor.add_needed_import(context, import_module, symbol)
            else:
                AddImportsVisitor.add_needed_import(context, import_module, original_name, asname=symbol)
        elif symbol in source_analysis.symbol_to_module:
            defining_module = source_analysis.symbol_to_module[symbol]
            AddImportsVisitor.add_needed_import(context, defining_module, symbol)
        else:
            AddImportsVisitor.add_needed_import(context, module_path, symbol)

    return module.visit(AddImportsVisitor(context))


def add_missing_imports(
    test_module: cst.Module, source_cst: cst.Module | dict[str, cst.Module], module_path: str
) -> cst.Module:
    """Add missing imports to test code based on source module symbols.

    Args:
        test_module: Parsed libcst.Module of the test code
        source_cst: Source module(s) - can be:
            - cst.Module: a single parsed module
            - dict[str, cst.Module]: file paths to parsed modules (multi-context)
        module_path: Default import path (e.g., "mypackage.mymodule")

    Returns:
        Modified test_module with imports added.
        On any error, returns original test_module unchanged.

    """
    try:
        source_analysis = get_source_analysis(source_cst, module_path)
        if source_analysis is None:
            return test_module

        if not source_analysis.defined_symbols and not source_analysis.imports and not source_analysis.referenced_names:
            return test_module

        test_analysis = analyze_test(test_module)
        definitions_to_remove = test_analysis.local_definitions & source_analysis.public_symbols

        modified_module = test_module
        if definitions_to_remove:
            remover = DefinitionRemover(definitions_to_remove)
            modified_module = modified_module.visit(remover)
            test_analysis = analyze_test(modified_module)

        symbols_to_import = collect_symbols_to_import(test_analysis, source_analysis, definitions_to_remove)
        if not symbols_to_import:
            return modified_module

        return apply_imports(modified_module, symbols_to_import, source_analysis, module_path)

    except Exception:  # noqa: BLE001
        logger.warning("Error in add_missing_imports, returning original module", exc_info=True)
        sentry_sdk.capture_exception()
        return test_module


# =============================================================================
# Forward reference handling
# =============================================================================


def extract_names_from_cst_annotation(annotation: cst.BaseExpression, names: set[str]) -> None:
    """Walk a CST annotation expression and collect all Name nodes.

    Handles: Name, Attribute, Subscript (generics), BinaryOperation (unions).
    """
    if isinstance(annotation, cst.Name):
        names.add(annotation.value)
    elif isinstance(annotation, cst.Attribute):
        # For X.Y.Z, we only care about the base name X
        extract_names_from_cst_annotation(annotation.value, names)
    elif isinstance(annotation, cst.Subscript):
        # e.g., List[X] or Dict[K, V] - visit the base and slice
        extract_names_from_cst_annotation(annotation.value, names)
        for element in annotation.slice:
            if isinstance(element, cst.SubscriptElement):
                slice_node = element.slice
                if isinstance(slice_node, cst.Index):
                    extract_names_from_cst_annotation(slice_node.value, names)
    elif isinstance(annotation, cst.BinaryOperation):
        # e.g., X | Y (union types)
        extract_names_from_cst_annotation(annotation.left, names)
        extract_names_from_cst_annotation(annotation.right, names)
    elif isinstance(annotation, (cst.Tuple, cst.List)):
        # e.g., tuple[X, Y] or list[X]
        for element in annotation.elements:
            if isinstance(element, (cst.Element, cst.StarredElement)):
                extract_names_from_cst_annotation(element.value, names)


class CSTAnnotationNameCollector(cst.CSTVisitor):
    """CST-based visitor to collect names used in type annotations."""

    def __init__(self) -> None:
        self.annotation_names: set[str] = set()
        self.defined_names: set[str] = set()
        self.imported_names: set[str] = set()

    def visit_Import(self, node: cst.Import) -> bool | None:
        self.imported_names.update(collect_imported_names_from_import(node))
        return False

    def visit_ImportFrom(self, node: cst.ImportFrom) -> bool | None:
        self.imported_names.update(collect_imported_names_from_import_from(node))
        return False

    def visit_FunctionDef(self, node: cst.FunctionDef) -> bool | None:
        self.defined_names.add(node.name.value)
        self.collect_annotation_names_from_function(node)
        return True

    def visit_ClassDef(self, node: cst.ClassDef) -> bool | None:
        self.defined_names.add(node.name.value)
        return True

    def visit_AnnAssign(self, node: cst.AnnAssign) -> bool | None:
        extract_names_from_cst_annotation(node.annotation.annotation, self.annotation_names)
        return True

    def collect_annotation_names_from_function(self, node: cst.FunctionDef) -> None:
        # Collect from parameters
        params = node.params
        for param in [*params.params, *params.kwonly_params, *params.posonly_params]:
            if param.annotation:
                extract_names_from_cst_annotation(param.annotation.annotation, self.annotation_names)
        if params.star_arg and isinstance(params.star_arg, cst.Param) and params.star_arg.annotation:
            extract_names_from_cst_annotation(params.star_arg.annotation.annotation, self.annotation_names)
        if params.star_kwarg and params.star_kwarg.annotation:
            extract_names_from_cst_annotation(params.star_kwarg.annotation.annotation, self.annotation_names)
        # Collect from return type
        if node.returns:
            extract_names_from_cst_annotation(node.returns.annotation, self.annotation_names)

    def get_undefined_annotation_names(self) -> set[str]:
        return self.annotation_names - self.defined_names - self.imported_names - BUILTIN_NAMES


def has_future_annotations_import(module: cst.Module) -> bool:
    """Check if module already has 'from __future__ import annotations'."""
    for stmt in module.body:
        if isinstance(stmt, cst.SimpleStatementLine):
            for small_stmt in stmt.body:
                if (
                    isinstance(small_stmt, cst.ImportFrom)
                    and small_stmt.module
                    and get_dotted_name(small_stmt.module) == "__future__"
                ):
                    if isinstance(small_stmt.names, cst.ImportStar):
                        continue
                    for alias in small_stmt.names:
                        if alias.name.value == "annotations":
                            return True
    return False


def add_future_annotations_import(module: cst.Module) -> cst.Module:
    """Add 'from __future__ import annotations' to module if needed.

    This fixes forward reference issues where type annotations reference
    names that aren't yet defined or imported. With this import, Python
    treats all annotations as strings (lazy evaluation), avoiding NameError.

    Args:
        module: The CST module to process

    Returns:
        The original module unchanged if import not needed, or a new
        module with the import added.

    """
    if has_future_annotations_import(module):
        return module

    collector = CSTAnnotationNameCollector()
    module.visit(collector)
    if not collector.get_undefined_annotation_names():
        return module

    context = CodemodContext()
    AddImportsVisitor.add_needed_import(context, "__future__", "annotations")
    return module.visit(AddImportsVisitor(context))
