"""Storage utilities for translation evaluations."""

from __future__ import annotations

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

from src.models.evaluation import (
    BookEvaluationSummary,
    ChapterEvaluation,
    EvaluationScore,
    VerseEvaluation,
)
from src.models.verse import VerseReference


class EvaluationStorage:
    """Manages persistence of translation evaluations."""

    def __init__(self, data_dir: Path | str) -> None:
        """Initialize the evaluation storage.

        Args:
            data_dir: Path to the data directory.
        """
        self.data_dir = Path(data_dir)
        self.output_dir = self.data_dir / "output"

    def _get_eval_path(self, book: str, chapter: int) -> Path:
        """Get path to evaluation file for a chapter."""
        return self.output_dir / book / f"chapter_{chapter:02d}_eval.json"

    def _serialize_datetime(self, obj: Any) -> Any:
        """Serialize datetime objects for JSON."""
        if isinstance(obj, datetime):
            return obj.isoformat()
        raise TypeError(f"Object of type {type(obj)} is not JSON serializable")

    def save_chapter_evaluation(self, evaluation: ChapterEvaluation) -> Path:
        """Save a chapter evaluation to JSON.

        Args:
            evaluation: ChapterEvaluation to save.

        Returns:
            Path to the saved file.
        """
        path = self._get_eval_path(evaluation.book, evaluation.chapter)
        path.parent.mkdir(parents=True, exist_ok=True)

        # Convert to dict for JSON serialization
        data = {
            "book": evaluation.book,
            "chapter": evaluation.chapter,
            "evaluator_model": evaluation.evaluator_model,
            "evaluated_at": evaluation.evaluated_at.isoformat(),
            "verses": {},
            "aggregate": {
                "scores": evaluation.aggregate_scores,
                "overall": evaluation.overall_score,
                "flagged_count": evaluation.flagged_count,
                "flagged_verses": evaluation.flagged_verses,
            },
        }

        for verse_num, verse_eval in evaluation.verses.items():
            data["verses"][str(verse_num)] = {
                "scores": {
                    dim: {
                        "score": score.score,
                        "confidence": score.confidence,
                        "reasoning": score.reasoning,
                    }
                    for dim, score in verse_eval.scores.items()
                },
                "overall_score": verse_eval.overall_score,
                "flags": verse_eval.flags,
                "notes": verse_eval.notes,
                "timestamp": verse_eval.timestamp.isoformat(),
            }

        with path.open("w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False, indent=2)

        return path

    def load_chapter_evaluation(
        self, book: str, chapter: int
    ) -> ChapterEvaluation | None:
        """Load a chapter evaluation from JSON.

        Args:
            book: Book name.
            chapter: Chapter number.

        Returns:
            ChapterEvaluation if found, None otherwise.
        """
        path = self._get_eval_path(book, chapter)
        if not path.exists():
            return None

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

        verses = {}
        for verse_num_str, verse_data in data.get("verses", {}).items():
            verse_num = int(verse_num_str)
            scores = {}
            for dim, score_data in verse_data.get("scores", {}).items():
                scores[dim] = EvaluationScore(
                    dimension=dim,
                    score=score_data["score"],
                    confidence=score_data["confidence"],
                    reasoning=score_data["reasoning"],
                )

            verses[verse_num] = VerseEvaluation(
                verse_number=verse_num,
                evaluator_model=data.get("evaluator_model", "unknown"),
                timestamp=datetime.fromisoformat(
                    verse_data.get("timestamp", datetime.utcnow().isoformat())
                ),
                scores=scores,
                flags=verse_data.get("flags", []),
                notes=verse_data.get("notes"),
            )

        return ChapterEvaluation(
            book=data["book"],
            chapter=data["chapter"],
            evaluator_model=data.get("evaluator_model", "unknown"),
            evaluated_at=datetime.fromisoformat(
                data.get("evaluated_at", datetime.utcnow().isoformat())
            ),
            verses=verses,
        )

    def get_flagged_verses(self, book: str) -> list[VerseReference]:
        """Get all flagged verses for a book.

        Args:
            book: Book name.

        Returns:
            List of VerseReference objects for flagged verses.
        """
        flagged = []
        book_dir = self.output_dir / book

        if not book_dir.exists():
            return flagged

        for eval_file in sorted(book_dir.glob("chapter_*_eval.json")):
            chapter_eval = self.load_chapter_evaluation(
                book, int(eval_file.stem.split("_")[1])
            )
            if chapter_eval:
                for verse_num in chapter_eval.flagged_verses:
                    flagged.append(
                        VerseReference(
                            book=book, chapter=chapter_eval.chapter, verse=verse_num
                        )
                    )

        return flagged

    def get_book_summary(self, book: str) -> BookEvaluationSummary | None:
        """Generate a summary of evaluations for an entire book.

        Args:
            book: Book name.

        Returns:
            BookEvaluationSummary if evaluations exist, None otherwise.
        """
        book_dir = self.output_dir / book
        if not book_dir.exists():
            return None

        eval_files = list(book_dir.glob("chapter_*_eval.json"))
        if not eval_files:
            return None

        all_scores: dict[str, list[float]] = {}
        total_verses = 0
        flagged_refs: list[str] = []

        for eval_file in sorted(eval_files):
            chapter_num = int(eval_file.stem.split("_")[1])
            chapter_eval = self.load_chapter_evaluation(book, chapter_num)

            if chapter_eval:
                total_verses += len(chapter_eval.verses)

                # Aggregate dimension scores
                for dim, score in chapter_eval.aggregate_scores.items():
                    if dim not in all_scores:
                        all_scores[dim] = []
                    all_scores[dim].append(score)

                # Collect flagged verses
                for verse_num in chapter_eval.flagged_verses:
                    flagged_refs.append(f"{book}:{chapter_num}:{verse_num}")

        if not all_scores:
            return None

        aggregate = {
            dim: round(sum(scores) / len(scores), 1)
            for dim, scores in all_scores.items()
        }
        overall = round(sum(aggregate.values()) / len(aggregate), 1)

        return BookEvaluationSummary(
            book=book,
            total_chapters=len(eval_files),
            total_verses=total_verses,
            aggregate_scores=aggregate,
            overall_score=overall,
            flagged_verses=flagged_refs,
            evaluated_at=datetime.utcnow(),
        )

    def is_chapter_evaluated(self, book: str, chapter: int) -> bool:
        """Check if a chapter has been evaluated.

        Args:
            book: Book name.
            chapter: Chapter number.

        Returns:
            True if evaluation exists.
        """
        return self._get_eval_path(book, chapter).exists()

    def save_verse_evaluation(
        self, book: str, chapter: int, evaluation: VerseEvaluation
    ) -> Path:
        """Save a single verse evaluation.

        Loads existing chapter evaluation if it exists, updates the verse,
        and saves back.

        Args:
            book: Book name.
            chapter: Chapter number.
            evaluation: The verse evaluation to save.

        Returns:
            Path to the saved file.
        """
        # Load existing or create new chapter evaluation
        chapter_eval = self.load_chapter_evaluation(book, chapter)
        if chapter_eval is None:
            chapter_eval = ChapterEvaluation(
                book=book,
                chapter=chapter,
                evaluator_model=evaluation.evaluator_model,
                verses={},
            )

        # Update the verse
        chapter_eval.verses[evaluation.verse_number] = evaluation

        # Save back
        return self.save_chapter_evaluation(chapter_eval)