"""Request and response models for the specialist inference service."""

from __future__ import annotations

from pydantic import BaseModel, Field

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

MODEL_VERSION = "gliner2-large-v1.0"
SCHEMA_VERSION = "specialist-v1"


# ---------------------------------------------------------------------------
# Request models
# ---------------------------------------------------------------------------


class ExtractionRequest(BaseModel):
    """Batch extraction request accepted by entity, relation, and structured endpoints."""

    texts: list[str] = Field(..., min_length=1, description="Texts to process")
    schema_labels: list[str] = Field(
        ..., min_length=1, description="Entity/relation/event labels to extract"
    )
    batch_id: str | None = Field(
        default=None, description="Optional caller-provided batch identifier"
    )


class ClassificationRequest(BaseModel):
    """Batch classification request accepted by the classify endpoint."""

    texts: list[str] = Field(..., min_length=1, description="Texts to classify")
    schema_labels: list[str] = Field(
        ..., min_length=1, description="Classification labels"
    )
    batch_id: str | None = Field(
        default=None, description="Optional caller-provided batch identifier"
    )


# ---------------------------------------------------------------------------
# Result models
# ---------------------------------------------------------------------------


class EntityResult(BaseModel):
    """A single extracted entity span."""

    text: str
    entity_type: str
    start_char: int
    end_char: int
    score: float = Field(..., ge=0.0, le=1.0)
    model_version: str = MODEL_VERSION
    schema_version: str = SCHEMA_VERSION


class RelationResult(BaseModel):
    """A single extracted relation."""

    subject: str
    subject_type: str
    subject_start: int
    subject_end: int
    relation: str
    object: str
    object_type: str
    object_start: int
    object_end: int
    score: float = Field(..., ge=0.0, le=1.0)
    model_version: str = MODEL_VERSION
    schema_version: str = SCHEMA_VERSION


class ClassificationResult(BaseModel):
    """A single classification result."""

    text: str
    label: str
    score: float = Field(..., ge=0.0, le=1.0)
    model_version: str = MODEL_VERSION
    schema_version: str = SCHEMA_VERSION


class StructuredResult(BaseModel):
    """A single structured extraction result with key-value facts."""

    text: str
    field: str
    value: str
    start_char: int
    end_char: int
    score: float = Field(..., ge=0.0, le=1.0)
    model_version: str = MODEL_VERSION
    schema_version: str = SCHEMA_VERSION


# ---------------------------------------------------------------------------
# Batch response
# ---------------------------------------------------------------------------


class BatchResponse(BaseModel):
    """Unified batch response wrapping results from any endpoint."""

    results: list[list[EntityResult]] | list[list[RelationResult]] | list[list[ClassificationResult]] | list[list[StructuredResult]]
    model_version: str = MODEL_VERSION
    schema_version: str = SCHEMA_VERSION
    processing_time_ms: float
    batch_id: str | None = None