from __future__ import annotations

from pydantic import BaseModel, ConfigDict, Field, model_validator


class DiarizationSpan(BaseModel):
    model_config = ConfigDict(extra="forbid")

    speaker: str = Field(min_length=1, max_length=128)
    start: float = Field(ge=0)
    end: float = Field(ge=0)

    @model_validator(mode="after")
    def validate_range(self) -> "DiarizationSpan":
        if self.end <= self.start:
            raise ValueError("diarization span end must be greater than start")
        return self


class TranscriptWord(BaseModel):
    model_config = ConfigDict(extra="forbid")
    word: str
    start: float = Field(ge=0)
    end: float = Field(ge=0)
    probability: float | None = Field(default=None, ge=0, le=1)
    speaker: str | None = Field(default=None, min_length=1, max_length=128)

    @model_validator(mode="after")
    def validate_range(self) -> "TranscriptWord":
        if self.end < self.start:
            raise ValueError("word end must be greater than or equal to start")
        return self


class TranscriptSegment(BaseModel):
    model_config = ConfigDict(extra="forbid")
    id: int
    start: float = Field(ge=0)
    end: float = Field(ge=0)
    text: str
    speaker: str | None = Field(default=None, min_length=1, max_length=128)
    words: list[TranscriptWord] = Field(default_factory=list)

    @model_validator(mode="after")
    def validate_range(self) -> "TranscriptSegment":
        if self.end < self.start:
            raise ValueError("segment end must be greater than or equal to start")
        return self


class Transcript(BaseModel):
    model_config = ConfigDict(extra="forbid")
    schema_version: str = "1.1"
    language: str
    language_probability: float = Field(ge=0, le=1)
    duration: float = Field(ge=0)
    task: str
    model: str
    text: str
    speakers: list[str] = Field(default_factory=list)
    segments: list[TranscriptSegment] = Field(default_factory=list)
    word_timestamps: bool = True