from __future__ import annotations

import ast
from dataclasses import dataclass
from typing import Any, Iterable

from ..schema import CallRef, FileFacts, ImportRef, Span, SymbolDef
from ..utils import normalize_whitespace


def _span_from_node(node: ast.AST) -> Span:
    # Python 3.8+ has end_lineno/end_col_offset for most nodes.
    lineno = getattr(node, "lineno", 1) or 1
    end_lineno = getattr(node, "end_lineno", lineno) or lineno
    col = getattr(node, "col_offset", 0) or 0
    end_col = getattr(node, "end_col_offset", col) or col
    return {
        "start_line": int(lineno),
        "end_line": int(end_lineno),
        "start_col": int(col),
        "end_col": int(end_col),
    }


def _dotted_name(expr: ast.AST) -> str | None:
    if isinstance(expr, ast.Name):
        return expr.id
    if isinstance(expr, ast.Attribute):
        base = _dotted_name(expr.value)
        if base is None:
            return None
        return f"{base}.{expr.attr}"
    return None


def _unparse(expr: ast.AST) -> str:
    try:
        return ast.unparse(expr)  # py>=3.9
    except Exception:
        return expr.__class__.__name__


def _first_doc_line(node: ast.AST) -> str:
    doc = ast.get_docstring(node, clean=True)
    if not doc:
        return ""
    return doc.splitlines()[0].strip()


def _format_args(args: ast.arguments) -> list[str]:
    out: list[str] = []

    def fmt_arg(a: ast.arg) -> str:
        if a.annotation is None:
            return a.arg
        return f"{a.arg}: {_unparse(a.annotation)}"

    # posonlyargs (py3.8+)
    for a in getattr(args, "posonlyargs", []):
        out.append(fmt_arg(a))
    if getattr(args, "posonlyargs", None):
        out.append("/")

    for a in args.args:
        out.append(fmt_arg(a))

    if args.vararg is not None:
        out.append(f"*{fmt_arg(args.vararg)}")
    elif args.kwonlyargs:
        out.append("*")

    for a in args.kwonlyargs:
        out.append(fmt_arg(a))

    if args.kwarg is not None:
        out.append(f"**{fmt_arg(args.kwarg)}")

    return out


def _symbol_id(path: str, qualname: str, kind: str, span: Span) -> str:
    return f"{path}::{kind}::{qualname}::{span['start_line']}:{span['end_line']}"


@dataclass
class ASTParser:
    """
    Python AST parser that extracts a "facts layer" from a file:
    - imports
    - class/function/method definitions with spans
    - lightweight call references (callee dotted name) with spans
    """

    def parse_python(self, code: str, path: str) -> FileFacts:
        facts: FileFacts = {
            "path": path,
            "language": "python",
            "imports": [],
            "symbols": [],
            "calls": [],
            "parse_errors": [],
        }

        try:
            tree = ast.parse(code, filename=path)
        except SyntaxError as e:
            facts["parse_errors"] = [f"SyntaxError: {e}"]
            return facts
        except Exception as e:
            facts["parse_errors"] = [f"ParseError: {e}"]
            return facts

        imports: list[ImportRef] = []
        symbols: list[SymbolDef] = []
        calls: list[CallRef] = []

        class_stack: list[str] = []
        current_symbol_stack: list[str] = []

        def qualname(name: str) -> str:
            if class_stack:
                return ".".join(class_stack + [name])
            return name

        def record_import(node: ast.AST) -> None:
            if isinstance(node, ast.Import):
                for alias in node.names:
                    imports.append(
                        {
                            "module": alias.name,
                            "alias": alias.asname or "",
                            "level": 0,
                            "span": _span_from_node(node),
                        }
                    )
            elif isinstance(node, ast.ImportFrom):
                mod = node.module or ""
                for alias in node.names:
                    imports.append(
                        {
                            "module": mod,
                            "name": alias.name,
                            "alias": alias.asname or "",
                            "level": int(node.level or 0),
                            "span": _span_from_node(node),
                        }
                    )

        def record_symbol(
            node: ast.AST, *, kind: str, name: str, signature: str, returns: str
        ) -> str:
            sp = _span_from_node(node)
            qn = qualname(name)
            sid = _symbol_id(path, qn, kind, sp)
            decorators: list[str] = []
            deco_nodes = getattr(node, "decorator_list", []) or []
            for d in deco_nodes:
                decorators.append(normalize_whitespace(_unparse(d)))

            sym: SymbolDef = {
                "symbol_id": sid,
                "kind": kind,  # type: ignore[typeddict-item]
                "name": name,
                "qualname": qn,
                "path": path,
                "span": sp,
                "signature": normalize_whitespace(signature),
                "decorators": decorators,
                "doc_firstline": _first_doc_line(node),
                "returns": returns,
            }
            symbols.append(sym)
            return sid

        def record_calls_in_body(body: Iterable[ast.stmt], caller_sid: str) -> None:
            for n in ast.walk(ast.Module(body=list(body), type_ignores=[])):  # type: ignore[arg-type]
                if isinstance(n, ast.Call):
                    callee = _dotted_name(n.func)
                    if not callee:
                        continue
                    calls.append(
                        {
                            "caller_symbol_id": caller_sid,
                            "callee": callee,
                            "span": _span_from_node(n),
                        }
                    )

        class Visitor(ast.NodeVisitor):
            def visit_Import(self, node: ast.Import) -> Any:
                record_import(node)

            def visit_ImportFrom(self, node: ast.ImportFrom) -> Any:
                record_import(node)

            def visit_ClassDef(self, node: ast.ClassDef) -> Any:
                # Class symbol
                signature = f"class {node.name}"
                sid = record_symbol(
                    node, kind="class", name=node.name, signature=signature, returns=""
                )
                class_stack.append(node.name)
                current_symbol_stack.append(sid)
                self.generic_visit(node)
                current_symbol_stack.pop()
                class_stack.pop()

            def visit_FunctionDef(self, node: ast.FunctionDef) -> Any:
                args = ", ".join(_format_args(node.args))
                ret = _unparse(node.returns) if node.returns is not None else ""
                signature = f"def {node.name}({args})"
                if ret:
                    signature += f" -> {ret}"
                kind = "method" if class_stack else "function"
                sid = record_symbol(
                    node, kind=kind, name=node.name, signature=signature, returns=ret
                )
                current_symbol_stack.append(sid)
                # Lightweight call refs inside function body
                record_calls_in_body(node.body, sid)
                self.generic_visit(node)
                current_symbol_stack.pop()

            def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> Any:
                args = ", ".join(_format_args(node.args))
                ret = _unparse(node.returns) if node.returns is not None else ""
                signature = f"async def {node.name}({args})"
                if ret:
                    signature += f" -> {ret}"
                kind = "method" if class_stack else "function"
                sid = record_symbol(
                    node, kind=kind, name=node.name, signature=signature, returns=ret
                )
                current_symbol_stack.append(sid)
                record_calls_in_body(node.body, sid)
                self.generic_visit(node)
                current_symbol_stack.pop()

        Visitor().visit(tree)

        facts["imports"] = imports
        facts["symbols"] = symbols
        facts["calls"] = calls
        return facts

