"""Storage utilities for translations and progress tracking."""

from __future__ import annotations

import json
from datetime import datetime
from pathlib import Path
from typing import TYPE_CHECKING, Any

from src.models.verse import TranslationOutput, VerseReference

if TYPE_CHECKING:
    from src.orchestration.provenance import TranslationProvenance
    from src.protos.converters import (
        TranslationPromptInputProto,
        TranslationResultProto,
    )


class Storage:
    """Manages storage of translations and progress tracking."""

    def __init__(
        self,
        data_dir: Path | str,
        enable_archive: bool = False,
    ) -> None:
        """Initialize storage.

        Args:
            data_dir: Path to the data directory.
            enable_archive: Whether to enable proto archive storage.
        """
        self.data_dir = Path(data_dir)
        self.output_dir = self.data_dir / "output"
        self.progress_dir = self.data_dir / "progress"
        self.archive_dir = self.data_dir / "archive"
        self._enable_archive = enable_archive
        self._ensure_dirs()

    @property
    def enable_archive(self) -> bool:
        """Get archive enabled status."""
        return self._enable_archive

    @enable_archive.setter
    def enable_archive(self, value: bool) -> None:
        """Set archive enabled status."""
        self._enable_archive = value
        if value:
            self.archive_dir.mkdir(parents=True, exist_ok=True)

    def _ensure_dirs(self) -> None:
        """Ensure required directories exist."""
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.progress_dir.mkdir(parents=True, exist_ok=True)
        if self._enable_archive:
            self.archive_dir.mkdir(parents=True, exist_ok=True)

    def _load_json(self, path: Path) -> dict[str, Any]:
        """Load a JSON file."""
        if not path.exists():
            return {}
        with path.open("r", encoding="utf-8") as f:
            return json.load(f)

    def _save_json(self, path: Path, data: dict[str, Any]) -> None:
        """Save data to a JSON file."""
        path.parent.mkdir(parents=True, exist_ok=True)
        with path.open("w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False, indent=2, default=str)

    def save_translation(self, output: TranslationOutput) -> None:
        """Save a translation output.

        Args:
            output: The translation output to save.
        """
        book = output.reference.book
        chapter = output.reference.chapter
        verse = output.reference.verse

        # Load existing chapter file or create new
        chapter_file = self.output_dir / book / f"chapter_{chapter:02d}.json"
        chapter_data = self._load_json(chapter_file)

        if "verses" not in chapter_data:
            chapter_data["verses"] = {}
            chapter_data["book"] = book
            chapter_data["chapter"] = chapter

        verse_key = str(verse)
        chapter_data["verses"][verse_key] = {
            "text": output.text,
            "text_with_furigana": output.text_with_furigana,
            "furigana_annotations": output.furigana_annotations,
            "model": output.model,
            "timestamp": output.timestamp.isoformat(),
            "prompt_version": output.prompt_version,
            "sources_used": output.sources_used,
            "glossary_terms_applied": output.glossary_terms_applied,
        }

        self._save_json(chapter_file, chapter_data)
        self._update_progress(book, chapter, verse)

    def load_translation(self, reference: VerseReference) -> TranslationOutput | None:
        """Load a saved translation.

        Args:
            reference: The verse reference to load.

        Returns:
            TranslationOutput if found, None otherwise.
        """
        chapter_file = (
            self.output_dir / reference.book / f"chapter_{reference.chapter:02d}.json"
        )
        chapter_data = self._load_json(chapter_file)

        verses = chapter_data.get("verses", {})
        verse_data = verses.get(str(reference.verse))
        if not verse_data:
            return None

        return TranslationOutput(
            reference=reference,
            text=verse_data["text"],
            text_with_furigana=verse_data.get("text_with_furigana", ""),
            furigana_annotations=verse_data.get("furigana_annotations", []),
            model=verse_data["model"],
            timestamp=datetime.fromisoformat(verse_data["timestamp"]),
            prompt_version=verse_data.get("prompt_version", "v1.0"),
            sources_used=verse_data.get("sources_used", []),
            glossary_terms_applied=verse_data.get("glossary_terms_applied", []),
        )

    def load_chapter_translations(
        self, book: str, chapter: int
    ) -> dict[int, TranslationOutput]:
        """Load all translations for a chapter.

        Returns:
            Dictionary mapping verse numbers to TranslationOutput.
        """
        chapter_file = self.output_dir / book / f"chapter_{chapter:02d}.json"
        chapter_data = self._load_json(chapter_file)

        translations = {}
        verses = chapter_data.get("verses", {})
        for verse_num_str, verse_data in verses.items():
            verse_num = int(verse_num_str)
            ref = VerseReference(book=book, chapter=chapter, verse=verse_num)
            translations[verse_num] = TranslationOutput(
                reference=ref,
                text=verse_data["text"],
                text_with_furigana=verse_data.get("text_with_furigana", ""),
                furigana_annotations=verse_data.get("furigana_annotations", []),
                model=verse_data["model"],
                timestamp=datetime.fromisoformat(verse_data["timestamp"]),
                prompt_version=verse_data.get("prompt_version", "v1.0"),
                sources_used=verse_data.get("sources_used", []),
                glossary_terms_applied=verse_data.get("glossary_terms_applied", []),
            )

        return translations

    def _update_progress(self, book: str, chapter: int, verse: int) -> None:
        """Update progress tracking for a book."""
        progress_file = self.progress_dir / f"{book}_progress.json"
        progress = self._load_json(progress_file)

        if "completed" not in progress:
            progress["completed"] = {}
            progress["book"] = book
            progress["last_updated"] = None

        chapter_key = str(chapter)
        if chapter_key not in progress["completed"]:
            progress["completed"][chapter_key] = []

        if verse not in progress["completed"][chapter_key]:
            progress["completed"][chapter_key].append(verse)
            progress["completed"][chapter_key].sort()

        progress["last_updated"] = datetime.utcnow().isoformat()
        self._save_json(progress_file, progress)

    def get_progress(self, book: str) -> dict[str, Any]:
        """Get progress for a book.

        Returns:
            Progress dictionary with completed verses per chapter.
        """
        progress_file = self.progress_dir / f"{book}_progress.json"
        return self._load_json(progress_file)

    def is_verse_translated(self, reference: VerseReference) -> bool:
        """Check if a verse has already been translated."""
        progress = self.get_progress(reference.book)
        completed = progress.get("completed", {})
        chapter_verses = completed.get(str(reference.chapter), [])
        return reference.verse in chapter_verses

    def get_completion_stats(self, book: str, total_verses: int) -> dict[str, Any]:
        """Get completion statistics for a book.

        Args:
            book: Book name.
            total_verses: Total verses expected in the book.

        Returns:
            Dictionary with completion statistics.
        """
        progress = self.get_progress(book)
        completed = progress.get("completed", {})

        completed_count = sum(len(verses) for verses in completed.values())
        percentage = (completed_count / total_verses * 100) if total_verses > 0 else 0

        return {
            "book": book,
            "completed": completed_count,
            "total": total_verses,
            "percentage": round(percentage, 2),
            "chapters_started": len(completed),
        }

    def save_translation_with_provenance(  # pragma: no cover
        self,
        output: TranslationOutput,
        provenance: TranslationProvenance,
    ) -> None:
        """Save a translation output with full provenance tracking.

        Args:
            output: The translation output to save.
            provenance: Full provenance record for this translation.
        """
        book = output.reference.book
        chapter = output.reference.chapter
        verse = output.reference.verse

        # Load existing chapter file or create new
        chapter_file = self.output_dir / book / f"chapter_{chapter:02d}.json"
        chapter_data = self._load_json(chapter_file)

        if "verses" not in chapter_data:
            chapter_data["verses"] = {}
            chapter_data["book"] = book
            chapter_data["chapter"] = chapter

        verse_key = str(verse)
        chapter_data["verses"][verse_key] = {
            "text": output.text,
            "text_with_furigana": output.text_with_furigana,
            "furigana_annotations": output.furigana_annotations,
            "model": output.model,
            "timestamp": output.timestamp.isoformat(),
            "prompt_version": output.prompt_version,
            "sources_used": output.sources_used,
            "glossary_terms_applied": output.glossary_terms_applied,
            "provenance": provenance.to_dict(),
        }

        self._save_json(chapter_file, chapter_data)
        self._update_progress(book, chapter, verse)

    def load_provenance(  # pragma: no cover
        self, reference: VerseReference
    ) -> TranslationProvenance | None:
        """Load provenance for a saved translation.

        Args:
            reference: The verse reference to load provenance for.

        Returns:
            TranslationProvenance if found, None otherwise.
        """
        from src.orchestration.provenance import TranslationProvenance

        chapter_file = (
            self.output_dir / reference.book / f"chapter_{reference.chapter:02d}.json"
        )
        chapter_data = self._load_json(chapter_file)

        verses = chapter_data.get("verses", {})
        verse_data = verses.get(str(reference.verse))
        if not verse_data:
            return None

        provenance_data = verse_data.get("provenance")
        if not provenance_data:
            return None

        return TranslationProvenance.from_dict(provenance_data)

    # =========================================================================
    # Proto Archive Methods
    # =========================================================================

    def _get_archive_path(self, reference: VerseReference) -> Path:
        """Get the archive directory path for a verse.

        Args:
            reference: The verse reference.

        Returns:
            Path to the archive directory for this verse.
        """
        return (
            self.archive_dir
            / reference.book
            / f"chapter_{reference.chapter:02d}"
        )

    def save_translation_archive(
        self,
        reference: VerseReference,
        input_proto: TranslationPromptInputProto,
        output_proto: TranslationResultProto,
    ) -> Path:
        """Save translation input/output protos to archive.

        Args:
            reference: The verse reference.
            input_proto: The input proto with all translation inputs.
            output_proto: The output proto with translation results.

        Returns:
            Path to the archive directory where files were saved.
        """
        from src.protos.serialization import serialize_proto

        archive_path = self._get_archive_path(reference)
        archive_path.mkdir(parents=True, exist_ok=True)

        verse_num = reference.verse
        input_file = archive_path / f"verse_{verse_num:02d}_input.msgpack"
        output_file = archive_path / f"verse_{verse_num:02d}_output.msgpack"

        input_bytes = serialize_proto(input_proto)
        output_bytes = serialize_proto(output_proto)

        with input_file.open("wb") as f:
            f.write(input_bytes)

        with output_file.open("wb") as f:
            f.write(output_bytes)

        return archive_path

    def load_translation_archive(
        self,
        reference: VerseReference,
    ) -> tuple[TranslationPromptInputProto, TranslationResultProto] | None:
        """Load translation input/output protos from archive.

        Args:
            reference: The verse reference.

        Returns:
            Tuple of (input_proto, output_proto) if found, None otherwise.
        """
        from src.protos.converters import (
            TranslationPromptInputProto,
            TranslationResultProto,
        )
        from src.protos.serialization import deserialize_proto

        archive_path = self._get_archive_path(reference)
        verse_num = reference.verse
        input_file = archive_path / f"verse_{verse_num:02d}_input.msgpack"
        output_file = archive_path / f"verse_{verse_num:02d}_output.msgpack"

        if not input_file.exists() or not output_file.exists():
            return None

        with input_file.open("rb") as f:
            input_bytes = f.read()

        with output_file.open("rb") as f:
            output_bytes = f.read()

        input_proto = deserialize_proto(input_bytes, TranslationPromptInputProto)
        output_proto = deserialize_proto(output_bytes, TranslationResultProto)

        return input_proto, output_proto