"""Pydantic output schemas for structured LLM responses."""

from __future__ import annotations
import enum
from typing import Literal  # noqa: F401 — used at runtime by Pydantic
from pydantic import BaseModel, create_model


class IncludeRating(enum.IntEnum):
    DEFINITELY_EXCLUDE = 1
    PROBABLY_EXCLUDE = 2
    UNSURE = 3
    PROBABLY_INCLUDE = 4
    DEFINITELY_INCLUDE = 5


class IncludeResponse(BaseModel):
    include_rating: IncludeRating


class ReasoningResponse(BaseModel):
    rating_reasoning: str


class ScreeningEvidenceSpan(BaseModel):
    text: str
    section: Literal["title", "abstract"]
    criterion: str  # Which criterion category (e.g. "Population")
    supports: Literal["include", "exclude"]


class ReasoningWithEvidenceResponse(BaseModel):
    rating_reasoning: str
    evidence: list[ScreeningEvidenceSpan]


# Dynamic schema builders for clustering (need criteria-dependent enums)
def build_cluster_schema(criteria_keys: list[str]):
    """Build dynamic Pydantic model for clustering output."""
    criteria_dict = {x: x for x in list(criteria_keys) + ["Other"]}
    CriteriaEnum = enum.Enum("CriteriaEnum", criteria_dict)
    ClusterSchema = create_model(
        "ClusterSchema",
        cluster_name=(str, ...),
        cluster_description=(str, ...),
        related_criteria=(CriteriaEnum, ...),
    )
    return create_model("ClusterResponse", cluster_list=(list[ClusterSchema], ...))


def build_cluster_selection_schema():
    """Build Pydantic model for cluster selection output."""
    return create_model("ClusterSelectionResponse", cluster_list=(list[int], ...))