from dataclasses import dataclass
from typing import Callable, Literal

from typesafe_sdk import Choice, TypeSafeClient

client = TypeSafeClient()


@dataclass
class Verdict:
    label: str
    confidence: float
    source: Literal["jev", "llm", "human"]


def cascade_judge(
    rubric: str,
    candidate: str,
    llm_judge: Callable[[str, str], Verdict],  # your stronger model, same label space
    jev_gate: float = 0.85,
    llm_gate: float = 0.80,
) -> Verdict:
    state = f"RUBRIC:\n{rubric}\n\nCANDIDATE RESPONSE:\n{candidate}"

    first = client.system_one(
        state=state,
        questions={
            "verdict": Choice(
                instructions="Does the candidate satisfy the rubric",
                criteria={
                    "pass": "Meets every requirement in the rubric",
                    "fail": "Misses or violates at least one requirement",
                },
            )
        },
    ).answers["verdict"]

    if first.confidence >= jev_gate:
        return Verdict(first.choice, first.confidence, "jev")

    second = llm_judge(rubric, candidate)
    if second.confidence >= llm_gate:
        return second

    return Verdict("needs_review", second.confidence, "human")