import ast
import builtins
import importlib.util
import sys
from dataclasses import dataclass, field
from pathlib import Path

from .types import CraftIssue, CraftIssueSeverity


_SKIP_DIRS = frozenset({
    ".venv",
    "venv",
    "env",
    ".env",
    "site-packages",
    "node_modules",
    "__pycache__",
    ".pytest_cache",
    ".mypy_cache",
    ".ruff_cache",
    ".git",
    "build",
    "dist",
})

_MODULE_GLOBALS = frozenset({
    "__file__",
    "__name__",
    "__package__",
    "__doc__",
    "__spec__",
    "__loader__",
    "__cached__",
    "__builtins__",
})

_BUILTIN_NAMES = frozenset(dir(builtins))

def _is_skip_dir(name: str) -> bool:
    return name in _SKIP_DIRS


def _parent_map(tree: ast.AST) -> dict[ast.AST, ast.AST]:
    parents: dict[ast.AST, ast.AST] = {}
    for parent in ast.walk(tree):
        for child in ast.iter_child_nodes(parent):
            parents[child] = parent
    return parents


def _is_soft_import_context(node: ast.AST, parents: dict[ast.AST, ast.AST]) -> bool:
    current = node
    while current in parents:
        parent = parents[current]
        if isinstance(parent, ast.If) and _is_type_checking_test(parent.test):
            return True
        if isinstance(parent, ast.Try) and _node_is_in_try_body(current, parent):
            if any(_handler_catches_import_error(handler) for handler in parent.handlers):
                return True
        current = parent
    return False


def _is_type_checking_test(node: ast.AST) -> bool:
    if isinstance(node, ast.Name):
        return node.id == "TYPE_CHECKING"
    if isinstance(node, ast.Attribute):
        return node.attr == "TYPE_CHECKING"
    return False


def _node_is_in_try_body(node: ast.AST, try_node: ast.Try) -> bool:
    return any(child is node for child in try_node.body)


def _handler_catches_import_error(handler: ast.ExceptHandler) -> bool:
    if handler.type is None:
        return False
    caught = handler.type.elts if isinstance(handler.type, ast.Tuple) else [handler.type]
    return any(_exception_name(exc) in {"ImportError", "ModuleNotFoundError"} for exc in caught)


def _exception_name(node: ast.AST) -> str:
    if isinstance(node, ast.Name):
        return node.id
    if isinstance(node, ast.Attribute):
        return node.attr
    return ""


@dataclass
class _FuncSig:
    """Lightweight function/method signature for workspace index."""
    name: str
    params: list[str] = field(default_factory=list)
    defaults_count: int = 0
    has_varargs: bool = False
    has_kwargs: bool = False
    is_method: bool = False


@dataclass
class _ClassInfo:
    """Lightweight class info for workspace index."""
    name: str
    bases: list[str] = field(default_factory=list)
    methods: dict[str, _FuncSig] = field(default_factory=dict)
    attributes: set[str] = field(default_factory=set)


def _has_decorator(node: ast.FunctionDef | ast.AsyncFunctionDef, name: str) -> bool:
    """Check if a function node has a decorator with the given name."""
    for decorator in node.decorator_list:
        if isinstance(decorator, ast.Name) and decorator.id == name:
            return True
        if isinstance(decorator, ast.Attribute) and decorator.attr == name:
            return True
    return False


class ReferenceChecker:
    def __init__(self):
        self._workspace_symbols: dict[str, set[str]] | None = None
        self._workspace_modules: set[str] | None = None
        self._workspace_root: Path | None = None
        # Phase 4: richer signature/class data
        self._workspace_functions: dict[str, _FuncSig] = {}       # module.funcname -> sig
        self._workspace_classes: dict[str, _ClassInfo] = {}       # module.ClassName -> info

    # ------------------------------------------------------------------
    # Workspace index
    # ------------------------------------------------------------------

    def invalidate_workspace_index(self, workspace_root: Path | None = None) -> None:
        """Clear cached workspace index so the next check rebuilds it lazily."""
        if workspace_root is not None and self._workspace_root != workspace_root:
            return
        self._workspace_symbols = None
        self._workspace_modules = None
        self._workspace_root = None
        self._workspace_functions = {}
        self._workspace_classes = {}

    def _build_workspace_index(self, workspace_root: Path) -> None:
        if self._workspace_symbols is not None and self._workspace_root == workspace_root:
            return

        self._workspace_root = workspace_root
        self._workspace_symbols = {}
        self._workspace_modules = set()
        self._workspace_functions = {}
        self._workspace_classes = {}

        import os
        for root, dirs, files in os.walk(workspace_root):
            dirs[:] = [d for d in dirs if not _is_skip_dir(d)]
            for file in files:
                if not file.endswith(".py"):
                    continue
                py_file = Path(root) / file
                try:
                    rel_path = py_file.relative_to(workspace_root)
                    parts = list(rel_path.parts)
                    if parts[-1] == "__init__.py":
                        parts.pop()
                    else:
                        parts[-1] = parts[-1][:-3]

                    module_path = ".".join(parts) if parts else ""
                    if module_path:
                        self._workspace_modules.add(module_path)

                    content = py_file.read_text(encoding="utf-8")
                    tree = ast.parse(content)
                    symbols = set()
                    for node in ast.walk(tree):
                        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
                            symbols.add(node.name)
                        elif isinstance(node, ast.Name) and isinstance(node.ctx, ast.Store):
                            symbols.add(node.id)
                        elif isinstance(node, ast.Import):
                            for alias in node.names:
                                symbols.add(alias.asname or alias.name.split('.')[0])
                        elif isinstance(node, ast.ImportFrom):
                            for alias in node.names:
                                symbols.add(alias.asname or alias.name)

                    if module_path:
                        self._workspace_symbols[module_path] = symbols

                    # Phase 4: collect function signatures and class info
                    for node in ast.iter_child_nodes(tree):
                        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                            sig = self._extract_signature(node)
                            key = f"{module_path}.{node.name}" if module_path else node.name
                            self._workspace_functions[key] = sig
                        elif isinstance(node, ast.ClassDef):
                            ci = self._extract_class_info(node)
                            key = f"{module_path}.{node.name}" if module_path else node.name
                            self._workspace_classes[key] = ci
                            # Also index methods as module.ClassName.method
                            for mname, msig in ci.methods.items():
                                mkey = f"{key}.{mname}"
                                self._workspace_functions[mkey] = msig
                except Exception:
                    continue

    @staticmethod
    def _extract_signature(node: ast.FunctionDef | ast.AsyncFunctionDef, is_method: bool = False) -> _FuncSig:
        params: list[str] = []
        defaults_count = 0
        has_varargs = False
        has_kwargs = False

        for arg in node.args.args:
            params.append(arg.arg)
        if node.args.vararg:
            has_varargs = True
            params.append(f"*{node.args.vararg.arg}")
        if node.args.kwarg:
            has_kwargs = True
            params.append(f"**{node.args.kwarg.arg}")

        # Count defaults (they align with the *last* positional params)
        defaults_count = len(node.args.defaults)

        return _FuncSig(
            name=node.name,
            params=params,
            defaults_count=defaults_count,
            has_varargs=has_varargs,
            has_kwargs=has_kwargs,
            is_method=is_method,
        )

    @staticmethod
    def _extract_class_info(node: ast.ClassDef) -> _ClassInfo:
        bases: list[str] = []
        for base in node.bases:
            if isinstance(base, ast.Name):
                bases.append(base.id)
            elif isinstance(base, ast.Attribute):
                bases.append(ast.unparse(base))

        methods: dict[str, _FuncSig] = {}
        attributes: set[str] = set()

        for child in ast.iter_child_nodes(node):
            if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)):
                is_method = not _has_decorator(child, "staticmethod")
                methods[child.name] = ReferenceChecker._extract_signature(child, is_method=is_method)
            elif isinstance(child, ast.Assign):
                for target in child.targets:
                    if isinstance(target, ast.Name):
                        attributes.add(target.id)
            elif isinstance(child, ast.AnnAssign) and isinstance(child.target, ast.Name):
                attributes.add(child.target.id)

        return _ClassInfo(name=node.name, bases=bases, methods=methods, attributes=attributes)

    # ------------------------------------------------------------------
    # Main check entry point
    # ------------------------------------------------------------------

    def check(self, capsule, workspace_root: Path | None = None) -> list[CraftIssue]:
        if getattr(capsule, "language", "python") != "python":
            return []

        if workspace_root:
            self._build_workspace_index(workspace_root)

        try:
            tree = ast.parse(capsule.proposed_code)
        except SyntaxError:
            return []

        package_parts: list[str] = []
        if getattr(capsule, "path", None) and workspace_root:
            try:
                rel_path = capsule.path.relative_to(workspace_root)
                parts = list(rel_path.parts)
                if parts:
                    package_parts = parts[:-1]
            except ValueError:
                pass

        local_defs = set(_MODULE_GLOBALS)
        for node in ast.walk(tree):
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
                local_defs.add(node.name)
            elif isinstance(node, ast.Name) and isinstance(node.ctx, ast.Store):
                local_defs.add(node.id)
            elif isinstance(node, ast.arg):
                local_defs.add(node.arg)
            elif isinstance(node, ast.ExceptHandler) and node.name:
                local_defs.add(node.name)

        imported_names: dict[str, str] = {}
        import_sources: list[tuple[str, ast.AST, str | None]] = []

        parent_map = _parent_map(tree)

        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                for alias in node.names:
                    name = alias.asname or alias.name.split('.')[0]
                    imported_names[name] = alias.name
                    if not _is_soft_import_context(node, parent_map):
                        import_sources.append((alias.name, node, None))
            elif isinstance(node, ast.ImportFrom):
                level = node.level or 0
                base_module = node.module or ""

                if level > 0 and package_parts:
                    if level == 1:
                        resolved_parts = list(package_parts)
                    else:
                        resolved_parts = package_parts[:-(level - 1)]

                    if base_module:
                        resolved_parts.append(base_module)
                    resolved_module = ".".join(resolved_parts)
                else:
                    resolved_module = base_module if level == 0 else ("." * level) + base_module

                if resolved_module:
                    for alias in node.names:
                        name = alias.asname or alias.name
                        imported_names[name] = f"{resolved_module}.{alias.name}"
                        if not _is_soft_import_context(node, parent_map):
                            import_sources.append((resolved_module, node, alias.name))

        issues: list[CraftIssue] = []
        builtin_names = set(dir(builtins))

        # 1. Undefined-name check
        for node in ast.walk(tree):
            if isinstance(node, ast.Name) and isinstance(node.ctx, ast.Load):
                name = node.id
                if name in builtin_names:
                    continue
                if name in local_defs:
                    continue
                if name in imported_names:
                    continue

                issues.append(
                    CraftIssue(
                        line=node.lineno,
                        column=node.col_offset,
                        code="undefined-name",
                        message=f"Name '{name}' is used but never defined or imported.",
                        suggestion="Define or import the name before using it.",
                        severity=CraftIssueSeverity.HARD,
                    )
                )

        # 2. Broken-import check
        stdlib_modules = getattr(sys, "stdlib_module_names", set())

        for mod_path, node, symbol in import_sources:
            base_module = mod_path.split('.')[0]
            resolved = False

            if base_module in stdlib_modules or mod_path in stdlib_modules:
                resolved = True

            if not resolved:
                try:
                    if importlib.util.find_spec(base_module) is not None:
                        resolved = True
                except Exception:
                    resolved = False

            if not resolved and self._workspace_modules is not None:
                if mod_path in self._workspace_modules or base_module in self._workspace_modules:
                    resolved = True

            if not resolved and workspace_root:
                rel_path = mod_path.replace('.', '/')
                if (workspace_root / f"{rel_path}.py").exists() or (workspace_root / rel_path / "__init__.py").exists():
                    resolved = True

            if not resolved:
                issues.append(
                    CraftIssue(
                        line=node.lineno,
                        column=node.col_offset,
                        code="broken-import",
                        message=f"Import source '{mod_path}' could not be resolved in workspace or stdlib.",
                        suggestion="Install the missing package or correct the import path.",
                        severity=CraftIssueSeverity.HARD,
                    )
                )
            elif symbol and symbol != "*":
                if self._workspace_symbols and mod_path in self._workspace_symbols:
                    if symbol not in self._workspace_symbols[mod_path]:
                        issues.append(
                            CraftIssue(
                                line=node.lineno,
                                column=node.col_offset,
                                code="broken-import",
                                message=f"Import source '{mod_path}' has no symbol '{symbol}'.",
                                suggestion="Check the spelling or define the missing symbol.",
                                severity=CraftIssueSeverity.HARD,
                            )
                        )
                elif base_module in stdlib_modules:
                    try:
                        mod = __import__(mod_path, fromlist=[symbol])
                        if not hasattr(mod, symbol):
                            issues.append(
                                CraftIssue(
                                    line=node.lineno,
                                    column=node.col_offset,
                                    code="broken-import",
                                    message=f"Import source '{mod_path}' has no symbol '{symbol}'.",
                                    suggestion="Check the spelling or correct the import path.",
                                    severity=CraftIssueSeverity.HARD,
                                )
                            )
                    except Exception:
                        continue

        # 3. Phase 4: signature mismatch and attribute checks
        if self._workspace_functions or self._workspace_classes:
            issues.extend(self._check_calls(tree, local_defs, imported_names))

        return issues

    # ------------------------------------------------------------------
    # Phase 4: Call signature & attribute checks
    # ------------------------------------------------------------------

    def _check_calls(
        self,
        tree: ast.Module,
        local_defs: set[str],
        imported_names: dict[str, str],
    ) -> list[CraftIssue]:
        issues: list[CraftIssue] = []
        local_functions, local_classes, instance_types = self._collect_local_call_context(tree, imported_names)

        for node in ast.walk(tree):
            if isinstance(node, ast.Call):
                issue = self._check_single_call(
                    node,
                    local_defs,
                    imported_names,
                    local_functions,
                    local_classes,
                    instance_types,
                )
                if issue:
                    issues.append(issue)

            elif isinstance(node, ast.Attribute) and isinstance(node.ctx, ast.Load):
                issue = self._check_attribute(node, imported_names, local_classes, instance_types)
                if issue:
                    issues.append(issue)

        return issues

    def _collect_local_call_context(
        self,
        tree: ast.Module,
        imported_names: dict[str, str],
    ) -> tuple[dict[str, _FuncSig], dict[str, _ClassInfo], dict[str, str]]:
        local_functions: dict[str, _FuncSig] = {}
        local_classes: dict[str, _ClassInfo] = {}
        instance_types: dict[str, str] = {}

        for node in ast.iter_child_nodes(tree):
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                local_functions[node.name] = self._extract_signature(node)
            elif isinstance(node, ast.ClassDef):
                local_classes[node.name] = self._extract_class_info(node)

        class_names = set(local_classes)
        class_names.update(name for name, target in imported_names.items() if self._find_class(target) is not None)

        for node in ast.walk(tree):
            if not isinstance(node, ast.Assign):
                continue
            if not isinstance(node.value, ast.Call):
                continue
            cls_name = self._class_name_from_constructor(node.value.func, imported_names, local_classes)
            if cls_name is None or cls_name not in class_names:
                continue
            for target in node.targets:
                if isinstance(target, ast.Name):
                    instance_types[target.id] = cls_name

        return local_functions, local_classes, instance_types

    def _resolve_callee(
        self,
        func_node: ast.expr,
        local_defs: set[str],
        imported_names: dict[str, str],
        local_functions: dict[str, _FuncSig],
        local_classes: dict[str, _ClassInfo],
        instance_types: dict[str, str],
    ) -> tuple[str | None, str | None, _FuncSig | None]:
        """Resolve a call target to (module_path, symbol_name, signature) if possible.

        Returns (None, None, None) when we cannot resolve with high confidence.
        """
        if isinstance(func_node, ast.Name):
            name = func_node.id
            if name in _BUILTIN_NAMES:
                return None, None, None
            if name in local_functions:
                return None, name, local_functions[name]
            if name in local_classes:
                return None, name, local_classes[name].methods.get("__init__")
            if name in imported_names:
                full = imported_names[name]
                parts = full.rsplit(".", 1)
                if len(parts) == 2:
                    return parts[0], parts[1], None
                return None, parts[0], None
            if name in local_defs:
                return None, None, None
            return None, None, None

        elif isinstance(func_node, ast.Attribute):
            if isinstance(func_node.value, ast.Name):
                imported_target = imported_names.get(func_node.value.id)
                if imported_target and f"{imported_target}.{func_node.attr}" in self._workspace_functions:
                    return imported_target, func_node.attr, None
            cls_info = self._resolve_attribute_owner(func_node.value, imported_names, local_classes, instance_types)
            if cls_info is not None:
                return cls_info.name, func_node.attr, cls_info.methods.get(func_node.attr)
            return None, None, None

        return None, None, None

    @staticmethod
    def _resolve_obj_name(node: ast.expr) -> str | None:
        """Best-effort resolve an expression to a simple name for class lookup."""
        if isinstance(node, ast.Name):
            return node.id
        if isinstance(node, ast.Attribute):
            return node.attr
        return None

    def _check_single_call(
        self,
        call_node: ast.Call,
        local_defs: set[str],
        imported_names: dict[str, str],
        local_functions: dict[str, _FuncSig],
        local_classes: dict[str, _ClassInfo],
        instance_types: dict[str, str],
    ) -> CraftIssue | None:
        module_path, symbol, resolved_sig = self._resolve_callee(
            call_node.func,
            local_defs,
            imported_names,
            local_functions,
            local_classes,
            instance_types,
        )
        if module_path is None and symbol is None:
            return None

        # Look up the function signature
        sig = resolved_sig or self._find_signature(module_path, symbol)
        if sig is None:
            return None

        # Count positional arguments (args before any keyword)
        pos_args = 0
        for arg in call_node.args:
            if isinstance(arg, ast.Starred):
                # *args unpacking — can't count, skip
                return None
            pos_args += 1

        # Count required positional params (total params minus defaults, minus *args/**kwargs)
        param_list = list(sig.params)
        # For methods, skip the first param (self/cls) — it's passed implicitly
        if sig.is_method and param_list:
            param_list.pop(0)
            # Adjust defaults_count: if the removed param had a default, reduce count
            if sig.defaults_count > 0 and sig.defaults_count == len(sig.params) - (len(param_list)):
                pass  # defaults align with last params; removing first doesn't change count
            # Actually, defaults always align with the last N params. Removing self (first param)
            # doesn't change how many defaults there are, but the total param count drops by 1.
            # If defaults_count exceeded the new param count, cap it.
            effective_defaults = min(sig.defaults_count, len(param_list))
        else:
            effective_defaults = sig.defaults_count

        total_params = len(param_list)
        required = total_params - effective_defaults
        if sig.has_varargs:
            required -= 1  # *args param
        if sig.has_kwargs:
            required -= 1  # **kwargs param

        # If the function has *args, any number of positional args is fine
        if sig.has_varargs:
            return None

        max_pos = total_params
        if sig.has_kwargs:
            max_pos -= 1

        keyword_names = {kw.arg for kw in call_node.keywords if kw.arg is not None}
        # Params covered by keywords reduce the positional requirement
        params_covered_by_kw = 0
        for p in param_list:
            if p in keyword_names:
                params_covered_by_kw += 1

        effective_required = max(0, required - params_covered_by_kw)
        effective_max = max_pos - params_covered_by_kw

        if pos_args < effective_required:
            display_name = symbol if module_path is None else f"{module_path}.{symbol}"
            return CraftIssue(
                line=call_node.lineno,
                column=call_node.col_offset,
                code="call-signature",
                message=(
                    f"Function '{display_name}' expects at least {effective_required} "
                    f"positional argument(s) but {pos_args} were provided."
                ),
                suggestion="Provide the required positional arguments.",
                severity=CraftIssueSeverity.SOFT,
            )

        if pos_args > effective_max:
            display_name = symbol if module_path is None else f"{module_path}.{symbol}"
            return CraftIssue(
                line=call_node.lineno,
                column=call_node.col_offset,
                code="call-signature",
                message=(
                    f"Function '{display_name}' accepts at most {effective_max} "
                    f"positional argument(s) but {pos_args} were provided."
                ),
                suggestion="Remove excess positional arguments or use keyword arguments.",
                severity=CraftIssueSeverity.SOFT,
            )

        return None

    def _find_signature(self, module_path: str | None, symbol: str) -> _FuncSig | None:
        """Find a function signature in the workspace index."""
        # Direct lookup: module.symbol
        if module_path:
            key = f"{module_path}.{symbol}"
            if key in self._workspace_functions:
                return self._workspace_functions[key]
            cls_key = f"{module_path}.{symbol}"
            if cls_key in self._workspace_classes:
                return self._workspace_classes[cls_key].methods.get("__init__")
            return None

        return None

    def _find_class(self, full_name: str) -> _ClassInfo | None:
        if full_name in self._workspace_classes:
            return self._workspace_classes[full_name]
        return None

    def _class_name_from_constructor(
        self,
        func_node: ast.expr,
        imported_names: dict[str, str],
        local_classes: dict[str, _ClassInfo],
    ) -> str | None:
        if not isinstance(func_node, ast.Name):
            return None
        name = func_node.id
        if name in local_classes:
            return name
        if name in imported_names and self._find_class(imported_names[name]) is not None:
            return name
        return None

    def _resolve_attribute_owner(
        self,
        value_node: ast.expr,
        imported_names: dict[str, str],
        local_classes: dict[str, _ClassInfo],
        instance_types: dict[str, str],
    ) -> _ClassInfo | None:
        if not isinstance(value_node, ast.Name):
            return None
        name = value_node.id
        if name in local_classes:
            return local_classes[name]
        if name in instance_types:
            cls_name = instance_types[name]
            if cls_name in local_classes:
                return local_classes[cls_name]
            imported_full = imported_names.get(cls_name)
            if imported_full:
                return self._find_class(imported_full)
        if name in imported_names:
            return self._find_class(imported_names[name])
        return None

    def _check_attribute(
        self,
        attr_node: ast.Attribute,
        imported_names: dict[str, str],
        local_classes: dict[str, _ClassInfo],
        instance_types: dict[str, str],
    ) -> CraftIssue | None:
        """Check that obj.attr exists when obj resolves to a known workspace class."""
        cls_info = self._resolve_attribute_owner(attr_node.value, imported_names, local_classes, instance_types)
        if cls_info is None:
            return None

        attr = attr_node.attr
        # Check methods and attributes
        if attr in cls_info.methods or attr in cls_info.attributes:
            return None

        if attr.startswith("__") and attr.endswith("__"):
            return None

        # If the class has bases, we can't be sure — skip
        if cls_info.bases:
            return None

        return CraftIssue(
            line=attr_node.lineno,
            column=attr_node.col_offset,
            code="missing-attribute",
            message=f"Class '{cls_info.name}' has no attribute '{attr}'.",
            suggestion="Check the attribute name or define it on the class.",
            severity=CraftIssueSeverity.SOFT,
        )
