"""Enforce strict typing contracts for src/auen."""

from __future__ import annotations

import argparse
import ast
import io
import tokenize
from dataclasses import dataclass
from pathlib import Path


@dataclass(frozen=True)
class Violation:
    kind: str
    path: str
    line: int
    column: int
    message: str


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


def _collect_annotation_nodes(tree: ast.AST) -> list[ast.AST]:
    nodes: list[ast.AST] = []

    class _Collector(ast.NodeVisitor):
        def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
            self._visit_signature(node)
            self.generic_visit(node)

        def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
            self._visit_signature(node)
            self.generic_visit(node)

        def _visit_signature(self, node: ast.FunctionDef | ast.AsyncFunctionDef) -> None:
            if node.returns is not None:
                nodes.append(node.returns)
            all_args = (
                list(node.args.posonlyargs)
                + list(node.args.args)
                + list(node.args.kwonlyargs)
            )
            if node.args.vararg is not None:
                all_args.append(node.args.vararg)
            if node.args.kwarg is not None:
                all_args.append(node.args.kwarg)
            for arg in all_args:
                if arg.annotation is not None:
                    nodes.append(arg.annotation)

        def visit_AnnAssign(self, node: ast.AnnAssign) -> None:
            nodes.append(node.annotation)
            if node.value is not None and _is_type_alias_marker(node.annotation):
                nodes.append(node.value)
            self.generic_visit(node)

    collector = _Collector()
    collector.visit(tree)

    type_alias_node_type = getattr(ast, "TypeAlias", None)
    if type_alias_node_type is not None:
        for node in ast.walk(tree):
            if isinstance(node, type_alias_node_type):
                value = getattr(node, "value", None)
                if isinstance(value, ast.AST):
                    nodes.append(value)

    return nodes


def _iter_forbidden_names(annotation_node: ast.AST) -> list[tuple[str, int, int]]:
    hits: list[tuple[str, int, int]] = []

    def _walk(node: ast.AST) -> None:
        for candidate in ast.walk(node):
            if isinstance(candidate, ast.Name) and candidate.id in {"Any", "object"}:
                hits.append(
                    (
                        candidate.id,
                        getattr(candidate, "lineno", 0),
                        getattr(candidate, "col_offset", 0),
                    )
                )
            elif isinstance(candidate, ast.Attribute) and candidate.attr in {
                "Any",
                "object",
            }:
                hits.append(
                    (
                        candidate.attr,
                        getattr(candidate, "lineno", 0),
                        getattr(candidate, "col_offset", 0),
                    )
                )
            elif isinstance(candidate, ast.Constant) and isinstance(candidate.value, str):
                try:
                    parsed = ast.parse(candidate.value, mode="eval")
                except SyntaxError:
                    continue
                _walk(parsed)

    _walk(annotation_node)
    return hits


def _load_allowlist(path: Path) -> set[tuple[str, str, int]]:
    entries: set[tuple[str, str, int]] = set()
    if not path.exists():
        return entries

    for raw_line in path.read_text(encoding="utf-8").splitlines():
        line = raw_line.strip()
        if not line or line.startswith("#"):
            continue
        parts = line.split(":", 3)
        if len(parts) < 3:
            msg = f"Invalid allowlist entry: {raw_line!r}"
            raise ValueError(msg)
        kind, rel_path, line_str = parts[:3]
        entries.add((kind, rel_path, int(line_str)))
    return entries


def _iter_type_ignores(source: str) -> list[tuple[int, int]]:
    ignores: list[tuple[int, int]] = []
    for token in tokenize.generate_tokens(io.StringIO(source).readline):
        if token.type != tokenize.COMMENT:
            continue
        if "type: ignore" in token.string:
            ignores.append(token.start)
    return ignores


def scan_sources(
    *,
    repo_root: Path,
    src_root: Path,
    allowlist_path: Path,
) -> list[Violation]:
    allowlist = _load_allowlist(allowlist_path)
    violations: list[Violation] = []

    for file_path in sorted(src_root.rglob("*.py")):
        rel_path = file_path.relative_to(repo_root).as_posix()
        source = file_path.read_text(encoding="utf-8")
        tree = ast.parse(source, filename=rel_path, type_comments=True)

        for annotation_node in _collect_annotation_nodes(tree):
            for forbidden_name, line, column in _iter_forbidden_names(annotation_node):
                kind = f"annotation-{forbidden_name.lower()}"
                if (kind, rel_path, line) in allowlist:
                    continue
                violations.append(
                    Violation(
                        kind=kind,
                        path=rel_path,
                        line=line,
                        column=column,
                        message=(
                            f"{forbidden_name} is forbidden in type annotations under "
                            "src/auen"
                        ),
                    )
                )

        for line, column in _iter_type_ignores(source):
            if ("type-ignore", rel_path, line) in allowlist:
                continue
            violations.append(
                Violation(
                    kind="type-ignore",
                    path=rel_path,
                    line=line,
                    column=column,
                    message="type: ignore is forbidden under src/auen",
                )
            )

    return violations


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--repo-root",
        type=Path,
        default=Path(__file__).resolve().parents[1],
        help="Repository root (default: inferred from script path)",
    )
    parser.add_argument(
        "--src-root",
        type=Path,
        default=Path("src/auen"),
        help="Source root to check (default: src/auen)",
    )
    parser.add_argument(
        "--allowlist",
        type=Path,
        default=Path("scripts/typing_contract_allowlist.txt"),
        help="Allowlist file with kind:path:line entries",
    )
    args = parser.parse_args()

    repo_root = args.repo_root.resolve()
    src_root = (repo_root / args.src_root).resolve()
    allowlist_path = (repo_root / args.allowlist).resolve()

    violations = scan_sources(
        repo_root=repo_root,
        src_root=src_root,
        allowlist_path=allowlist_path,
    )

    if not violations:
        return 0

    for violation in violations:
        print(
            f"{violation.path}:{violation.line}:{violation.column + 1} "
            f"[{violation.kind}] {violation.message}"
        )
    return 1


if __name__ == "__main__":
    raise SystemExit(main())
