enhancer.py

"""Flashcard enhancement engine using Claude AI and Azure TTS."""

import asyncio
import csv
import hashlib
import json
import logging
import os
import re
import shutil
import time
import zipfile
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, List, Optional

import anthropic
import azure.cognitiveservices.speech as speechsdk
from azure.identity import DefaultAzureCredential
from dotenv import load_dotenv
from rich.console import Console
from rich.progress import BarColumn, Progress, SpinnerColumn, TextColumn

from models import (
    AdditionalMeaning,
    AudioFiles,
    BatchInfo,
    BatchProcessingStatus,
    BatchProgress,
    BatchResult,
    CardData,
    EnhancedCard,
    EnhancedContent,
    EnhancementConfig,
    ExampleSentence,
    ProcessingProgress,
    current_version,
)

load_dotenv()

console = Console()


class FlashcardEnhancer:
    """Enhanced flashcard generator using Claude AI and Azure TTS."""

    def __init__(self, config: EnhancementConfig):
        self.config = config
        self._setup_logging()
        self._setup_directories()

        # Initialize Claude client
        self.client = anthropic.Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))

        # Initialize Azure TTS if enabled
        self._speech_config = None
        if self.config.tts.enabled:
            self._init_azure_tts()

    def _setup_logging(self):
        self.config.log_dir.mkdir(parents=True, exist_ok=True)
        logging.basicConfig(
            level=logging.WARNING,
            format="%(asctime)s - %(levelname)s - %(message)s",
            handlers=[
                logging.FileHandler(self.config.log_dir / "enhancement.log"),
                logging.StreamHandler(),
            ],
        )
        self.logger = logging.getLogger(__name__)

    def _setup_directories(self):
        for d in [
            self.config.cached_cards_dir,
            self.config.audio_dir,
            self.config.output_dir,
            self.config.input_dir,
            self.config.log_dir,
        ]:
            d.mkdir(parents=True, exist_ok=True)

    def _init_azure_tts(self):
        endpoint = os.environ.get("SPEECH_ENDPOINT")
        resource_id = os.environ.get("SPEECH_RESOURCE_ID")
        if not endpoint or not resource_id:
            self.logger.warning(
                "SPEECH_ENDPOINT or SPEECH_RESOURCE_ID not set, TTS disabled"
            )
            self.config.tts.enabled = False
            return

        try:
            self._speech_config = speechsdk.SpeechConfig(endpoint=endpoint)
            aad_token = (
                DefaultAzureCredential()
                .get_token("https://cognitiveservices.azure.com/.default")
                .token
            )
            self._speech_config.authorization_token = (
                f"aad#{resource_id}#{aad_token}"
            )
            self._speech_config.set_speech_synthesis_output_format(
                speechsdk.SpeechSynthesisOutputFormat.Audio16Khz32KBitRateMonoMp3
            )
            console.print("[green]Azure TTS initialized[/green]")
        except Exception as e:
            self.logger.warning(f"Failed to init Azure TTS: {e}")
            self.config.tts.enabled = False

    def _get_target_language(self) -> str:
        """Infer target language from the Azure voice name."""
        voice = self.config.tts.azure.male_voice
        # Voice names like zu-ZA-ThembaNeural -> extract language
        lang_map = {
            "zu": "Zulu",
            "af": "Afrikaans",
            "sw": "Swahili",
            "ar": "Arabic",
            "fa": "Persian",
            "ja": "Japanese",
            "ko": "Korean",
            "zh": "Chinese",
            "es": "Spanish",
            "fr": "French",
            "de": "German",
            "pt": "Portuguese",
            "hi": "Hindi",
            "ru": "Russian",
        }
        lang_code = voice.split("-")[0] if "-" in voice else "en"
        return lang_map.get(lang_code, lang_code)

    # --- CSV Loading ---

    def load_cards(self) -> List[CardData]:
        """Load cards from all CSV files in the frequency directory."""
        freq_dir = Path(self.config.frequency_dir)
        if not freq_dir.exists():
            console.print(
                f"[red]Frequency directory not found: {freq_dir}[/red]"
            )
            return []

        cards = []
        csv_files = sorted(freq_dir.glob("*.csv"))

        for csv_file in csv_files:
            console.print(f"[blue]Loading {csv_file.name}...[/blue]")
            with open(csv_file, "r", encoding="utf-8") as f:
                reader = csv.reader(f)
                for row in reader:
                    if len(row) < 2:
                        continue
                    word = row[0].strip()
                    english = row[1].strip()
                    if not word or not english:
                        continue
                    card_id = hashlib.md5(word.encode()).hexdigest()[:8]
                    cards.append(CardData(id=card_id, word=word, english=english))

        console.print(f"[green]Loaded {len(cards)} cards from {len(csv_files)} files[/green]")
        return cards

    # --- Claude Enhancement ---

    def _build_enhancement_prompt(self, card: CardData) -> str:
        """Build the Claude prompt for enhancing a card."""
        target_lang = self._get_target_language()
        return f"""Provide detailed linguistic information for this {target_lang} word/phrase.

{target_lang} word/phrase: {card.word}
English meaning: {card.english}

Please provide the following in JSON format:

1. **word**: The {target_lang} word/phrase (as provided)
2. **romanization**: Romanization for an English speaker
3. **pronunciation_ipa**: IPA phonetic transcription in square brackets
4. **main_part_of_speech**: Part of speech (e.g., "n.", "adj.", "v.", "pron.", "prep.")
5. **examples**: Array of 3 example sentences, each with "target_language" ({target_lang} sentence) and "english" (translation) keys
6. **etymology**: Brief origin and development (2-3 sentences)
7. **additional_meanings**: Array of alternative meanings, each with "meaning" and "part_of_speech" keys
8. **pos_tags**: Array of part-of-speech tags (e.g., ["noun"], ["verb", "transitive"])

Format as valid JSON only:
{{
  "word": "{card.word}",
  "romanization": "romanization here",
  "pronunciation_ipa": "[IPA here]",
  "main_part_of_speech": "n.",
  "examples": [
    {{"target_language": "{target_lang} sentence", "english": "English translation"}},
    {{"target_language": "{target_lang} sentence", "english": "English translation"}},
    {{"target_language": "{target_lang} sentence", "english": "English translation"}}
  ],
  "etymology": "Etymology explanation...",
  "additional_meanings": [
    {{"meaning": "alternative meaning", "part_of_speech": "n."}}
  ],
  "pos_tags": ["part_of_speech"]
}}

Focus on accuracy. Provide authentic {target_lang} examples that naturally use this word."""

    def _parse_enhancement_response(
        self, response_text: str, card: CardData
    ) -> Optional[EnhancedContent]:
        """Parse a Claude response into EnhancedContent."""
        json_start = response_text.find("{")
        json_end = response_text.rfind("}") + 1

        if json_start < 0 or json_end <= json_start:
            self.logger.error(f"No JSON in response for '{card.word}'")
            return None

        json_text = response_text[json_start:json_end]
        try:
            data = json.loads(json_text)
        except json.JSONDecodeError:
            cleaned = json_text.replace("\n", "\\n").replace("\r", "\\r")
            try:
                data = json.loads(cleaned)
            except json.JSONDecodeError:
                self.logger.error(f"JSON parse failed for '{card.word}'")
                return None

        examples = [
            ExampleSentence(
                target_language=ex.get("target_language", ""),
                english=ex.get("english", ""),
            )
            for ex in data.get("examples", [])
        ]

        additional_meanings = []
        for md in data.get("additional_meanings", []):
            if isinstance(md, dict):
                additional_meanings.append(
                    AdditionalMeaning(
                        meaning=md.get("meaning", ""),
                        part_of_speech=md.get("part_of_speech", "n."),
                    )
                )

        return EnhancedContent(
            word=data.get("word", card.word),
            romanization=data.get("romanization", ""),
            pronunciation_ipa=data.get("pronunciation_ipa", ""),
            main_part_of_speech=data.get("main_part_of_speech", "n."),
            examples=examples,
            etymology=data.get("etymology", ""),
            additional_meanings=additional_meanings,
            pos_tags=data.get("pos_tags", []),
        )

    def enhance_card_with_claude(
        self, card: CardData
    ) -> Optional[EnhancedContent]:
        """Generate enhanced content using Claude AI (single request)."""
        prompt = self._build_enhancement_prompt(card)

        try:
            message = self.client.messages.create(
                model=self.config.model,
                max_tokens=1000,
                temperature=0.3,
                messages=[{"role": "user", "content": prompt}],
            )
            return self._parse_enhancement_response(message.content[0].text, card)

        except Exception as e:
            self.logger.error(f"Claude API error for '{card.word}': {e}")
            return None

    # --- Azure TTS ---

    def generate_tts_audio(
        self, text: str, voice_name: str, filename: str
    ) -> Optional[str]:
        """Generate TTS audio using Azure Speech SDK."""
        if not self.config.tts.enabled or not self._speech_config:
            return None

        audio_path = self.config.audio_dir / filename
        if audio_path.exists():
            console.print(f"[cyan]Using existing audio: {filename}[/cyan]")
            return str(audio_path)

        try:
            audio_config = speechsdk.audio.AudioOutputConfig(filename=str(audio_path))
            synthesizer = speechsdk.SpeechSynthesizer(
                speech_config=self._speech_config, audio_config=audio_config
            )

            # Extract language from voice name (e.g. zu-ZA-ThembaNeural -> zu-ZA)
            parts = voice_name.split("-")
            lang = f"{parts[0]}-{parts[1]}" if len(parts) >= 2 else "en-US"

            ssml = (
                f'<speak version="1.0" xmlns="http://www.w3.org/2001/10/synthesis" xml:lang="{lang}">'
                f'<voice name="{voice_name}">'
                f"{text}"
                f"</voice></speak>"
            )
            result = synthesizer.speak_ssml_async(ssml).get()

            if result.reason == speechsdk.ResultReason.SynthesizingAudioCompleted:
                console.print(f"[green]Generated audio: {filename}[/green]")
                return str(audio_path)
            else:
                details = result.cancellation_details
                self.logger.error(
                    f"TTS failed for '{text}': {details.reason} - {details.error_details}"
                )
                return None

        except Exception as e:
            self.logger.error(f"TTS error for '{text}': {e}")
            return None

    def generate_card_audio(
        self, card: CardData, enhanced: EnhancedContent, card_index: int
    ) -> AudioFiles:
        """Generate all audio files for a card (word + examples)."""
        audio_files = AudioFiles()

        if not self.config.tts.enabled:
            return audio_files

        try:
            # Generate male voice for main word
            word_filename = f"word_{card.id}.mp3"
            word_audio = self.generate_tts_audio(
                enhanced.word, self.config.tts.azure.male_voice, word_filename
            )
            if word_audio:
                audio_files.word = word_audio

            # Generate example sentence audio (alternate voices)
            for i, example in enumerate(enhanced.examples):
                use_male = (card_index + i) % 2 == 0
                voice = (
                    self.config.tts.azure.male_voice
                    if use_male
                    else self.config.tts.azure.female_voice
                )
                voice_label = "male" if use_male else "female"
                example_filename = f"example_{card.id}_{i}_{voice_label}.mp3"

                example_audio = self.generate_tts_audio(
                    example.target_language, voice, example_filename
                )
                if example_audio:
                    audio_files.examples.append(example_audio)
                    example.audio_file = example_audio

        except Exception as e:
            self.logger.error(f"Audio generation failed for card {card.id}: {e}")

        return audio_files

    def generate_all_audio_parallel(
        self, enhanced_cards: List[EnhancedCard], card_start_index: int = 0
    ) -> None:
        """Generate all audio files in parallel using ThreadPoolExecutor."""
        if not self.config.tts.enabled:
            return

        # Collect all audio tasks: (text, voice, filename, card_ref, slot)
        # slot is ("word",) or ("example", example_index)
        tasks = []
        for i, ec in enumerate(enhanced_cards):
            card_index = card_start_index + i
            card = ec.original

            # Word audio
            word_filename = f"word_{card.id}.mp3"
            tasks.append((
                ec, enhanced_cards, card_index,
                ec.enhanced.word, self.config.tts.azure.male_voice,
                word_filename, ("word",)
            ))

            # Example audio
            for j, example in enumerate(ec.enhanced.examples):
                use_male = (card_index + j) % 2 == 0
                voice = (
                    self.config.tts.azure.male_voice
                    if use_male
                    else self.config.tts.azure.female_voice
                )
                voice_label = "male" if use_male else "female"
                example_filename = f"example_{card.id}_{j}_{voice_label}.mp3"
                tasks.append((
                    ec, enhanced_cards, card_index,
                    example.target_language, voice,
                    example_filename, ("example", j)
                ))

        # Filter out tasks where audio already exists
        pending = []
        for task in tasks:
            audio_path = self.config.audio_dir / task[5]
            if not audio_path.exists():
                pending.append(task)
            else:
                # Apply existing path
                self._apply_audio_result(task, str(audio_path))

        if not pending:
            console.print("[cyan]All audio files already exist[/cyan]")
            return

        console.print(
            f"[blue]Generating {len(pending)} audio files "
            f"({len(tasks) - len(pending)} cached) "
            f"with {self.config.audio_concurrency} workers...[/blue]"
        )

        completed = 0
        failed = 0

        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
            console=console,
        ) as prog:
            progress_task = prog.add_task("Generating audio...", total=len(pending))

            with ThreadPoolExecutor(max_workers=self.config.audio_concurrency) as executor:
                future_to_task = {
                    executor.submit(
                        self.generate_tts_audio, t[3], t[4], t[5]
                    ): t
                    for t in pending
                }

                for future in as_completed(future_to_task):
                    task_info = future_to_task[future]
                    try:
                        result = future.result()
                        if result:
                            self._apply_audio_result(task_info, result)
                            completed += 1
                        else:
                            failed += 1
                    except Exception as e:
                        self.logger.error(f"Audio task failed: {e}")
                        failed += 1
                    prog.advance(progress_task)

        console.print(
            f"[green]Audio generation complete: {completed} succeeded, {failed} failed[/green]"
        )

    def _apply_audio_result(self, task_info: tuple, audio_path: str) -> None:
        """Apply an audio generation result back to the EnhancedCard."""
        ec = task_info[0]
        slot = task_info[6]

        if slot[0] == "word":
            ec.enhanced.audio_files.word = audio_path
        elif slot[0] == "example":
            example_idx = slot[1]
            # Ensure examples list is long enough
            while len(ec.enhanced.audio_files.examples) <= example_idx:
                ec.enhanced.audio_files.examples.append("")
            ec.enhanced.audio_files.examples[example_idx] = audio_path
            if example_idx < len(ec.enhanced.examples):
                ec.enhanced.examples[example_idx].audio_file = audio_path

    # --- Cleanup ---

    def clean_empty_audio(self) -> int:
        """Delete empty (0-byte) audio files from the audio directory."""
        removed = 0
        for f in self.config.audio_dir.glob("*.mp3"):
            if f.stat().st_size == 0:
                f.unlink()
                removed += 1
        if removed:
            console.print(f"[yellow]Removed {removed} empty audio files[/yellow]")
        else:
            console.print("[green]No empty audio files found[/green]")
        return removed

    # --- Progress ---

    def load_progress(self) -> ProcessingProgress:
        if self.config.progress_file.exists():
            try:
                with open(self.config.progress_file, "r") as f:
                    return ProcessingProgress(**json.load(f))
            except Exception:
                pass
        return ProcessingProgress()

    def save_progress(self, progress: ProcessingProgress):
        with open(self.config.progress_file, "w") as f:
            json.dump(progress.model_dump(), f, indent=2)

    # --- Batch API ---

    def load_batch_progress(self) -> BatchProgress:
        if self.config.batch_progress_file.exists():
            try:
                with open(self.config.batch_progress_file, "r") as f:
                    return BatchProgress(**json.load(f))
            except Exception:
                pass
        return BatchProgress()

    def save_batch_progress(self, progress: BatchProgress):
        with open(self.config.batch_progress_file, "w") as f:
            json.dump(progress.model_dump(), f, indent=2)

    def submit_batch(self, cards: List[CardData]) -> Optional[BatchInfo]:
        """Submit a batch of cards for Claude processing."""
        console.print(f"[blue]Submitting batch of {len(cards)} cards...[/blue]")

        requests = []
        card_id_mapping = {}

        for card in cards:
            sanitized_id = card.id
            card_id_mapping[sanitized_id] = card.id
            prompt = self._build_enhancement_prompt(card)

            requests.append({
                "custom_id": sanitized_id,
                "params": {
                    "model": self.config.model,
                    "max_tokens": 1000,
                    "temperature": 0.3,
                    "messages": [{"role": "user", "content": prompt}],
                },
            })

        try:
            batch_response = self.client.messages.batches.create(requests=requests)

            batch_info = BatchInfo(
                batch_id=batch_response.id,
                status=BatchProcessingStatus(batch_response.processing_status),
                created_at=str(batch_response.created_at),
                expires_at=str(batch_response.expires_at) if batch_response.expires_at else None,
                total_requests=len(requests),
                card_id_mapping=card_id_mapping,
            )

            console.print(f"[green]Batch submitted: {batch_info.batch_id}[/green]")
            return batch_info

        except Exception as e:
            self.logger.error(f"Failed to submit batch: {e}")
            console.print(f"[red]Failed to submit batch: {e}[/red]")
            return None

    def check_batch_status(self, batch_id: str) -> Optional[BatchInfo]:
        """Check the status of a submitted batch."""
        try:
            resp = self.client.messages.batches.retrieve(batch_id)

            status_mapping = {
                "in_progress": BatchProcessingStatus.IN_PROGRESS,
                "canceling": BatchProcessingStatus.CANCELING,
                "ended": BatchProcessingStatus.ENDED,
            }

            return BatchInfo(
                batch_id=resp.id,
                status=status_mapping.get(
                    resp.processing_status, BatchProcessingStatus.IN_PROGRESS
                ),
                created_at=str(resp.created_at),
                expires_at=str(resp.expires_at) if resp.expires_at else None,
                ended_at=str(resp.ended_at) if resp.ended_at else None,
                results_url=resp.results_url,
                total_requests=(
                    resp.request_counts.processing
                    + resp.request_counts.succeeded
                    + resp.request_counts.errored
                    + resp.request_counts.canceled
                    + resp.request_counts.expired
                ),
                succeeded=resp.request_counts.succeeded,
                failed=resp.request_counts.errored + resp.request_counts.canceled + resp.request_counts.expired,
                processing=resp.request_counts.processing,
            )

        except Exception as e:
            self.logger.error(f"Failed to check batch status: {e}")
            return None

    def download_batch_results(self, batch_info: BatchInfo) -> List[BatchResult]:
        """Download and parse batch results."""
        if not batch_info.results_url:
            self.logger.error("No results URL available")
            return []

        try:
            console.print("[blue]Downloading batch results...[/blue]")

            import httpx

            with httpx.Client() as client:
                response = client.get(
                    batch_info.results_url,
                    headers={
                        "x-api-key": os.getenv("ANTHROPIC_API_KEY"),
                        "anthropic-version": "2023-06-01",
                    },
                )
                response.raise_for_status()

            results = []
            for line in response.text.strip().split("\n"):
                if line:
                    data = json.loads(line)
                    results.append(
                        BatchResult(
                            custom_id=data["custom_id"],
                            result=data.get("result"),
                            error=data.get("error"),
                        )
                    )

            console.print(f"[green]Downloaded {len(results)} results[/green]")
            return results

        except Exception as e:
            self.logger.error(f"Failed to download batch results: {e}")
            console.print(f"[red]Failed to download results: {e}[/red]")
            return []

    def _build_batch(
        self,
        cards: List[CardData],
        start: int = 0,
        build_start: Optional[float] = None,
    ) -> bool:
        """Run the build pipeline using the Claude Batch API."""
        batch_progress = self.load_batch_progress()
        batch_progress.total_cards = len(cards)

        # Separate cached vs uncached cards
        cached_cards = []
        uncached_cards = []
        for card in cards:
            cached_path = self.config.cached_cards_dir / f"{card.id}.json"
            if cached_path.exists():
                try:
                    with open(cached_path, "r", encoding="utf-8") as f:
                        ec = EnhancedCard(**json.load(f))
                    cached_cards.append(ec)
                    continue
                except Exception:
                    pass
            uncached_cards.append(card)

        if cached_cards:
            console.print(f"[cyan]{len(cached_cards)} cards loaded from cache[/cyan]")

        if not uncached_cards:
            console.print("[cyan]All cards already cached[/cyan]")
            all_enhanced = cached_cards
        else:
            # Submit uncached cards in chunks of batch_api_size
            cards_map = {c.id: c for c in uncached_cards}

            for i in range(0, len(uncached_cards), self.config.batch_size):
                chunk = uncached_cards[i : i + self.config.batch_size]
                batch_info = self.submit_batch(chunk)
                if batch_info:
                    batch_progress.batches.append(batch_info)
                    batch_progress.cards_by_batch[batch_info.batch_id] = [
                        c.id for c in chunk
                    ]
                    self.save_batch_progress(batch_progress)

            # Poll until all batches complete
            while batch_progress.active_batches:
                active = batch_progress.active_batches
                console.print(
                    f"[blue]Waiting for {len(active)} batch(es)... "
                    f"checking every {self.config.batch_check_interval}s[/blue]"
                )

                for batch_info in active:
                    updated = self.check_batch_status(batch_info.batch_id)
                    if updated:
                        for j, b in enumerate(batch_progress.batches):
                            if b.batch_id == batch_info.batch_id:
                                # Preserve card_id_mapping from original submission
                                updated.card_id_mapping = b.card_id_mapping
                                batch_progress.batches[j] = updated
                                break

                        console.print(
                            f"  Batch {updated.batch_id[:12]}... "
                            f"status={updated.status.value} "
                            f"succeeded={updated.succeeded} "
                            f"failed={updated.failed} "
                            f"processing={updated.processing}"
                        )

                        if updated.status == BatchProcessingStatus.ENDED:
                            console.print(
                                f"[green]Batch {updated.batch_id[:12]}... completed![/green]"
                            )

                self.save_batch_progress(batch_progress)

                if batch_progress.active_batches:
                    time.sleep(self.config.batch_check_interval)

            # Download and process results from all completed batches
            new_enhanced = []
            for batch_info in batch_progress.batches:
                if batch_info.status != BatchProcessingStatus.ENDED:
                    continue
                if not batch_info.results_url:
                    # Re-fetch to get results_url
                    batch_info = self.check_batch_status(batch_info.batch_id)
                    if not batch_info or not batch_info.results_url:
                        continue

                results = self.download_batch_results(batch_info)
                for result in results:
                    card = cards_map.get(result.custom_id)
                    if not card:
                        continue

                    if result.error:
                        self.logger.error(
                            f"Batch error for {card.word}: {result.error}"
                        )
                        batch_progress.failed_cards.append(card.id)
                        continue

                    try:
                        response_text = result.result["message"]["content"][0]["text"]
                        enhanced = self._parse_enhancement_response(
                            response_text, card
                        )
                        if enhanced:
                            ec = EnhancedCard(original=card, enhanced=enhanced)
                            new_enhanced.append(ec)
                            batch_progress.completed_cards.append(card.id)
                        else:
                            batch_progress.failed_cards.append(card.id)
                    except Exception as e:
                        self.logger.error(
                            f"Failed to process batch result for {card.word}: {e}"
                        )
                        batch_progress.failed_cards.append(card.id)

            self.save_batch_progress(batch_progress)

            console.print(
                f"[green]Batch processing complete: "
                f"{len(new_enhanced)} succeeded, "
                f"{len(batch_progress.failed_cards)} failed[/green]"
            )

            # Generate audio in parallel for new cards
            if new_enhanced:
                self.generate_all_audio_parallel(new_enhanced, card_start_index=start)

                # Save to cache
                for ec in new_enhanced:
                    cached_path = self.config.cached_cards_dir / f"{ec.original.id}.json"
                    with open(cached_path, "w", encoding="utf-8") as f:
                        json.dump(ec.model_dump(), f, indent=2, ensure_ascii=False)

            all_enhanced = cached_cards + new_enhanced

        # Generate audio for cached cards that may be missing audio
        if cached_cards:
            self.generate_all_audio_parallel(cached_cards, card_start_index=start)
            for ec in cached_cards:
                cached_path = self.config.cached_cards_dir / f"{ec.original.id}.json"
                with open(cached_path, "w", encoding="utf-8") as f:
                    json.dump(ec.model_dump(), f, indent=2, ensure_ascii=False)

        console.print(
            f"\n[green]Enhanced {len(all_enhanced)}/{len(cards)} cards[/green]"
        )

        if all_enhanced:
            self.create_mochi_file(all_enhanced)

        self._print_timing_summary(build_start, len(all_enhanced))
        return len(all_enhanced) > 0

    # --- Review Data Preservation ---

    def extract_review_data_from_mochi(self, mochi_path: Path) -> Dict[str, Dict]:
        """Extract review data from an existing .mochi file."""
        review_data = {}
        if not mochi_path.exists():
            return review_data

        temp_dir = self.config.build_dir / "temp_extraction"
        try:
            temp_dir.mkdir(exist_ok=True)
            with zipfile.ZipFile(mochi_path, "r") as zf:
                zf.extractall(temp_dir)

            data_file = temp_dir / "data.json"
            if not data_file.exists():
                return review_data

            with open(data_file, "r", encoding="utf-8") as f:
                data = json.load(f)

            for deck in data.get("~:decks", []):
                if not isinstance(deck, dict):
                    continue
                cards = deck.get("~:cards", {})
                card_list = (
                    cards.get("~#list", [])
                    if isinstance(cards, dict)
                    else cards if isinstance(cards, list) else []
                )

                for card in card_list:
                    if not isinstance(card, dict):
                        continue
                    fields = card.get("~:fields", {})
                    name_field = fields.get("~:name", {})
                    card_name = (
                        name_field.get("~:value", "")
                        if isinstance(name_field, dict)
                        else ""
                    )
                    # Strip audio markers from name
                    card_name_clean = re.sub(
                        r"\n\n!\[[^\]]*\]\([^)]+\)", "", card_name
                    ).strip()

                    if not card_name_clean:
                        continue

                    card_review = {
                        "reviews": card.get("~:reviews", []),
                        "reverse_reviews": card.get("~:reverse-reviews", []),
                        "cloze_reviews": card.get("~:cloze/reviews", {}),
                        "needs_rereview": card.get("~:needs-rereview?", False),
                        "reverse_needs_rereview": card.get(
                            "~:reverse/needs-rereview?", False
                        ),
                        "cloze_needs_rereview": card.get(
                            "~:cloze/needs-rereview?", {}
                        ),
                        "created_at": card.get("~:created-at", {}),
                        "cloze_indexes": card.get(
                            "~:cloze/indexes", {"~#set": []}
                        ),
                        "references": card.get("~:references", {"~#set": []}),
                    }

                    if (
                        card_review["reviews"]
                        or card_review["reverse_reviews"]
                        or card_review["cloze_reviews"]
                    ):
                        review_data[card_name_clean] = card_review

        except Exception as e:
            self.logger.error(f"Failed to extract review data: {e}")
        finally:
            shutil.rmtree(temp_dir, ignore_errors=True)

        return review_data

    def load_existing_review_data(self) -> Dict[str, Dict]:
        """Load review data from all .mochi files in the input directory."""
        all_review_data = {}
        if not self.config.input_dir.exists():
            return all_review_data

        for mochi_file in self.config.input_dir.glob("*.mochi"):
            console.print(f"[blue]Loading review data from {mochi_file.name}...[/blue]")
            review_data = self.extract_review_data_from_mochi(mochi_file)
            all_review_data.update(review_data)

        if all_review_data:
            console.print(
                f"[green]Found review data for {len(all_review_data)} cards[/green]"
            )
        return all_review_data

    # --- Mochi File Creation ---

    def create_mochi_file(
        self, enhanced_cards: List[EnhancedCard], deck_name: Optional[str] = None
    ) -> bool:
        """Create a .mochi file from enhanced cards."""
        try:
            final_deck_name = deck_name or self.config.deck_name
            existing_review_data = self.load_existing_review_data()

            def det_id(name: str) -> str:
                return f"~:{hashlib.md5(name.encode()).hexdigest()[:8]}"

            deck_id = det_id(f"deck_{final_deck_name}")
            template_id = det_id(f"template_{final_deck_name}")
            timestamp_ms = int(datetime.now().timestamp() * 1000)

            # Field IDs
            english_fid = det_id("english_field")
            word_helper_fid = det_id("word_helper_field")
            pos_fid = det_id("main_pos_field")
            add_meanings_fid = det_id("add_meanings_field")
            roman_fid = det_id("romanization_field")
            ipa_fid = det_id("ipa_field")
            examples_fid = det_id("examples_field")
            etymology_fid = det_id("etymology_field")
            final_english_fid = det_id("final_english_field")

            template = {
                "~:id": template_id,
                "~:pos": "A",
                "~:name": final_deck_name,
                "~:content": (
                    "# <<Word>>\n\n---\n\n### <<WordHelper>>\n\n"
                    "(<<MainPartOfSpeech>> <<English>>)\n<<AdditionalMeanings>>\n\n"
                    "<<Romanization>>\n<<IPA>>\n\n"
                    "**Examples:**\n\n<<Examples>>\n\n"
                    "**Etymology:**\n<<Etymology>>\n\n"
                    "---\n\n# <<FinalEnglish>>"
                ),
                "~:cloze?": None,
                "~:fields": {
                    "~:name": {
                        "~:id": "~:name",
                        "~:pos": "k0",
                        "~:name": "Word",
                        "~:type": "~:text",
                    },
                    word_helper_fid: {
                        "~:id": word_helper_fid,
                        "~:pos": "k1",
                        "~:name": "WordHelper",
                        "~:type": "~:text",
                    },
                    english_fid: {
                        "~:id": english_fid,
                        "~:pos": "k2",
                        "~:name": "English",
                        "~:type": "~:text",
                    },
                    pos_fid: {
                        "~:id": pos_fid,
                        "~:pos": "k3",
                        "~:name": "MainPartOfSpeech",
                        "~:type": "~:text",
                    },
                    add_meanings_fid: {
                        "~:id": add_meanings_fid,
                        "~:pos": "k4",
                        "~:name": "AdditionalMeanings",
                        "~:type": "~:text",
                    },
                    roman_fid: {
                        "~:id": roman_fid,
                        "~:pos": "k5",
                        "~:name": "Romanization",
                        "~:type": "~:text",
                    },
                    ipa_fid: {
                        "~:id": ipa_fid,
                        "~:pos": "k6",
                        "~:name": "IPA",
                        "~:type": "~:text",
                    },
                    examples_fid: {
                        "~:id": examples_fid,
                        "~:pos": "k7",
                        "~:name": "Examples",
                        "~:type": "~:text",
                    },
                    etymology_fid: {
                        "~:id": etymology_fid,
                        "~:pos": "k8",
                        "~:name": "Etymology",
                        "~:type": "~:text",
                    },
                    final_english_fid: {
                        "~:id": final_english_fid,
                        "~:pos": "k9",
                        "~:name": "FinalEnglish",
                        "~:type": "~:text",
                    },
                },
            }

            def generate_pos(index):
                chars = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
                if index < len(chars):
                    return chars[index]
                result = ""
                while index >= len(chars):
                    result = chars[index % len(chars)] + result
                    index = index // len(chars)
                result = chars[index] + result
                return result

            mochi_cards = []
            audio_files_to_include = []
            used_filenames = set()

            for i, ec in enumerate(enhanced_cards):
                card_id = det_id(f"card_{ec.original.id}")
                card_attachments = {}

                # --- Word audio ---
                word_audio_link = ""
                if ec.enhanced.audio_files.word and Path(ec.enhanced.audio_files.word).exists():
                    src = Path(ec.enhanced.audio_files.word)
                    unique_fn = f"{card_id.replace('~:', '')}_word{src.suffix}"
                    if unique_fn in used_filenames:
                        unique_fn = f"{card_id.replace('~:', '')}_word_1{src.suffix}"
                    used_filenames.add(unique_fn)

                    word_audio_link = f"\n\n![audio]({unique_fn})"
                    audio_files_to_include.append((str(src), unique_fn))
                    card_attachments[unique_fn] = {
                        "~:name": unique_fn,
                        "~:type": "audio/mpeg",
                        "~:size": src.stat().st_size,
                    }

                # --- Example audio ---
                examples_text = ""
                for j, ex in enumerate(ec.enhanced.examples):
                    example_text = f"{ex.target_language}\n\n*{ex.english}*"

                    if ex.audio_file and Path(ex.audio_file).exists():
                        src = Path(ex.audio_file)
                        unique_fn = f"{card_id.replace('~:', '')}_ex_{j}{src.suffix}"
                        if unique_fn in used_filenames:
                            unique_fn = f"{card_id.replace('~:', '')}_ex_{j}_1{src.suffix}"
                        used_filenames.add(unique_fn)

                        example_text += f"\n\n![audio]({unique_fn})"
                        audio_files_to_include.append((str(src), unique_fn))
                        card_attachments[unique_fn] = {
                            "~:name": unique_fn,
                            "~:type": "audio/mpeg",
                            "~:size": src.stat().st_size,
                        }

                    examples_text += example_text
                    if j < len(ec.enhanced.examples) - 1:
                        examples_text += "\n\n"

                # --- Review data ---
                word_clean = ec.enhanced.word.strip()
                existing_reviews = existing_review_data.get(word_clean, {})

                if existing_reviews:
                    reviews = existing_reviews.get("reviews", [])
                    reverse_reviews = existing_reviews.get("reverse_reviews", [])
                    cloze_reviews = existing_reviews.get("cloze_reviews", {})
                    needs_rereview = existing_reviews.get("needs_rereview", False)
                    reverse_needs_rereview = existing_reviews.get(
                        "reverse_needs_rereview", False
                    )
                    cloze_needs_rereview = existing_reviews.get(
                        "cloze_needs_rereview", {}
                    )
                    cloze_indexes = existing_reviews.get(
                        "cloze_indexes", {"~#set": []}
                    )
                    references = existing_reviews.get("references", {"~#set": []})
                    created_at = existing_reviews.get(
                        "created_at", {"~#dt": timestamp_ms}
                    )
                else:
                    reviews = []
                    reverse_reviews = []
                    cloze_reviews = {}
                    needs_rereview = False
                    reverse_needs_rereview = False
                    cloze_needs_rereview = {}
                    cloze_indexes = {"~#set": []}
                    references = {"~#set": []}
                    created_at = {"~#dt": timestamp_ms}

                # Additional meanings text
                add_meanings_text = ""
                if ec.enhanced.additional_meanings:
                    add_meanings_text = "\n".join(
                        f"({m.part_of_speech} {m.meaning})"
                        for m in ec.enhanced.additional_meanings
                    )

                final_english = ec.original.english
                if ec.enhanced.additional_meanings:
                    final_english += ", " + ", ".join(
                        m.meaning for m in ec.enhanced.additional_meanings
                    )

                mochi_card = {
                    "~:tags": {"~#set": []},
                    "~:content": "",
                    "~:cloze/reviews": cloze_reviews,
                    "~:name": ec.enhanced.word,
                    "~:deck-id": deck_id,
                    "~:needs-rereview?": needs_rereview,
                    "~:fields": {
                        "~:name": {
                            "~:id": "~:name",
                            "~:value": ec.enhanced.word + word_audio_link,
                        },
                        word_helper_fid: {
                            "~:id": word_helper_fid,
                            "~:value": ec.original.word + word_audio_link,
                        },
                        english_fid: {
                            "~:id": english_fid,
                            "~:value": ec.original.english,
                        },
                        pos_fid: {
                            "~:id": pos_fid,
                            "~:value": ec.enhanced.main_part_of_speech,
                        },
                        add_meanings_fid: {
                            "~:id": add_meanings_fid,
                            "~:value": add_meanings_text,
                        },
                        roman_fid: {
                            "~:id": roman_fid,
                            "~:value": ec.enhanced.romanization,
                        },
                        ipa_fid: {
                            "~:id": ipa_fid,
                            "~:value": ec.enhanced.pronunciation_ipa,
                        },
                        examples_fid: {
                            "~:id": examples_fid,
                            "~:value": examples_text,
                        },
                        etymology_fid: {
                            "~:id": etymology_fid,
                            "~:value": ec.enhanced.etymology,
                        },
                        final_english_fid: {
                            "~:id": final_english_fid,
                            "~:value": final_english,
                        },
                    },
                    "~:attachments": card_attachments,
                    "~:cloze/indexes": cloze_indexes,
                    "~:pos": generate_pos(i),
                    "~:references": references,
                    "~:reverse-reviews": reverse_reviews,
                    "~:id": card_id,
                    "~:reviews": reviews,
                    "~:reverse/needs-rereview?": reverse_needs_rereview,
                    "~:cloze/needs-rereview?": cloze_needs_rereview,
                    "~:created-at": created_at,
                    "~:template-id": template_id,
                }
                mochi_cards.append(mochi_card)

            # Build deck structure
            mochi_data = {
                "~:version": 2,
                "~:templates": {"~#list": [template]},
                "~:decks": [
                    {
                        "~:name": final_deck_name,
                        "~:id": deck_id,
                        "~:sort": 1,
                        "~:cards": {"~#list": mochi_cards},
                    }
                ],
                "~:files": [],
            }

            safe_name = "".join(
                c for c in final_deck_name if c.isalnum() or c in (" ", "-", "_")
            ).replace(" ", "_")
            output_path = self.config.output_dir / f"{safe_name}.mochi"

            with zipfile.ZipFile(output_path, "w", zipfile.ZIP_DEFLATED) as zf:
                json_content = json.dumps(mochi_data, ensure_ascii=False, indent=4)
                zf.writestr("data.json", json_content.encode("utf-8"))

                added = 0
                for audio_src, unique_fn in audio_files_to_include:
                    src_path = Path(audio_src)
                    if src_path.exists():
                        zf.write(src_path, unique_fn)
                        added += 1

            console.print(f"[green]Created {output_path}[/green]")
            console.print(
                f"[green]  {len(enhanced_cards)} cards, {added} audio files[/green]"
            )

            cards_with_reviews = len(
                [c for c in enhanced_cards if c.enhanced.word.strip() in existing_review_data]
            )
            if cards_with_reviews > 0:
                console.print(
                    f"[cyan]Preserved review data for {cards_with_reviews} cards[/cyan]"
                )

            return True

        except Exception as e:
            import traceback

            self.logger.error(f"Failed to create .mochi file: {e}")
            self.logger.error(traceback.format_exc())
            console.print(f"[red]Failed to create .mochi file: {e}[/red]")
            return False

    # --- Timing ---

    def _print_timing_summary(self, build_start: Optional[float], card_count: int) -> None:
        """Print a timing summary at the end of a build."""
        if build_start is None or card_count == 0:
            return
        elapsed = time.time() - build_start
        minutes, seconds = divmod(int(elapsed), 60)
        per_card = elapsed / card_count
        console.print(
            f"\n[bold green]Done in {minutes}m {seconds}s "
            f"({per_card:.1f}s/card, {card_count} cards)[/bold green]"
        )

    # --- Main Build ---

    def build(
        self,
        size: Optional[int] = None,
        render_only: bool = False,
        start: int = 0,
    ) -> bool:
        """Run the full build pipeline in batches."""
        self.clean_empty_audio()
        cards = self.load_cards()
        if not cards:
            console.print("[red]No cards to process[/red]")
            return False

        if size:
            cards = cards[start : start + size]
        elif start > 0:
            cards = cards[start:]

        console.print(
            f"[blue]Processing {len(cards)} cards with model {self.config.model} "
            f"(batch_size={self.config.batch_size})[/blue]"
        )

        if render_only:
            return self._render_from_cache(cards)

        build_start = time.time()

        if self.config.use_batch_api:
            return self._build_batch(cards, start=start, build_start=build_start)

        progress = self.load_progress()
        progress.total_cards = len(cards)
        all_enhanced_cards = []
        batch_size = self.config.batch_size

        # Split cards into batches
        batches = [
            cards[i : i + batch_size] for i in range(0, len(cards), batch_size)
        ]

        for batch_num, batch in enumerate(batches):
            batch_start_index = start + batch_num * batch_size
            console.print(
                f"\n[blue]Batch {batch_num + 1}/{len(batches)} "
                f"({len(batch)} cards)[/blue]"
            )

            # Phase 1: Claude enhancement
            batch_enhanced = []
            with Progress(
                SpinnerColumn(),
                TextColumn("[progress.description]{task.description}"),
                BarColumn(),
                TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
                console=console,
            ) as prog:
                task = prog.add_task("Enhancing...", total=len(batch))

                for card in batch:
                    prog.update(task, description=f"Enhancing: {card.word[:30]}...")

                    # Check cache first
                    cached_path = self.config.cached_cards_dir / f"{card.id}.json"
                    if cached_path.exists() and card.id not in progress.failed_ids:
                        try:
                            with open(cached_path, "r", encoding="utf-8") as f:
                                cached_data = json.load(f)
                            ec = EnhancedCard(**cached_data)
                            batch_enhanced.append(ec)
                            if card.id not in progress.processed_ids:
                                progress.processed_ids.append(card.id)
                            prog.advance(task)
                            continue
                        except Exception:
                            pass

                    enhanced = self.enhance_card_with_claude(card)
                    if enhanced:
                        ec = EnhancedCard(original=card, enhanced=enhanced)
                        batch_enhanced.append(ec)
                        progress.processed_ids.append(card.id)
                    else:
                        progress.failed_ids.append(card.id)

                    prog.advance(task)
                    time.sleep(self.config.rate_limit_delay)

            # Phase 2: Parallel audio generation for this batch
            if batch_enhanced:
                self.generate_all_audio_parallel(
                    batch_enhanced, card_start_index=batch_start_index
                )

                # Save cache with audio paths
                for ec in batch_enhanced:
                    cached_path = self.config.cached_cards_dir / f"{ec.original.id}.json"
                    with open(cached_path, "w", encoding="utf-8") as f:
                        json.dump(ec.model_dump(), f, indent=2, ensure_ascii=False)

            all_enhanced_cards.extend(batch_enhanced)
            self.save_progress(progress)

        console.print(
            f"\n[green]Enhanced {len(all_enhanced_cards)}/{len(cards)} cards[/green]"
        )

        if all_enhanced_cards:
            self.create_mochi_file(all_enhanced_cards)

        self._print_timing_summary(build_start, len(all_enhanced_cards))
        return len(all_enhanced_cards) > 0

    def _render_from_cache(self, cards: List[CardData]) -> bool:
        """Re-render .mochi file from cached cards only."""
        enhanced_cards = []
        for card in cards:
            cached_path = self.config.cached_cards_dir / f"{card.id}.json"
            if cached_path.exists():
                try:
                    with open(cached_path, "r", encoding="utf-8") as f:
                        ec = EnhancedCard(**json.load(f))
                    enhanced_cards.append(ec)
                except Exception as e:
                    self.logger.error(f"Failed to load cached card {card.id}: {e}")

        console.print(f"[blue]Loaded {len(enhanced_cards)} cached cards[/blue]")
        if enhanced_cards:
            return self.create_mochi_file(enhanced_cards)
        return False