"""Canonical schemas for every v3 LLM state-changing command.

Providers may parse structured output first, but this module validates the raw
tool arguments a second time before a command reaches a handler.
"""

from __future__ import annotations

from copy import deepcopy
import hashlib
import json
import math
from typing import Annotated, Any, Literal, TypeVar

from jsonschema import Draft202012Validator
from pydantic import BaseModel, ConfigDict, Field, JsonValue, ValidationError, model_validator

from app.cad_agent.domain.errors import ErrorCode, WorkflowError


class StrictDto(BaseModel):
    model_config = ConfigDict(extra="forbid", strict=True, str_strip_whitespace=True)


ShortText = Annotated[str, Field(min_length=1, max_length=360)]
Identifier = Annotated[str, Field(pattern=r"^[a-z][a-z0-9_:-]{0,95}$")]


class AcceptanceClaimInput(StrictDto):
    claim_kind: Identifier
    expected: dict[str, JsonValue] = Field(min_length=0, max_length=24)


class RequirementInput(StrictDto):
    source_ids: list[Identifier] = Field(min_length=1, max_length=32)
    statement: Annotated[str, Field(min_length=1, max_length=1000)]
    assumptions: list[Annotated[str, Field(min_length=1, max_length=360)]] = Field(default_factory=list, max_length=16)
    acceptance_claims: list[AcceptanceClaimInput] = Field(min_length=1, max_length=16)

    @model_validator(mode="after")
    def _source_ids_are_unique(self) -> "RequirementInput":
        if len(self.source_ids) != len(set(self.source_ids)):
            raise ValueError("source_ids must not contain duplicates")
        return self


class RequirementsDraftBatch(StrictDto):
    items: list[RequirementInput] = Field(min_length=1, max_length=8)


class RequirementsPatch(StrictDto):
    target_draft_id: Identifier = Field(description="Current server-assigned draft ID to patch. This field belongs inside one patches[] entry.")
    op: Literal["replace", "remove"] = Field(description="replace supplies a complete replacement item; remove supplies a reason instead.")
    item: RequirementInput | None = Field(default=None, description="Complete replacement RequirementInput for op=replace. It must not include draft_id because the server preserves that ID.")
    reason: str | None = Field(default=None, min_length=1, max_length=360, description="Required only for op=remove; explains why the current draft item is removed.")

    @model_validator(mode="after")
    def _complete_patch(self) -> "RequirementsPatch":
        if self.op == "replace" and self.item is None:
            raise ValueError("replace requires a complete item")
        if self.op == "remove" and (self.item is not None or self.reason is None):
            raise ValueError("remove requires a reason and forbids item")
        return self


class RequirementsPatchBatch(StrictDto):
    patches: list[RequirementsPatch] = Field(min_length=1, max_length=8, description="Patch entries. Example shape: {\"patches\":[{\"target_draft_id\":\"draft_001\",\"op\":\"replace\",\"item\":{...}}]}.")


class EmptyCommand(StrictDto):
    pass


class ReviewNormalization(StrictDto):
    rule_id: Literal["full_circle_equal_spacing"]
    count: int = Field(ge=2, le=1024)
    declared_spacing_degrees: float = Field(gt=0, le=360)
    full_circle: bool


class ReviewFinding(StrictDto):
    draft_id: Identifier
    source_ids: list[Identifier] = Field(min_length=1, max_length=32)
    finding_type: Literal[
        "missing_source_semantics",
        "claim_mismatch",
        "verification_gap",
        "derivable_conflict",
        "ambiguous_conflict",
    ]
    description: Annotated[str, Field(min_length=1, max_length=1000)]
    question: str | None = Field(default=None, min_length=1, max_length=360)
    normalization: ReviewNormalization | None = None

    @model_validator(mode="after")
    def _finding_payload_matches_type(self) -> "ReviewFinding":
        if len(self.source_ids) != len(set(self.source_ids)):
            raise ValueError("source_ids must not contain duplicates")
        if self.finding_type == "ambiguous_conflict" and self.question is None:
            raise ValueError("ambiguous_conflict requires an answerable question")
        if self.finding_type == "derivable_conflict" and self.normalization is None:
            raise ValueError("derivable_conflict requires structured normalization data")
        if self.finding_type != "derivable_conflict" and self.normalization is not None:
            raise ValueError("normalization is allowed only for derivable_conflict")
        return self


class RequirementsReview(StrictDto):
    findings: list[ReviewFinding] = Field(default_factory=list, max_length=128)


class NextAction(StrictDto):
    working_head: Annotated[str, Field(pattern=r"^[a-z0-9_:-]{5,192}$")]
    intent: ShortText
    requirement_ids: list[Identifier] = Field(min_length=1, max_length=5)
    atomic_id: Identifier
    expected_change: ShortText

    @model_validator(mode="after")
    def _requirement_ids_are_unique(self) -> "NextAction":
        if len(self.requirement_ids) != len(set(self.requirement_ids)):
            raise ValueError("requirement_ids must not contain duplicates")
        return self


class TopologyRequest(StrictDto):
    working_head: Annotated[str, Field(pattern=r"^[a-z0-9_:-]{5,192}$")]
    kind: Literal["face", "edge", "vertex", "plane", "axis", "body"] | None = None
    limit: int = Field(default=16, ge=1, le=64)


class OperationContractRequest(StrictDto):
    working_head: Annotated[str, Field(pattern=r"^[a-z0-9_:-]{5,192}$")]
    atomic_id: Identifier


class GeometryConclusion(StrictDto):
    working_head: Annotated[str, Field(pattern=r"^[a-z0-9_:-]{5,192}$")]
    evidence_refs: list[Identifier] = Field(min_length=1, max_length=16)
    root_cause: Annotated[str, Field(min_length=1, max_length=360)]
    decision: Literal["return_to_action_selection", "rollback"]
    corrective_intent: str | None = Field(default=None, min_length=1, max_length=360)

    @model_validator(mode="after")
    def _evidence_refs_are_unique(self) -> "GeometryConclusion":
        if len(self.evidence_refs) != len(set(self.evidence_refs)):
            raise ValueError("evidence_refs must not contain duplicates")
        return self


class RollbackCheckpoint(StrictDto):
    working_head: Annotated[str, Field(pattern=r"^[a-z0-9_:-]{5,192}$")]
    checkpoint_token: Identifier
    reason: Annotated[str, Field(min_length=1, max_length=360)]


class ClaimCoverage(StrictDto):
    claim_id: Identifier
    status: Literal["pass", "pending", "fail", "not_applicable"]
    evidence_refs: list[Identifier] = Field(default_factory=list, max_length=16)

    @model_validator(mode="after")
    def _evidence_refs_are_unique(self) -> "ClaimCoverage":
        if len(self.evidence_refs) != len(set(self.evidence_refs)):
            raise ValueError("evidence_refs must not contain duplicates")
        return self


class CandidateReview(StrictDto):
    candidate_id: Identifier = Field(description="Server-issued candidate ID from the review facts.")
    working_head: Annotated[str, Field(pattern=r"^[a-z0-9_:-]{5,192}$", description="Current server-issued working head from the review facts.")]
    verdict: Literal["accept", "reject"] = Field(description="Required independent decision. Set accept only when the supplied candidate evidence supports every covered claim; otherwise set reject.")
    claim_coverage: list[ClaimCoverage] = Field(min_length=1, max_length=128, description="Required coverage decision for every claim ID in the supplied candidate facts.")
    evidence: list[Annotated[str, Field(min_length=1, max_length=360)]] = Field(default_factory=list, max_length=16)
    issues: list[Annotated[str, Field(min_length=1, max_length=360)]] = Field(default_factory=list, max_length=16)


class FinalReview(StrictDto):
    working_head: Annotated[str, Field(pattern=r"^[a-z0-9_:-]{5,192}$", description="Current server-issued working head from the final review facts.")]
    verdict: Literal["pass", "repair"] = Field(description="Required independent final decision. Set pass only when the supplied evidence supports every claim; otherwise set repair.")
    claim_coverage: list[ClaimCoverage] = Field(min_length=1, max_length=128, description="Required coverage decision for every claim ID in the supplied final-review facts.")
    evidence: list[Annotated[str, Field(min_length=1, max_length=360)]] = Field(default_factory=list, max_length=16)
    issues: list[Annotated[str, Field(min_length=1, max_length=360)]] = Field(default_factory=list, max_length=16)


def requirements_draft_schema(claim_one_of: dict[str, Any], source_ids: list[str]) -> dict[str, Any]:
    """Bind claim and source enums for the current requirements snapshot."""
    schema = RequirementsDraftBatch.model_json_schema()
    requirement = schema.get("$defs", {}).get("RequirementInput")
    if isinstance(requirement, dict):
        properties = requirement.get("properties", {})
        source_items = properties.get("source_ids", {}).get("items") if isinstance(properties, dict) and isinstance(properties.get("source_ids"), dict) else None
        if isinstance(source_items, dict):
            source_items["enum"] = source_ids
        claims = requirement.get("properties", {}).get("acceptance_claims")
        if isinstance(claims, dict):
            claims["items"] = deepcopy(claim_one_of)
    return schema


def requirements_patch_schema(claim_one_of: dict[str, Any], source_ids: list[str], draft_ids: list[str]) -> dict[str, Any]:
    schema = RequirementsPatchBatch.model_json_schema()
    definitions = schema.get("$defs", {})
    requirement = definitions.get("RequirementInput") if isinstance(definitions, dict) else None
    if isinstance(requirement, dict):
        properties = requirement.get("properties", {})
        source_items = properties.get("source_ids", {}).get("items") if isinstance(properties, dict) and isinstance(properties.get("source_ids"), dict) else None
        if isinstance(source_items, dict):
            source_items["enum"] = source_ids
        claims = requirement.get("properties", {}).get("acceptance_claims")
        if isinstance(claims, dict):
            claims["items"] = deepcopy(claim_one_of)
    patch = definitions.get("RequirementsPatch") if isinstance(definitions, dict) else None
    if isinstance(patch, dict):
        target = patch.get("properties", {}).get("target_draft_id")
        if isinstance(target, dict):
            target["enum"] = draft_ids
    return schema


def requirements_review_schema(source_ids: list[str], draft_ids: list[str]) -> dict[str, Any]:
    schema = RequirementsReview.model_json_schema()
    definitions = schema.get("$defs", {})
    finding = definitions.get("ReviewFinding") if isinstance(definitions, dict) else None
    if isinstance(finding, dict):
        properties = finding.get("properties", {})
        if isinstance(properties.get("draft_id"), dict):
            properties["draft_id"] = {"enum": draft_ids}
        source_items = properties.get("source_ids", {}).get("items") if isinstance(properties.get("source_ids"), dict) else None
        if isinstance(source_items, dict):
            source_items.clear()
            source_items.update({"enum": source_ids})
    return schema


def next_action_schema(working_head: str, requirement_ids: list[str], atomic_ids: list[str]) -> dict[str, Any]:
    schema = NextAction.model_json_schema()
    properties = schema.get("properties", {})
    if isinstance(properties, dict):
        properties["working_head"] = {"const": working_head}
        if isinstance(properties.get("requirement_ids"), dict):
            properties["requirement_ids"]["items"] = {"enum": requirement_ids}
        properties["atomic_id"] = {"enum": atomic_ids}
    return schema


def candidate_review_schema(candidate_id: str, working_head: str, claim_ids: list[str]) -> dict[str, Any]:
    """Bind an independent candidate review to immutable candidate facts."""
    schema = CandidateReview.model_json_schema()
    properties = schema.get("properties", {})
    if isinstance(properties, dict):
        properties["candidate_id"] = {"const": candidate_id}
        properties["working_head"] = {"const": working_head}
    _bind_claim_coverage_ids(schema, claim_ids)
    return schema


def final_review_schema(working_head: str, claim_ids: list[str]) -> dict[str, Any]:
    """Bind final review output to the currently reviewable revision."""
    schema = FinalReview.model_json_schema()
    properties = schema.get("properties", {})
    if isinstance(properties, dict):
        properties["working_head"] = {"const": working_head}
    _bind_claim_coverage_ids(schema, claim_ids)
    return schema


def topology_request_schema(working_head: str) -> dict[str, Any]:
    schema = TopologyRequest.model_json_schema()
    properties = schema.get("properties", {})
    if isinstance(properties, dict):
        properties["working_head"] = {"const": working_head}
    return schema


def operation_contract_request_schema(working_head: str, atomic_id: str) -> dict[str, Any]:
    schema = OperationContractRequest.model_json_schema()
    properties = schema.get("properties", {})
    if isinstance(properties, dict):
        properties["working_head"] = {"const": working_head}
        properties["atomic_id"] = {"const": atomic_id}
    return schema


def geometry_conclusion_schema(working_head: str, evidence_refs: list[str]) -> dict[str, Any]:
    """Bind a diagnostic conclusion to evidence generated for this head."""
    schema = GeometryConclusion.model_json_schema()
    properties = schema.get("properties", {})
    if isinstance(properties, dict):
        properties["working_head"] = {"const": working_head}
        evidence = properties.get("evidence_refs")
        if isinstance(evidence, dict):
            evidence["items"] = {"enum": evidence_refs}
    return schema


def rollback_checkpoint_schema(working_head: str, checkpoint_tokens: list[str]) -> dict[str, Any]:
    """Bind a rollback request to immutable checkpoints in the active lineage."""
    schema = RollbackCheckpoint.model_json_schema()
    properties = schema.get("properties", {})
    if isinstance(properties, dict):
        properties["working_head"] = {"const": working_head}
        properties["checkpoint_token"] = {"enum": checkpoint_tokens}
    return schema


def _bind_claim_coverage_ids(schema: dict[str, Any], claim_ids: list[str]) -> None:
    definitions = schema.get("$defs", {})
    coverage = definitions.get("ClaimCoverage") if isinstance(definitions, dict) else None
    if not isinstance(coverage, dict):
        return
    properties = coverage.get("properties", {})
    if isinstance(properties, dict):
        properties["claim_id"] = {"enum": claim_ids}


Dto = TypeVar("Dto", bound=StrictDto)


def raw_arguments_hash(raw_arguments_json: str) -> str:
    return hashlib.sha256(raw_arguments_json.encode("utf-8")).hexdigest()


def json_depth(value: Any, current: int = 0) -> int:
    if isinstance(value, dict):
        return max([current, *(json_depth(item, current + 1) for item in value.values())])
    if isinstance(value, list):
        return max([current, *(json_depth(item, current + 1) for item in value)])
    return current


def canonical_json_object(raw_arguments_json: str, *, max_bytes: int = 48_000, max_depth: int = 16) -> dict[str, Any] | WorkflowError:
    """Bound and parse raw arguments before any schema-specific validation."""
    if len(raw_arguments_json.encode("utf-8")) > max_bytes:
        return WorkflowError(ErrorCode.AUTHOR_FORMAT_INVALID, "Tool arguments exceed the byte limit.")
    try:
        value = json.loads(raw_arguments_json)
    except json.JSONDecodeError as error:
        return WorkflowError(ErrorCode.AUTHOR_FORMAT_INVALID, "Tool arguments are not valid JSON.", field_errors=({"path": "/", "message": error.msg},))
    if not isinstance(value, dict):
        return WorkflowError(ErrorCode.AUTHOR_FORMAT_INVALID, "Tool arguments must be a JSON object.")
    if json_depth(value) > max_depth:
        return WorkflowError(ErrorCode.AUTHOR_FORMAT_INVALID, "Tool arguments exceed the nesting-depth limit.")
    if _contains_non_finite_number(value):
        return WorkflowError(ErrorCode.AUTHOR_FORMAT_INVALID, "Tool arguments must not contain NaN or infinite numbers.")
    return value


def _contains_non_finite_number(value: Any) -> bool:
    if isinstance(value, float):
        return not math.isfinite(value)
    if isinstance(value, dict):
        return any(_contains_non_finite_number(item) for item in value.values())
    if isinstance(value, list):
        return any(_contains_non_finite_number(item) for item in value)
    return False


def canonical_validate_schema(raw_arguments_json: str, schema: dict[str, Any]) -> WorkflowError | None:
    """Revalidate raw arguments against the current dynamic JSON Schema."""
    value = canonical_json_object(raw_arguments_json)
    if isinstance(value, WorkflowError):
        return value
    errors = [
        {"path": "/" + "/".join(str(part) for part in error.absolute_path), "message": error.message}
        for error in sorted(Draft202012Validator(schema).iter_errors(value), key=lambda item: (list(item.absolute_path), item.message))
    ]
    if errors:
        return WorkflowError(ErrorCode.AUTHOR_FORMAT_INVALID, "Tool arguments do not match the active dynamic schema.", field_errors=tuple(errors))
    return None


def canonical_validate(raw_arguments_json: str, model: type[Dto], *, max_bytes: int = 48_000, max_depth: int = 16) -> Dto | WorkflowError:
    """Parse raw tool arguments once and return field-level DTO failures safely."""
    value = canonical_json_object(raw_arguments_json, max_bytes=max_bytes, max_depth=max_depth)
    if isinstance(value, WorkflowError):
        return value
    try:
        return model.model_validate(value)
    except ValidationError as error:
        fields = tuple({"path": "/" + "/".join(str(part) for part in issue["loc"]), "message": issue["msg"]} for issue in error.errors())
        return WorkflowError(ErrorCode.AUTHOR_FORMAT_INVALID, "Tool arguments do not match the canonical schema.", field_errors=fields)


def validate_one_tool_call(tool_calls: list[dict[str, Any]], allowed_name: str) -> tuple[str, str] | WorkflowError:
    """Require exactly one known tool call; no provider parser is trusted."""
    if len(tool_calls) != 1:
        return WorkflowError(ErrorCode.AUTHOR_FORMAT_INVALID, "Exactly one tool call is required.")
    function = tool_calls[0].get("function") if isinstance(tool_calls[0], dict) else None
    name = str(function.get("name") or "") if isinstance(function, dict) else ""
    raw = str(function.get("arguments") or "") if isinstance(function, dict) else ""
    if name != allowed_name:
        return WorkflowError(ErrorCode.AUTHOR_FORMAT_INVALID, "The returned tool is not allowed in this workflow state.", details={"expected_tool": allowed_name, "actual_tool": name})
    return name, raw