models.py

"""Data models for Language Flashcard Generator."""

from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import Dict, List, Optional

from pydantic import BaseModel, Field

current_version = 1


class ExampleSentence(BaseModel):
    """An example sentence in the target language with English translation."""

    target_language: str = Field(..., description="Example sentence in target language")
    english: str = Field(..., description="English translation")
    audio_file: Optional[str] = Field(None, description="Path to audio file")


class AdditionalMeaning(BaseModel):
    """An additional meaning with its part of speech."""

    meaning: str = Field(..., description="The additional meaning")
    part_of_speech: str = Field(
        ..., description="Part of speech (e.g., 'n.', 'adj.', 'v.')"
    )


class AudioFiles(BaseModel):
    """Audio files associated with a card."""

    word: Optional[str] = Field(None, description="Audio for the main word")
    examples: List[str] = Field(
        default_factory=list, description="Audio files for example sentences"
    )


class EnhancedContent(BaseModel):
    """AI-generated enhanced content for a card."""

    word: str = Field(..., description="Word in target language")
    romanization: str = Field(..., description="Romanization of the word")
    pronunciation_ipa: str = Field(..., description="IPA phonetic transcription")
    main_part_of_speech: str = Field(..., description="Part of speech (e.g., 'n.')")
    examples: List[ExampleSentence] = Field(
        default_factory=list, description="Example sentences"
    )
    etymology: str = Field(default="", description="Etymology and word origin")
    additional_meanings: List[AdditionalMeaning] = Field(
        default_factory=list, description="Additional meanings"
    )
    pos_tags: List[str] = Field(
        default_factory=list, description="Part-of-speech tags"
    )
    audio_files: AudioFiles = Field(
        default_factory=AudioFiles, description="Associated audio files"
    )


class CardData(BaseModel):
    """Original card data from CSV."""

    id: str = Field(..., description="Unique card identifier")
    word: str = Field(..., description="Word in target language")
    english: str = Field(..., description="English translation")


class EnhancedCard(BaseModel):
    """Complete enhanced card with original and AI-generated content."""

    original: CardData
    enhanced: EnhancedContent
    version: int = Field(default=current_version, description="Card version")

    def to_mochi_content(self) -> str:
        """Convert to Mochi markdown format with audio links."""
        examples_list = []
        for ex in self.enhanced.examples:
            example_text = f"{ex.target_language}\n\n*{ex.english}*"
            if ex.audio_file and Path(ex.audio_file).exists():
                example_text += f"\n\n![audio]({Path(ex.audio_file).name})"
            examples_list.append(example_text)

        examples_text = "\n\n".join(examples_list)

        additional_meanings_text = ""
        if self.enhanced.additional_meanings:
            meaning_parts = [
                f"({m.part_of_speech} {m.meaning})"
                for m in self.enhanced.additional_meanings
            ]
            additional_meanings_text = "\n" + "\n".join(meaning_parts)

        word_audio = ""
        if (
            self.enhanced.audio_files.word
            and Path(self.enhanced.audio_files.word).exists()
        ):
            word_audio = (
                f"\n\n![audio]({Path(self.enhanced.audio_files.word).name})"
            )

        additional_meanings_simple = ""
        if self.enhanced.additional_meanings:
            additional_meanings_simple = ", " + ", ".join(
                [m.meaning for m in self.enhanced.additional_meanings]
            )

        content = f"""# {self.enhanced.word}{word_audio}

---

### {self.enhanced.word}{word_audio}

({self.enhanced.main_part_of_speech} {self.original.english}){additional_meanings_text}

{self.enhanced.romanization}
{self.enhanced.pronunciation_ipa}

**Examples:**

{examples_text}

**Etymology:**
{self.enhanced.etymology}

---

# {self.original.english}{additional_meanings_simple}"""

        return content


class ProcessingProgress(BaseModel):
    """Track processing progress and failures."""

    processed_ids: List[str] = Field(default_factory=list)
    failed_ids: List[str] = Field(default_factory=list)
    last_batch: int = Field(default=0)
    total_cards: int = Field(default=0)

    @property
    def completion_rate(self) -> float:
        if self.total_cards == 0:
            return 0.0
        return len(self.processed_ids) / self.total_cards


# --- Batch API models ---


class BatchProcessingStatus(str, Enum):
    """Status of batch processing."""

    NOT_STARTED = "not_started"
    IN_PROGRESS = "in_progress"
    CANCELING = "canceling"
    ENDED = "ended"
    FAILED = "failed"


class BatchRequest(BaseModel):
    """Individual request within a batch."""

    custom_id: str = Field(..., description="Unique identifier for this request")
    params: Dict = Field(..., description="Message creation parameters")


class BatchInfo(BaseModel):
    """Information about a submitted batch."""

    batch_id: str = Field(..., description="Batch identifier from Claude")
    status: BatchProcessingStatus = Field(..., description="Current batch status")
    created_at: str = Field(..., description="When batch was created")
    expires_at: Optional[str] = Field(None, description="When batch expires")
    ended_at: Optional[str] = Field(None, description="When batch completed")
    results_url: Optional[str] = Field(None, description="URL to download results")
    total_requests: int = Field(..., description="Total number of requests in batch")
    succeeded: int = Field(default=0, description="Number of successful requests")
    failed: int = Field(default=0, description="Number of failed requests")
    processing: int = Field(default=0, description="Number of processing requests")
    card_id_mapping: Optional[Dict[str, str]] = Field(
        default=None,
        description="Mapping from sanitized custom_id to original card_id",
    )


class BatchProgress(BaseModel):
    """Track batch processing progress across multiple batches."""

    batches: List[BatchInfo] = Field(default_factory=list)
    cards_by_batch: Dict[str, List[str]] = Field(default_factory=dict)
    completed_cards: List[str] = Field(default_factory=list)
    failed_cards: List[str] = Field(default_factory=list)
    total_cards: int = Field(default=0)

    @property
    def completion_rate(self) -> float:
        if self.total_cards == 0:
            return 0.0
        return len(self.completed_cards) / self.total_cards

    @property
    def active_batches(self) -> List[BatchInfo]:
        return [
            b
            for b in self.batches
            if b.status
            in [BatchProcessingStatus.IN_PROGRESS, BatchProcessingStatus.CANCELING]
        ]


class BatchResult(BaseModel):
    """Result from a batch request."""

    custom_id: str = Field(..., description="The custom ID from the request")
    result: Optional[Dict] = Field(None, description="The successful result")
    error: Optional[Dict] = Field(None, description="Error information if failed")


# --- Config models ---


class AzureTTSConfig(BaseModel):
    """Azure TTS voice configuration."""

    male_voice: str = Field(
        default="zu-ZA-ThembaNeural", description="Male voice name"
    )
    female_voice: str = Field(
        default="zu-ZA-ThandoNeural", description="Female voice name"
    )


class TTSConfig(BaseModel):
    """TTS configuration."""

    enabled: bool = Field(default=True, description="Whether to generate TTS")
    azure: AzureTTSConfig = Field(
        default_factory=AzureTTSConfig, description="Azure TTS settings"
    )


class EnhancementConfig(BaseModel):
    """Configuration for a deck build."""

    key: str = Field(..., description="Unique key for build isolation")
    deck_name: str = Field(
        default="Enhanced Vocabulary", description="Name of the Mochi deck"
    )
    frequency_dir: str = Field(..., description="Path to frequency list CSV directory")

    ai_provider: str = Field(default="anthropic", description="AI provider")
    model: str = Field(
        default="claude-3-5-haiku-20241022", description="Model name"
    )
    batch_size: int = Field(default=10, description="Cards per processing batch")
    rate_limit_delay: float = Field(
        default=1.0, description="Delay between API calls in seconds"
    )
    max_retries: int = Field(default=3, description="Max retry attempts")

    use_batch_api: bool = Field(default=False, description="Use Claude batch API")
    batch_check_interval: int = Field(
        default=10, description="Seconds between batch status checks"
    )

    audio_concurrency: int = Field(
        default=16, description="Max parallel Azure TTS workers"
    )

    tts: TTSConfig = Field(default_factory=TTSConfig, description="TTS configuration")

    enhancement_fields: Dict[str, bool] = Field(
        default_factory=lambda: {
            "romanization": True,
            "pronunciation_ipa": True,
            "example_sentences": True,
            "etymology": False,
            "additional_meanings": True,
            "part_of_speech": True,
        }
    )

    @property
    def build_dir(self) -> Path:
        """Root directory for all build artifacts for this config."""
        return Path("builds") / self.key

    @property
    def cached_cards_dir(self) -> Path:
        return self.build_dir / "cached_cards"

    @property
    def audio_dir(self) -> Path:
        return self.build_dir / "audio_files"

    @property
    def output_dir(self) -> Path:
        return self.build_dir / "output"

    @property
    def input_dir(self) -> Path:
        return self.build_dir / "input"

    @property
    def log_dir(self) -> Path:
        return self.build_dir / "logs"

    @property
    def progress_file(self) -> Path:
        return self.build_dir / "progress.json"

    @property
    def batch_progress_file(self) -> Path:
        return self.build_dir / "batch_progress.json"