path: root/src/gap_pipeline/kernel_models.py

"""Typed contracts for the literal five-stage GAP kernel pipeline."""

from __future__ import annotations

import json
import re
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator

from .models import SCHEMA_VERSION


class StrictModel(BaseModel):
    model_config = ConfigDict(extra="forbid")


class ProofNode(StrictModel):
    node_id: str
    claim: str
    dependencies: list[str] = Field(default_factory=list)

    @model_validator(mode="after")
    def validate_content(self) -> "ProofNode":
        if not re.fullmatch(r"n[1-9][0-9]*", self.node_id):
            raise ValueError("proof node IDs must be n1, n2, ...")
        if not self.claim.strip():
            raise ValueError("proof-node claim must be non-empty")
        if len(self.dependencies) != len(set(self.dependencies)):
            raise ValueError(f"{self.node_id} has duplicate dependencies")
        return self


class ProofDAG(StrictModel):
    nodes: list[ProofNode] = Field(min_length=1)
    terminal_node_id: str

    def node_ids(self) -> list[str]:
        return [node.node_id for node in self.nodes]

    def leaf_node_ids(self) -> list[str]:
        """Return source leaves: nodes with no prerequisite dependencies."""

        return [node.node_id for node in self.nodes if not node.dependencies]

    @model_validator(mode="after")
    def validate_graph(self) -> "ProofDAG":
        node_ids = self.node_ids()
        expected = [f"n{index}" for index in range(1, len(node_ids) + 1)]
        if node_ids != expected:
            raise ValueError("proof DAG nodes must be topologically ordered n1 through nN")
        known: set[str] = set()
        for node in self.nodes:
            if node.node_id in node.dependencies:
                raise ValueError(f"{node.node_id} cannot depend on itself")
            unknown = set(node.dependencies) - known
            if unknown:
                raise ValueError(
                    f"{node.node_id} has non-prior dependencies {sorted(unknown)}"
                )
            known.add(node.node_id)
        if self.terminal_node_id not in known:
            raise ValueError("terminal node is not present in the DAG")

        by_id = {node.node_id: node for node in self.nodes}
        ancestors: set[str] = set()
        pending = [self.terminal_node_id]
        while pending:
            node_id = pending.pop()
            if node_id in ancestors:
                continue
            ancestors.add(node_id)
            pending.extend(by_id[node_id].dependencies)
        if ancestors != known:
            raise ValueError("every proof node must contribute to the terminal node")
        return self


class MethodNode(StrictModel):
    node_id: str
    method_label: str

    @model_validator(mode="after")
    def validate_content(self) -> "MethodNode":
        if not self.method_label.strip():
            raise ValueError("method label must be non-empty")
        return self


class MethodPlan(StrictModel):
    nodes: list[MethodNode] = Field(min_length=1)

    def validate_against(self, dag: ProofDAG) -> "MethodPlan":
        if [node.node_id for node in self.nodes] != dag.node_ids():
            raise ValueError("method-plan IDs must exactly match proof-DAG IDs")
        return self


class ReplacementChange(StrictModel):
    slot_id: str
    source_node_id: str
    description: str
    original_value: str
    replacement_value: str
    guard_condition: str
    guard_justification: str

    @model_validator(mode="after")
    def validate_change(self) -> "ReplacementChange":
        text_fields = [
            self.slot_id,
            self.source_node_id,
            self.description,
            self.original_value,
            self.replacement_value,
            self.guard_condition,
            self.guard_justification,
        ]
        if any(not value.strip() for value in text_fields):
            raise ValueError("replacement fields must be non-empty")
        if not re.fullmatch(r"slot[1-9][0-9]*", self.slot_id):
            raise ValueError("replacement slot IDs must be slot1, slot2, ...")
        if self.original_value.strip() == self.replacement_value.strip():
            raise ValueError("replacement must differ from the original value")
        return self


class ReplacementPlan(StrictModel):
    changes: list[ReplacementChange] = Field(min_length=1)
    closure_statement: str

    @model_validator(mode="after")
    def validate_content(self) -> "ReplacementPlan":
        slot_ids = [change.slot_id for change in self.changes]
        if len(slot_ids) != len(set(slot_ids)):
            raise ValueError("replacement slot IDs must be unique")
        if not self.closure_statement.strip():
            raise ValueError("replacement plan must state its closure guarantee")
        return self

    def validate_against(self, dag: ProofDAG) -> "ReplacementPlan":
        known = set(dag.node_ids())
        unknown = {
            change.source_node_id
            for change in self.changes
            if change.source_node_id not in known
        }
        if unknown:
            raise ValueError(f"replacement plan references unknown nodes {sorted(unknown)}")
        leaves = set(dag.leaf_node_ids())
        non_leaf = {
            change.source_node_id
            for change in self.changes
            if change.source_node_id not in leaves
        }
        if non_leaf:
            raise ValueError(
                "replacement plan must target source leaf nodes; "
                f"received {sorted(non_leaf)}"
            )
        return self

    def validate_repair_of(
        self,
        previous: "ReplacementPlan",
    ) -> "ReplacementPlan":
        current_targets = [
            (change.slot_id, change.source_node_id) for change in self.changes
        ]
        previous_targets = [
            (change.slot_id, change.source_node_id)
            for change in previous.changes
        ]
        if current_targets != previous_targets:
            raise ValueError(
                "repair must preserve replacement slot IDs and source leaf nodes"
            )
        return self


class DiffusedProofNode(StrictModel):
    node_id: str
    dependencies: list[str] = Field(default_factory=list)
    method_label: str
    instantiated_claim: str
    justification: str

    @model_validator(mode="after")
    def validate_content(self) -> "DiffusedProofNode":
        if not self.method_label.strip():
            raise ValueError("diffused method label must be non-empty")
        if not self.instantiated_claim.strip() or not self.justification.strip():
            raise ValueError("diffused claim and justification must be non-empty")
        return self


class DiffusedProof(StrictModel):
    nodes: list[DiffusedProofNode] = Field(min_length=1)
    terminal_node_id: str
    terminal_answer: str

    def validate_against(
        self,
        dag: ProofDAG,
        method_plan: MethodPlan,
    ) -> "DiffusedProof":
        method_plan.validate_against(dag)
        if [node.node_id for node in self.nodes] != dag.node_ids():
            raise ValueError("diffused proof must contain exactly one row per DAG node")
        dag_by_id = {node.node_id: node for node in dag.nodes}
        methods = {node.node_id: node.method_label for node in method_plan.nodes}
        for node in self.nodes:
            if node.dependencies != dag_by_id[node.node_id].dependencies:
                raise ValueError(
                    f"{node.node_id} dependencies changed during DAG diffusion"
                )
            if node.method_label != methods[node.node_id]:
                raise ValueError(
                    f"{node.node_id} method label changed during DAG diffusion"
                )
        if self.terminal_node_id != dag.terminal_node_id:
            raise ValueError("diffused terminal node must match the source DAG")
        if not self.terminal_answer.strip():
            raise ValueError("diffused proof must expose a terminal answer")
        return self


class RenderedVariant(StrictModel):
    question: str
    solution: str
    node_order: list[str] = Field(min_length=1)
    terminal_answer: str

    @model_validator(mode="after")
    def validate_content(self) -> "RenderedVariant":
        if not self.question.strip() or not self.solution.strip():
            raise ValueError("rendered question and solution must be non-empty")
        if not self.terminal_answer.strip():
            raise ValueError("rendered terminal answer must be non-empty")
        return self

    def validate_against(
        self,
        dag: ProofDAG,
        diffused_proof: DiffusedProof,
    ) -> "RenderedVariant":
        if self.node_order != dag.node_ids():
            raise ValueError("rendered solution must cite every proof node in order")
        missing_markers = [
            node_id
            for node_id in self.node_order
            if f"[{node_id}]" not in self.solution
        ]
        if missing_markers:
            raise ValueError(
                f"rendered solution is missing node markers {missing_markers}"
            )
        if self.terminal_answer.strip() != diffused_proof.terminal_answer.strip():
            raise ValueError("rendered terminal answer differs from diffused proof")
        return self


class CandidateBundle(StrictModel):
    replacement_plan: ReplacementPlan
    diffused_proof: DiffusedProof
    variant: RenderedVariant


class JudgeVerdict(StrictModel):
    verdict: Literal["accept", "reject"]
    step_by_step_check: str
    blocking_issues: str = ""
    patch_suggestion: str = ""

    @field_validator(
        "step_by_step_check",
        "blocking_issues",
        "patch_suggestion",
        mode="before",
    )
    @classmethod
    def normalize_text_fields(cls, value: Any) -> str:
        if value is None:
            return ""
        if isinstance(value, str):
            return value
        return json.dumps(value, ensure_ascii=False, sort_keys=True)

    @field_validator("blocking_issues", "patch_suggestion", mode="after")
    @classmethod
    def normalize_absence_sentinels(cls, value: str) -> str:
        normalized = value.strip().lower().rstrip(".")
        if normalized in {
            "",
            "none",
            "none detected",
            "n/a",
            "no issues",
            "no blocking issues",
        }:
            return ""
        return value

    @model_validator(mode="after")
    def validate_verdict(self) -> "JudgeVerdict":
        if self.verdict == "accept" and self.blocking_issues.strip():
            raise ValueError("accept verdict cannot contain blocking issues")
        if self.verdict == "reject" and not self.blocking_issues.strip():
            raise ValueError("reject verdict must identify a blocking issue")
        return self

    def validate_coverage(
        self,
        dag: ProofDAG,
    ) -> "JudgeVerdict":
        missing_nodes = [
            node_id
            for node_id in dag.node_ids()
            if re.search(
                rf"(?<![A-Za-z0-9_]){re.escape(node_id)}(?![A-Za-z0-9_])",
                self.step_by_step_check,
            )
            is None
        ]
        if missing_nodes:
            raise ValueError(f"judge coverage incomplete: nodes={missing_nodes}")
        return self


class VerificationIteration(StrictModel):
    iteration: int
    bundle_sha256: str
    candidate_sha256: str
    verdicts: list[JudgeVerdict]
    unanimous: bool
    pass_streak_after: int
    repaired_from_previous: bool = False


class KernelRunResult(StrictModel):
    item_id: str
    status: Literal["accepted", "rejected"]
    proof_dag: ProofDAG
    method_plan: MethodPlan
    accepted_replacement_plan: ReplacementPlan | None = None
    accepted_diffused_proof: DiffusedProof | None = None
    accepted_candidate: RenderedVariant | None = None
    iterations: list[VerificationIteration]
    rejection_reason: str = ""
    accepted_candidate_sha256: str | None = None
    accepted_bundle_sha256: str | None = None
    schema_version: str = f"{SCHEMA_VERSION}-literal-five-stage"

    @model_validator(mode="after")
    def validate_terminal_state(self) -> "KernelRunResult":
        if self.status == "accepted":
            required = [
                self.accepted_replacement_plan,
                self.accepted_diffused_proof,
                self.accepted_candidate,
                self.accepted_candidate_sha256,
                self.accepted_bundle_sha256,
            ]
            if any(value is None for value in required):
                raise ValueError("accepted run must include complete stage provenance")
        elif not self.rejection_reason:
            raise ValueError("rejected run must state a reason")
        return self