"""Type definitions for classification results."""

from typing import Optional
from pydantic import BaseModel, Field


class PredictionDetail(BaseModel):
    """Individual domain prediction with confidence."""

    domain: str = Field(..., description="Domain name (e.g., 'Web_Search', 'TI', 'Generic')")
    confidence: float = Field(..., ge=0.0, le=1.0, description="Confidence score between 0 and 1")


class ClassificationResult(BaseModel):
    """Result of domain classification.

    Attributes:
        primary: The primary domain with highest confidence
        confidence: Confidence score for the primary domain (0-1)
        includes_generic: Whether Generic should be included due to low confidence
        all_predictions: Optional list of all predictions above threshold (multi-label mode)
    """

    primary: str = Field(..., description="Primary classified domain")
    confidence: float = Field(..., ge=0.0, le=1.0, description="Confidence score for primary domain")
    includes_generic: bool = Field(
        default=False, description="True if confidence is below threshold (Generic should be considered)"
    )
    all_predictions: Optional[list[PredictionDetail]] = Field(
        default=None, description="All predictions above threshold (multi-label mode only)"
    )

    class Config:
        json_schema_extra = {
            "example": {
                "primary": "Web_Search",
                "confidence": 0.98,
                "includes_generic": False,
                "all_predictions": [
                    {"domain": "Web_Search", "confidence": 0.98},
                    {"domain": "Generic", "confidence": 0.75},
                ],
            }
        }
