"""
Data models for dataset augmentation CLI tool.
"""

from typing import Any, Dict, List, Optional

from pydantic import BaseModel, Field


class DatasetSample(BaseModel):
    """Represents a dataset sample from Braintrust"""
    id: str
    input: Optional[Any] = None
    expected: Optional[Any] = None
    metadata: Dict[str, Any] = Field(default_factory=dict)


class GapAnalysisSuggestion(BaseModel):
    """Represents an LLM suggestion for dataset gaps"""
    title: str = Field(description="Short descriptive title for the gap")
    description: str = Field(description="Detailed description of what this gap covers")
    rationale: str = Field(description="Why this gap is important to fill")


class GapAnalysisResult(BaseModel):
    """Result from LLM gap analysis"""
    suggestions: List[GapAnalysisSuggestion] = Field(
        description="List of exactly 5 suggestions for dataset gaps"
    )
    overall_assessment: str = Field(
        description="Overall assessment of current dataset coverage"
    )


class InferredSchema(BaseModel):
    """Schema analysis result from examining dataset samples"""
    input_schema: Dict[str, Any] = Field(description="JSON Schema object for the input field - only what is actually observed")
    expected_schema: Dict[str, Any] = Field(description="JSON Schema object for the expected field - only what is actually observed")
    metadata_schema: Dict[str, Any] = Field(description="JSON Schema object for the metadata field - only what is actually observed")
    observed_patterns: List[str] = Field(description="Concrete patterns observed in the data (no extrapolation)")
    field_relationships: List[str] = Field(description="Observed relationships between input/expected fields (no guessing)")


class GeneratedSample(BaseModel):
    """A newly generated dataset sample"""
    input: Any = Field(description="Input data matching the dataset structure")
    expected: Any = Field(description="Expected output matching the dataset structure")
    metadata: Dict[str, Any] = Field(description="Metadata including test_name")


class CaseAbstract(BaseModel):
    """High-level abstract description of a test case to generate"""
    title: str = Field(description="Short descriptive title")
    description: str = Field(description="Detailed description of the test case")
    expected_input_characteristics: str = Field(description="What the input should look like")
    expected_output_characteristics: str = Field(description="What the expected output should look like")


class CaseAbstractList(BaseModel):
    """List of case abstracts generated by LLM"""
    abstracts: List[CaseAbstract] = Field(description="List of case abstracts to generate")
    generation_notes: str = Field(description="Notes about the generation process and any assumptions made")
