"""Pydantic models for company-specific sentiment analysis."""

from __future__ import annotations

from pydantic import BaseModel, Field, field_validator


class EvidenceGroup(BaseModel):
    """A group of evidence spans associated with a specific company.

    A single evidence span can appear in multiple groups when the
    underlying text mentions multiple companies.
    """

    company_id: str = Field(description="Resolved company identifier")
    evidence_ids: list[str] = Field(description="IDs of evidence spans in this group")
    texts: list[str] = Field(description="Text snippets from evidence spans")

    @field_validator("evidence_ids")
    @classmethod
    def evidence_ids_non_empty(cls, v: list[str]) -> list[str]:
        if not v:
            raise ValueError("evidence_ids must not be empty")
        return v

    @field_validator("texts")
    @classmethod
    def texts_non_empty(cls, v: list[str]) -> list[str]:
        if not v:
            raise ValueError("texts must not be empty")
        return v


class TextSentiment(BaseModel):
    """Per-text sentiment probability distribution with evidence linkage.

    Stores the raw FinBERT output for a single evidence span text,
    enabling full provenance from probability to source evidence.
    """

    evidence_id: str = Field(description="Evidence span ID this score belongs to")
    positive_prob: float = Field(ge=0.0, le=1.0, description="Probability of positive sentiment")
    negative_prob: float = Field(ge=0.0, le=1.0, description="Probability of negative sentiment")
    neutral_prob: float = Field(ge=0.0, le=1.0, description="Probability of neutral sentiment")

    @property
    def dominant_label(self) -> str:
        """Return the label with highest probability."""
        if self.positive_prob >= self.negative_prob and self.positive_prob >= self.neutral_prob:
            return "positive"
        elif self.negative_prob >= self.positive_prob and self.negative_prob >= self.neutral_prob:
            return "negative"
        return "neutral"


class CompanySentimentResult(BaseModel):
    """Sentiment classification result for a single company.

    Contains full probability distribution, supporting evidence IDs,
    per-text scores, and model/calibration versioning for lineage tracking.
    """

    company_id: str = Field(description="Resolved company identifier")
    label: str = Field(description="Derived label: positive, negative, neutral, or mixed")
    positive_prob: float = Field(ge=0.0, le=1.0, description="Probability of positive sentiment")
    negative_prob: float = Field(ge=0.0, le=1.0, description="Probability of negative sentiment")
    neutral_prob: float = Field(ge=0.0, le=1.0, description="Probability of neutral sentiment")
    evidence_ids: list[str] = Field(description="All evidence span IDs contributing to this result")
    is_mixed: bool = Field(default=False, description="Whether mixed sentiment was detected from disagreement")
    per_text_scores: list[TextSentiment] = Field(
        default_factory=list,
        description="Full probability distributions per evidence text",
    )
    model_version: str = Field(description="Sentiment model name and version")
    calibration_version: str = Field(default="uncalibrated", description="Calibration artifact version")

    @field_validator("label")
    @classmethod
    def label_valid(cls, v: str) -> str:
        valid_labels = {"positive", "negative", "neutral", "mixed"}
        if v not in valid_labels:
            raise ValueError(f"label must be one of {valid_labels}, got '{v}'")
        return v


class SentimentBatchResult(BaseModel):
    """Result of sentiment classification for a batch of companies."""

    results: list[CompanySentimentResult] = Field(description="Per-company sentiment results")
    model_version: str = Field(description="Sentiment model version used for batch")
    processing_time_ms: int = Field(ge=0, description="Total processing time in milliseconds")