from __future__ import annotations

import argparse
import ast
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable


REQUIRED_HANDLERS: dict[str, tuple[str, str]] = {
    "get_status": ("query", "QUERY_GET_STATUS"),
    "get_events": ("query", "QUERY_GET_EVENTS"),
}
MIXIN_COMPATIBLE_BASES = {"PercepticWorkflow", "AgentWorkflowBase"}


@dataclass(frozen=True)
class Diagnostic:
    path: Path
    line: int
    col: int
    code: str
    message: str


def _base_name(base: ast.expr) -> str | None:
    if isinstance(base, ast.Name):
        return base.id
    if isinstance(base, ast.Attribute):
        return base.attr
    if isinstance(base, ast.Subscript):
        return _base_name(base.value)
    return None


def _is_workflow_decorator_call(decorator: ast.expr) -> tuple[str | None, ast.Call | None]:
    if not isinstance(decorator, ast.Call):
        return None, None
    func = decorator.func
    if (
        isinstance(func, ast.Attribute)
        and isinstance(func.value, ast.Name)
        and func.value.id == "workflow"
        and func.attr in {"query", "update"}
    ):
        return func.attr, decorator
    return None, None


def _extract_name_kwarg(call: ast.Call) -> str | None:
    for keyword in call.keywords:
        if keyword.arg != "name":
            continue

        value = keyword.value
        if isinstance(value, ast.Constant) and isinstance(value.value, str):
            return value.value
        if (
            isinstance(value, ast.Attribute)
            and isinstance(value.value, ast.Name)
            and value.value.id == "PercepticWorkflow"
        ):
            return value.attr
    return None


def _method_decorator_info(node: ast.AST) -> tuple[str | None, str | None]:
    if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
        return None, None

    for decorator in node.decorator_list:
        handler_kind, call = _is_workflow_decorator_call(decorator)
        if handler_kind is None or call is None:
            continue

        return handler_kind, _extract_name_kwarg(call)

    return None, None


def _class_diagnostics(path: Path, class_node: ast.ClassDef) -> list[Diagnostic]:
    method_nodes = {
        node.name: node for node in class_node.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
    }

    diagnostics: list[Diagnostic] = []
    for method_name, (expected_kind, expected_name_constant) in REQUIRED_HANDLERS.items():
        method_node = method_nodes.get(method_name)
        if method_node is None:
            continue

        actual_kind, actual_name = _method_decorator_info(method_node)
        if actual_kind is None:
            diagnostics.append(
                Diagnostic(
                    path=path,
                    line=method_node.lineno,
                    col=method_node.col_offset,
                    code="PWF001",
                    message=(
                        f"Method '{method_name}' is missing a Temporal decorator. "
                        f"Expected @workflow.{expected_kind}(name=PercepticWorkflow.{expected_name_constant})."
                    ),
                )
            )
            continue

        if actual_kind != expected_kind:
            diagnostics.append(
                Diagnostic(
                    path=path,
                    line=method_node.lineno,
                    col=method_node.col_offset,
                    code="PWF002",
                    message=(
                        f"Method '{method_name}' uses @workflow.{actual_kind}, but "
                        f"@workflow.{expected_kind} is required."
                    ),
                )
            )

        if actual_name != expected_name_constant:
            actual_name_label = actual_name if actual_name is not None else "<missing>"
            diagnostics.append(
                Diagnostic(
                    path=path,
                    line=method_node.lineno,
                    col=method_node.col_offset,
                    code="PWF003",
                    message=(
                        f"Method '{method_name}' uses handler name '{actual_name_label}', but expected "
                        f"PercepticWorkflow.{expected_name_constant}."
                    ),
                )
            )

    return diagnostics


def _mixin_candidate_class_names(class_nodes: list[ast.ClassDef]) -> set[str]:
    candidate_names: set[str] = set()
    remaining = {class_node.name: class_node for class_node in class_nodes}

    while remaining:
        progressed = False
        for class_name, class_node in list(remaining.items()):
            base_names = {_base_name(base) for base in class_node.bases}
            base_names.discard(None)

            if base_names & MIXIN_COMPATIBLE_BASES:
                candidate_names.add(class_name)
                del remaining[class_name]
                progressed = True
                continue

            if base_names & candidate_names:
                candidate_names.add(class_name)
                del remaining[class_name]
                progressed = True

        if not progressed:
            break

    return candidate_names


def _iter_python_files(paths: Iterable[Path]) -> Iterable[Path]:
    for path in paths:
        if path.is_dir():
            for child in path.rglob("*.py"):
                if ".venv" in child.parts:
                    continue
                yield child
        elif path.suffix == ".py" and path.exists():
            yield path


def check_paths(paths: Iterable[Path]) -> list[Diagnostic]:
    diagnostics: list[Diagnostic] = []
    for file_path in _iter_python_files(paths):
        try:
            source = file_path.read_text(encoding="utf-8")
        except OSError:
            continue

        try:
            tree = ast.parse(source, filename=str(file_path))
        except SyntaxError:
            continue

        class_nodes = [node for node in tree.body if isinstance(node, ast.ClassDef)]
        candidate_names = _mixin_candidate_class_names(class_nodes)
        for node in class_nodes:
            if node.name in candidate_names:
                diagnostics.extend(_class_diagnostics(file_path, node))
    return diagnostics


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Check PercepticWorkflow overrides for Temporal decorator correctness."
    )
    parser.add_argument("paths", nargs="*", default=["src"], help="Files or directories to scan.")
    args = parser.parse_args()

    input_paths = [Path(path).resolve() for path in args.paths]
    diagnostics = check_paths(input_paths)
    for diagnostic in diagnostics:
        print(f"{diagnostic.path}:{diagnostic.line}:{diagnostic.col}: {diagnostic.code} {diagnostic.message}")

    return 1 if diagnostics else 0


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