"""Extraction output contract (design §7.1). The model must return exactly this shape."""
from __future__ import annotations
from typing import Literal
from pydantic import BaseModel, Field, field_validator, model_validator

ElementType = Literal["line", "heading", "equation", "chem", "figure", "table", "label"]
Flag = Literal["struck", "inserted", "uncertain", "illegible", "rough"]
Lang = Literal["bn", "en", "ar"]
PageKind = Literal["cover", "answer", "supplementary", "rough", "blank", "unreadable"]


class Element(BaseModel):
    type: ElementType
    bbox: list[float] | None = Field(None, description="[x0,y0,x1,y1] normalized 0..1, approximate")
    text: str | None = Field(None, description="Verbatim text. For 'heading': the raw label as written.")
    latex: str | None = None
    description: str | None = Field(None, description="For figure/table: what is actually drawn.")
    table: list[list[str]] | None = None
    lang: list[Lang] = Field(default_factory=list)
    confidence: float = Field(ge=0, le=1)
    flags: list[Flag] = Field(default_factory=list)
    candidates: list[str] | None = Field(None, description="Up to 3 alternative readings when uncertain")

    @field_validator("bbox")
    @classmethod
    def _bbox(cls, v):
        if v is None:
            return v
        if len(v) != 4:
            raise ValueError("bbox must have 4 numbers")
        # Accept 0..1000 integer grids too (some models do better with ints) and normalize.
        if max(v) > 1.0001:
            v = [x / 1000.0 for x in v]
        x0, y0, x1, y1 = (min(max(x, 0.0), 1.0) for x in v)
        if not (x0 < x1 and y0 < y1):
            raise ValueError("bbox must satisfy x0<x1 and y0<y1")
        return [x0, y0, x1, y1]

    @model_validator(mode="after")
    def _content(self):
        if "illegible" in self.flags:
            return self
        if self.type in ("line", "heading", "label") and not (self.text or "").strip():
            raise ValueError(f"{self.type} element needs text unless flagged illegible")
        if self.type in ("figure",) and not self.description:
            raise ValueError("figure element needs a description")
        if self.candidates and len(self.candidates) > 3:
            self.candidates = self.candidates[:3]
        return self


class CoverFields(BaseModel):
    roll: str | None = None
    registration: str | None = None
    subject_code: str | None = None


class ExtractionPage(BaseModel):
    kind: PageKind
    cover: CoverFields | None = None
    legibility: int = Field(ge=0, le=10)
    elements: list[Element] = Field(default_factory=list)
    notes: str | None = None

    @model_validator(mode="after")
    def _kind_consistency(self):
        if self.kind in ("answer", "supplementary") and not self.elements:
            raise ValueError("answer/supplementary page returned with no elements")
        return self