# contracts/verify_possibility.py
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Literal, Protocol, Annotated, Union, Any

from pydantic import BaseModel, ConfigDict, Field

from contracts.base import _StrictModel
from contracts.coders.GatherOutput import GatherOutput

@dataclass(frozen=True)
class VerifyPossibilityInput:
    gather: GatherOutput


def sanitize_for_openai(schema: dict[str, Any]) -> dict[str, Any]:
    """
    Recursively rewrite a Pydantic-generated JSON schema into the subset
    OpenAI strict mode accepts.

    - oneOf          -> anyOf   (OpenAI rejects oneOf)
    - discriminator  -> removed (OpenAI rejects it; the const fields
                                 already make branches disjoint)
    - adds additionalProperties: false to every object missing it
    - strips None defaults
    """
    return _walk(schema)


def _walk(node: Any) -> Any:
    if isinstance(node, list):
        return [_walk(x) for x in node]
    if not isinstance(node, dict):
        return node

    # Recurse into every nested schema container first.
    for key in ("properties", "$defs", "definitions"):
        if key in node and isinstance(node[key], dict):
            node[key] = {k: _walk(v) for k, v in node[key].items()}

    for key in ("items", "additionalProperties"):
        if key in node and isinstance(node[key], (dict, list)):
            node[key] = _walk(node[key])

    for key in ("anyOf", "allOf"):
        if key in node and isinstance(node[key], list):
            node[key] = [_walk(v) for v in node[key]]

    # ── the two rewrites that matter ──────────────────────────
    if "oneOf" in node and isinstance(node["oneOf"], list):
        existing = node.get("anyOf", [])
        if not isinstance(existing, list):
            existing = []
        node["anyOf"] = existing + [_walk(v) for v in node["oneOf"]]
        node.pop("oneOf")

    node.pop("discriminator", None)   # OpenAI doesn't accept it

    # ── strict-mode hygiene ───────────────────────────────────
    if node.get("type") == "object":
        node.setdefault("additionalProperties", False)
        props = node.get("properties")
        if isinstance(props, dict):
            # Strict mode: every property must be in required.
            node["required"] = list(props.keys())

    if node.get("default", object()) is None:
        node.pop("default", None)

    return node

class Option(_StrictModel):
    label: str = Field(min_length=1, max_length=80)
    description: str = Field(min_length=1)
    tradeoff: str | None = Field(
        description="Optional trade-off; null if none.",
    )

class _BaseVerdict(_StrictModel):
    reasoning: str = Field(
        min_length=1,
        description="2–5 sentences explaining the verdict.",
    )
    confidence: float = Field(
        ge=0.0, le=1.0,
        description="Your honest probability that the verdict is correct.",
    )


class FeasibleVerdict(_BaseVerdict):
    verdict: Literal["feasible"]


class AmbiguousVerdict(_BaseVerdict):
    verdict: Literal["ambiguous"]
    question: str = Field(
        min_length=1,
        description="Exactly one clarifying question that unblocks planning.",
    )


class InfeasibleVerdict(_BaseVerdict):
    verdict: Literal["infeasible"]
    options: list[Option] = Field(
        min_length=1, max_length=3,
        description="1–3 concrete alternatives.",
    )


Verdict = Annotated[
    Union[FeasibleVerdict, AmbiguousVerdict, InfeasibleVerdict],
    Field(discriminator="verdict"),
]


class VerifyPossibilityOutput(_StrictModel):
    """Top-level object required by OpenAI strict mode."""
    result: Verdict = Field(
        description="The feasibility verdict.",
    )