"""
AST parsing utilities to extract spec components.
"""

import ast
import contextlib
import doctest
import hashlib
import inspect
import linecache
import re
import textwrap
from collections.abc import Callable
from pathlib import Path
from typing import Any, cast


class SpecExtractor:
    """Extract spec components from a function."""

    def __init__(self, func: Callable[..., Any]):
        self.func = func
        self._func_obj = cast(Any, func)
        raw_source = getattr(self._func_obj, "__vibesafe_source__", None)
        if raw_source is None:
            raw_source = self._load_source(func)
            self._func_obj.__vibesafe_source__ = raw_source

        # Dedent source to handle nested definitions or REPL stubs
        self.source = textwrap.dedent(raw_source)
        self.tree = ast.parse(self.source)
        self.module = inspect.getmodule(func)
        self._func_def = self._find_function_def()
        self.inferred_type = self._infer_type()

    @staticmethod
    def _load_source(func: Callable[..., Any]) -> str:
        """Best-effort retrieval of function source code."""

        try:
            return inspect.getsource(func)
        except (OSError, TypeError):
            from_cache = SpecExtractor._source_from_linecache(func)
            if from_cache is not None:
                return from_cache
        return SpecExtractor._synthesized_source(func)

    @staticmethod
    def _source_from_linecache(func: Callable[..., Any]) -> str | None:
        """Attempt to reconstruct source using linecache for interactive sessions."""

        func_obj = cast(Any, func)
        filename = inspect.getsourcefile(func) or func_obj.__code__.co_filename
        if not filename:
            return None

        lines = linecache.getlines(filename)
        if not lines:
            return None

        lineno = func_obj.__code__.co_firstlineno
        if lineno <= 0 or lineno > len(lines):
            return None

        block = inspect.getblock(lines[lineno - 1 :])
        if not block:
            return None

        return "".join(block)

    @staticmethod
    def _synthesized_source(func: Callable[..., Any]) -> str:
        """Fallback source stub when real source is unavailable."""

        signature = inspect.signature(func)
        prefix = "async def" if inspect.iscoroutinefunction(func) else "def"
        func_obj = cast(Any, func)
        return f"{prefix} {func_obj.__name__}{signature}:\n    ...\n"

    def extract_signature(self) -> str:
        """
        Extract function signature as a string.

        Returns:
            Signature string (e.g., "def foo(a: int, b: str) -> int")
        """
        sig = inspect.signature(self.func)
        func_obj = cast(Any, self.func)
        name = func_obj.__name__

        # Build parameter string
        params = []
        for param_name, param in sig.parameters.items():
            param_str = param_name
            if param.annotation != inspect.Parameter.empty:
                # Get annotation string
                anno = param.annotation
                if hasattr(anno, "__name__"):
                    param_str += f": {anno.__name__}"
                else:
                    param_str += f": {anno}"
            if param.default != inspect.Parameter.empty:
                param_str += f" = {param.default!r}"
            params.append(param_str)

        param_str = ", ".join(params)

        # Build return annotation
        return_anno = ""
        if sig.return_annotation != inspect.Signature.empty:
            ret = sig.return_annotation
            return_anno = f" -> {ret.__name__}" if hasattr(ret, "__name__") else f" -> {ret}"

        # Check if async
        prefix = "async def" if inspect.iscoroutinefunction(self.func) else "def"

        return f"{prefix} {name}({param_str}){return_anno}"

    def extract_docstring(self) -> str:
        """
        Extract docstring from function.

        Returns:
            Docstring or empty string if none
        """
        doc = inspect.getdoc(self.func)
        return doc or ""

    def extract_hypothesis_blocks(self) -> list[str]:
        """Extract fenced hypothesis blocks from the docstring."""

        doc = self.extract_docstring()
        if not doc:
            return []

        pattern = re.compile(r"```hypothesis\n(.*?)\n```", re.IGNORECASE | re.DOTALL)
        blocks = []
        for match in pattern.findall(doc):
            blocks.append(textwrap.dedent(match).strip())
        return blocks

    def extract_body_before_handled(self) -> str:
        """
        Extract function body code before VibesafeHandled() marker.

        Returns:
            Code string before yield/return VibesafeHandled()
        """
        source_lines = self.source.split("\n")

        # Find the function definition line
        func_def_line = 0
        func_obj = cast(Any, self.func)

        for i, line in enumerate(source_lines):
            if "def " in line and func_obj.__name__ in line:
                func_def_line = i
                break

        # Find body start (after docstring if present)
        body_start = func_def_line + 1

        # Skip docstring
        if func_obj.__doc__:
            in_docstring = False
            for i in range(func_def_line + 1, len(source_lines)):
                line = source_lines[i].strip()
                if '"""' in line or "'''" in line:
                    if not in_docstring:
                        in_docstring = True
                        if line.count('"""') == 2 or line.count("'''") == 2:
                            # Single line docstring
                            body_start = i + 1
                            break
                    else:
                        body_start = i + 1
                        break

        # Extract body lines before VibeCoded
        body_lines = []
        sentinel_markers = {
            "pass",
            "...",
            "return ...",
            "return Ellipsis",
        }

        for i in range(body_start, len(source_lines)):
            line = source_lines[i]
            stripped = line.strip()

            # Stop at raise VibeCoded() or legacy VibesafeHandled
            if (
                "VibeCoded" in stripped
                or "VibesafeHandled" in stripped
                or "VibeHandled" in stripped
            ):
                break

            if stripped in sentinel_markers:
                break

            # Only include lines with actual content (strip empty/whitespace)
            if stripped:
                body_lines.append(line)

        # Join and dedent
        body_code = "\n".join(body_lines)
        with contextlib.suppress(Exception):
            body_code = inspect.cleandoc(body_code)

        return body_code

    def extract_doctests(self) -> list[doctest.Example]:
        """
        Extract doctest examples from docstring.

        Returns:
            List of doctest Example objects
        """
        docstring = self.extract_docstring()
        if not docstring:
            return []

        parser = doctest.DocTestParser()
        try:
            examples = parser.get_examples(docstring)
            return examples
        except Exception:
            return []

    def extract_dependencies(self) -> dict[str, dict[str, str]]:
        """
        Extract static dependencies (names referenced in spec body and signature).

        Returns:
            Dictionary mapping name -> source code (if available)
        """
        body_code = self.extract_body_before_handled()
        # if not body_code or self.module is None:
        #     return {}
        if self.module is None:
            return {}

        names: set[str] = set()

        # 1. Scan body for name usage
        if body_code:
            try:
                body_tree = ast.parse(textwrap.dedent(body_code))

                class _NameCollector(ast.NodeVisitor):
                    def visit_Name(self, node: ast.Name) -> None:  # type: ignore[override]
                        if isinstance(node.ctx, ast.Load):
                            names.add(node.id)
                        self.generic_visit(node)

                _NameCollector().visit(body_tree)
            except SyntaxError:
                pass

        # 2. Scan signature for types and defaults
        func_obj = cast(Any, self.func)
        try:
            sig = inspect.signature(func_obj)
            for param in sig.parameters.values():
                if param.annotation != inspect.Parameter.empty:
                    self._extract_names_from_annotation(param.annotation, names)
                if param.default != inspect.Parameter.empty:
                    self._extract_names_from_value(param.default, names)
            if sig.return_annotation != inspect.Signature.empty:
                self._extract_names_from_annotation(sig.return_annotation, names)
        except (ValueError, TypeError):
            pass

        module_dict = getattr(self.module, "__dict__", {})
        dependencies: dict[str, dict[str, str]] = {}

        for name in sorted(names):
            if name == func_obj.__name__:
                continue
            if name in module_dict:
                obj = module_dict[name]

                # Vibesafe-Awareness: Check if it's a registered unit
                if getattr(obj, "__vibesafe_unit_id__", None):
                    # It's a vibesafe unit! Extract its interface.
                    # We use a fresh extractor to get the "real" signature/docstring
                    # from the underlying function if possible, or just use what we have.
                    # Note: The wrapper usually hides the real signature unless functools.wraps
                    # was used perfectly. But our @vibesafe decorator sets __vibesafe_unit_id__
                    # on the wrapper. We stored the original func in the registry, but
                    # here we might only have the wrapper.
                    # Actually, our decorator sets __vibesafe_unit_id__ on the wrapper.
                    # Let's try to get the real signature from the wrapper (functools.wraps).

                    # We can construct a "Spec Interface" string
                    try:
                        # Use SpecExtractor on the dependency to get its clean signature/docstring
                        # This works because SpecExtractor handles unwrapping/source finding
                        dep_extractor = SpecExtractor(obj)
                        source = f"{dep_extractor.extract_signature()}\n{dep_extractor.extract_docstring()}"
                        file_path = inspect.getfile(obj)
                    except Exception:
                        # Fallback to standard source extraction
                        try:
                            source = inspect.getsource(obj)
                            file_path = inspect.getfile(obj)
                        except (OSError, TypeError):
                            continue
                else:
                    # Standard object
                    try:
                        source = inspect.getsource(obj)
                        file_path = inspect.getfile(obj)
                    except (OSError, TypeError):
                        continue

                normalized_source = textwrap.dedent(source).strip()
                file_hash = ""
                try:
                    if file_path:
                        with open(file_path, "rb") as fh:
                            file_hash = hashlib.sha256(fh.read()).hexdigest()
                except Exception:
                    file_hash = ""

                dependencies[name] = {
                    "source": normalized_source,
                    "path": str(Path(file_path).resolve()) if file_path else "",
                    "file_hash": file_hash,
                }

        return dependencies

    def _extract_names_from_annotation(self, annotation: Any, names: set[str]) -> None:
        """Helper to extract names from type annotations."""
        if isinstance(annotation, str):
            # Forward reference string - simple heuristic
            # We just grab the first word if it looks like a name
            # This is imperfect but covers common cases like "List[str]" -> "List"
            # Better: parse it as AST if possible
            try:
                tree = ast.parse(annotation)
                for node in ast.walk(tree):
                    if isinstance(node, ast.Name):
                        names.add(node.id)
            except SyntaxError:
                pass
            return

        if hasattr(annotation, "__name__"):
            names.add(annotation.__name__)

        # Handle generics like List[int] -> List, int
        if hasattr(annotation, "__origin__"):
            self._extract_names_from_annotation(annotation.__origin__, names)
        if hasattr(annotation, "__args__"):
            for arg in annotation.__args__:
                self._extract_names_from_annotation(arg, names)

    def _extract_names_from_value(self, value: Any, names: set[str]) -> None:
        """Helper to extract names from default values."""
        # If the value is a class or function, add its name
        if inspect.isclass(value) or inspect.isfunction(value):
            names.add(value.__name__)

    # --- Inference helpers -------------------------------------------------

    def _find_function_def(self) -> ast.AST | None:
        """Locate the AST node for the target function by name."""
        target = self._func_obj.__name__

        class _Finder(ast.NodeVisitor):
            def __init__(self) -> None:
                self.found: ast.AST | None = None

            def visit_FunctionDef(self, node: ast.FunctionDef) -> Any:  # type: ignore[override]
                if node.name == target and self.found is None:
                    self.found = node
                self.generic_visit(node)

            def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> Any:  # type: ignore[override]
                if node.name == target and self.found is None:
                    self.found = node
                self.generic_visit(node)

        finder = _Finder()
        finder.visit(self.tree)
        return finder.found

    def _infer_type(self) -> str:
        """
        Infer the surface type of the spec from decorators.

        Returns "http" if FastAPI-style route decorators are detected,
        otherwise "function". (CLI/Typer handling deferred.)
        """
        if not isinstance(self._func_def, (ast.FunctionDef, ast.AsyncFunctionDef)):
            return "function"

        if self._has_fastapi_route_decorator(self._func_def):
            return "http"

        return "function"

    def _has_fastapi_route_decorator(self, func_def: ast.AST) -> bool:
        """Return True if decorator list contains a FastAPI route-like call."""
        FASTAPI_METHODS = {
            "get",
            "post",
            "put",
            "delete",
            "patch",
            "options",
            "head",
            "route",
            "api_route",
            "websocket",
        }

        fastapi_imported = "fastapi" in self.source

        if not fastapi_imported:
            return False

        decorators = getattr(func_def, "decorator_list", [])
        for deco in decorators:
            call = deco if isinstance(deco, ast.Call) else None
            if call is None:
                continue

            dotted = self._dotted_name(call.func)
            if not dotted:
                continue

            # Require an explicit FastAPI-ish prefix to reduce false positives.
            if not (
                dotted.startswith(("app.", "router.", "api_router.", "fastapi.", "api."))
                or ".router." in dotted
                or ".app." in dotted
            ):
                continue

            if any(dotted.endswith(f".{method}") for method in FASTAPI_METHODS):
                return True

        return False

    def _dotted_name(self, node: ast.AST) -> str | None:
        """Convert Attribute/Name chain to dotted string."""
        parts: list[str] = []
        while isinstance(node, ast.Attribute):
            parts.append(node.attr)
            node = node.value
        if isinstance(node, ast.Name):
            parts.append(node.id)
        else:
            return None
        return ".".join(reversed(parts))

    def to_dict(self) -> dict[str, Any]:
        """
        Extract all spec components as a dictionary.

        Returns:
            Dictionary with signature, docstring, body, doctests, etc.
        """
        return {
            "signature": self.extract_signature(),
            "docstring": self.extract_docstring(),
            "body_before_handled": self.extract_body_before_handled(),
            "doctests": self.extract_doctests(),
            "hypothesis_blocks": self.extract_hypothesis_blocks(),
            "dependencies": self.extract_dependencies(),
            "source": self.source,
            "type": self.inferred_type,
        }


def extract_spec(func: Callable[..., Any]) -> dict[str, Any]:
    """
    Convenience function to extract spec from a function.

    Args:
        func: Function to extract spec from

    Returns:
        Dictionary with spec components
    """
    extractor = SpecExtractor(func)
    return extractor.to_dict()
