"""Cached, role-scoped structured-output conformance checks."""

from __future__ import annotations

from hashlib import sha256
import json
from typing import Any, Literal

from app.cad_agent.application.llm_contracts import (
    EmptyCommand,
    ImageObservation,
    MarkdownDocument,
    StatelessCandidateReview,
    StatelessGeometryConclusion,
    StatelessRollbackCheckpoint,
    StatelessTopologyRequest,
    compiled_requirements_schema,
    stateless_final_review_schema,
    stateless_next_action_schema,
)
from app.cad_agent.domain.feature_plan import FeaturePlan
from app.cad_agent.domain.operation_contract import fragment_schema
from app.cad_agent.domain.verifier_registry import default_registry
from app.cad_agent.ports import CadRuntime, ModelGateway


CapabilityRole = Literal["author", "reviewer"]


def conformance_tools(runtime: CadRuntime, *, role: CapabilityRole) -> list[dict[str, Any]]:
    if role == "reviewer":
        return [
            _tool("observe_images", ImageObservation.model_json_schema()),
            # Compatibility conformance probe; regular DAG execution never
            # calls a per-node reviewer.
            _tool("review_candidate", StatelessCandidateReview.model_json_schema()),
            _tool("review_final", stateless_final_review_schema(1)),
        ]
    atomic_ids = list(runtime.supported_atomic_ids())
    if not atomic_ids:
        raise RuntimeError("Runtime has no operations for conformance")
    tools = [
        _tool("write_requirements_document", MarkdownDocument.model_json_schema()),
        _tool("write_completion_target", MarkdownDocument.model_json_schema()),
        _tool("compile_requirements_spec", compiled_requirements_schema(default_registry().expected_one_of_schema(exclude_claim_kinds=frozenset({"coaxial", "coplanar"})), 1)),
        _tool("write_feature_plan", FeaturePlan.model_json_schema()),
        # Compatibility probe only; production v3.2 workflow never exposes it.
        _tool("write_modeling_plan", MarkdownDocument.model_json_schema()),
        _tool("inspect_topology", StatelessTopologyRequest.model_json_schema()),
        _tool("record_geometry_conclusion", StatelessGeometryConclusion.model_json_schema()),
        _tool("rollback_checkpoint", StatelessRollbackCheckpoint.model_json_schema()),
        _tool("complete_task", EmptyCommand.model_json_schema()),
    ]
    for atomic_id in atomic_ids:
        contract = runtime.operation_contract(atomic_id)
        tools.append(_tool(
            f"conformance_{atomic_id}",
            fragment_schema(contract, selector_tokens=["sel_conformance"], reference_tokens=["ref_conformance"]),
        ))
    return tools


def conformance_hash(tools: list[dict[str, Any]], *, role: CapabilityRole) -> str:
    payload = {"protocol": "cad.v3.2.feature-dag", "role": role, "tools": tools}
    return sha256(json.dumps(payload, ensure_ascii=True, sort_keys=True, separators=(",", ":")).encode("utf-8")).hexdigest()


def cached_model_capability(
    repository: Any,
    runtime: CadRuntime,
    *,
    provider_id: str,
    model_id: str,
    role: CapabilityRole,
) -> dict[str, Any] | None:
    tools = conformance_tools(runtime, role=role)
    schema_hash = conformance_hash(tools, role=role)
    cached = repository.model_capability(provider_id, model_id, schema_hash)
    if cached is None:
        return None
    return {"cached": True, "schema_hash": schema_hash, "role": role, **cached["report"]}


async def verify_model_capability(
    repository: Any,
    runtime: CadRuntime,
    models: ModelGateway,
    *,
    provider_id: str,
    model_id: str,
    role: CapabilityRole,
    force: bool = False,
) -> dict[str, Any]:
    tools = conformance_tools(runtime, role=role)
    schema_hash = conformance_hash(tools, role=role)
    cached = repository.model_capability(provider_id, model_id, schema_hash)
    if cached is not None and not force:
        return {"cached": True, "schema_hash": schema_hash, "role": role, **cached["report"]}
    report = await models.conformance(provider_id=provider_id, model_id=model_id, tools=tools)
    report = {"schema_hash": schema_hash, "role": role, "tool_count": len(tools), **report}
    if not report.get("probe_unavailable"):
        repository.record_model_capability(provider_id, model_id, schema_hash, report)
    return report


def _tool(name: str, parameters: dict[str, Any]) -> dict[str, Any]:
    return {
        "type": "function",
        "function": {
            "name": name,
            "description": "Structured output conformance probe. Return one schema-valid call.",
            "parameters": parameters,
        },
    }