"""Pydantic models for dataset analysis."""

from typing import Literal

from pydantic import BaseModel, Field


class QualityIssue(BaseModel):
    """A data quality issue found in the dataset."""

    issue_type: str  # "empty_row", "duplicate", "encoding_error"
    severity: Literal["warning", "error"]
    message: str
    row_indices: list[int] = Field(default_factory=list)
    count: int = 1


class PIIFinding(BaseModel):
    """A PII finding in the dataset."""

    pii_type: str  # "email", "phone_tr", "tc_kimlik", "credit_card"
    count: int
    sample_indices: list[int] = Field(default_factory=list)


class DatasetInsights(BaseModel):
    """LLM-generated dataset analysis insights."""

    quality_summary: str = ""
    recommended_preprocessing: list[str] = Field(default_factory=list)
    recommended_tokenizer: str | None = None
    training_suitability: Literal["excellent", "good", "fair", "poor"] = "fair"
    warnings: list[str] = Field(default_factory=list)
    suggestions: list[str] = Field(default_factory=list)


class DatasetProfile(BaseModel):
    """Complete analysis profile for a dataset."""

    path: str
    format: Literal["alpaca", "sharegpt", "custom", "unknown"] = "unknown"
    num_rows: int = 0
    num_columns: int = 0
    columns: list[str] = Field(default_factory=list)
    text_column: str | None = None
    avg_tokens: float = 0.0
    min_tokens: int = 0
    max_tokens: int = 0
    p50_tokens: int = 0
    p90_tokens: int = 0
    p99_tokens: int = 0
    tokenizer_name: str | None = None
    token_method: Literal["tokenizer", "char_estimate"] = "char_estimate"
    language: str | None = None
    language_distribution: dict[str, float] = Field(default_factory=dict)
    duplicate_ratio: float = 0.0
    empty_ratio: float = 0.0
    quality_issues: list[QualityIssue] = Field(default_factory=list)
    pii_findings: list[PIIFinding] = Field(default_factory=list)
    estimated_size_mb: float = 0.0
