"""
Benchmark Dataset Management

Handles loading and managing benchmark datasets for
evaluating the SAGE RAG system.
"""

import json
import logging
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path

logger = logging.getLogger(__name__)


class QuestionCategory(str, Enum):
    """Categories of benchmark questions."""

    FACTUAL = "factual"  # Direct fact lookup
    REASONING = "reasoning"  # Requires inference
    CALCULATION = "calculation"  # Numerical/calculation
    REGULATORY = "regulatory"  # Regulatory interpretation
    COMPARATIVE = "comparative"  # Compare multiple sources
    PROCEDURAL = "procedural"  # Step-by-step procedures


class Difficulty(str, Enum):
    """Question difficulty levels."""

    EASY = "easy"  # Single source, direct answer
    MEDIUM = "medium"  # Multiple sources or inference
    HARD = "hard"  # Complex reasoning or edge cases


@dataclass
class BenchmarkQuestion:
    """
    A single benchmark question with expected answer.

    Contains all information needed to evaluate a RAG response
    against ground truth.
    """

    id: str
    question: str
    expected_answer: str
    expected_citations: list[str] = field(default_factory=list)

    # Classification
    category: QuestionCategory = QuestionCategory.FACTUAL
    difficulty: Difficulty = Difficulty.MEDIUM

    # Additional context
    key_points: list[str] = field(default_factory=list)
    source_documents: list[str] = field(default_factory=list)
    notes: str = ""

    # Metadata
    tags: list[str] = field(default_factory=list)
    created_by: str = ""

    def to_dict(self) -> dict:
        """Convert to dictionary."""
        return {
            "id": self.id,
            "question": self.question,
            "expected_answer": self.expected_answer,
            "expected_citations": self.expected_citations,
            "category": self.category.value,
            "difficulty": self.difficulty.value,
            "key_points": self.key_points,
            "source_documents": self.source_documents,
            "notes": self.notes,
            "tags": self.tags,
            "created_by": self.created_by,
        }

    @classmethod
    def from_dict(cls, data: dict) -> "BenchmarkQuestion":
        """Create from dictionary."""
        return cls(
            id=data["id"],
            question=data["question"],
            expected_answer=data["expected_answer"],
            expected_citations=data.get("expected_citations", []),
            category=QuestionCategory(data.get("category", "factual")),
            difficulty=Difficulty(data.get("difficulty", "medium")),
            key_points=data.get("key_points", []),
            source_documents=data.get("source_documents", []),
            notes=data.get("notes", ""),
            tags=data.get("tags", []),
            created_by=data.get("created_by", ""),
        )


@dataclass
class BenchmarkDataset:
    """
    A collection of benchmark questions.

    Represents a complete benchmark suite for evaluation.
    """

    name: str
    description: str = ""
    version: str = "1.0"
    questions: list[BenchmarkQuestion] = field(default_factory=list)

    # Metadata
    created_date: str = ""
    domain: str = "ncs"

    def __len__(self) -> int:
        return len(self.questions)

    def __iter__(self):
        return iter(self.questions)

    def filter_by_category(self, category: QuestionCategory) -> "BenchmarkDataset":
        """Get subset by category."""
        filtered = [q for q in self.questions if q.category == category]
        return BenchmarkDataset(
            name=f"{self.name}_{category.value}",
            description=f"Filtered by category: {category.value}",
            questions=filtered,
        )

    def filter_by_difficulty(self, difficulty: Difficulty) -> "BenchmarkDataset":
        """Get subset by difficulty."""
        filtered = [q for q in self.questions if q.difficulty == difficulty]
        return BenchmarkDataset(
            name=f"{self.name}_{difficulty.value}",
            description=f"Filtered by difficulty: {difficulty.value}",
            questions=filtered,
        )

    def filter_by_tags(self, tags: list[str]) -> "BenchmarkDataset":
        """Get subset by tags (any match)."""
        tag_set = set(tags)
        filtered = [q for q in self.questions if tag_set & set(q.tags)]
        return BenchmarkDataset(
            name=f"{self.name}_tagged",
            description=f"Filtered by tags: {tags}",
            questions=filtered,
        )

    def get_question(self, question_id: str) -> BenchmarkQuestion | None:
        """Get question by ID."""
        for q in self.questions:
            if q.id == question_id:
                return q
        return None

    def to_dict(self) -> dict:
        """Convert to dictionary."""
        return {
            "name": self.name,
            "description": self.description,
            "version": self.version,
            "domain": self.domain,
            "created_date": self.created_date,
            "questions": [q.to_dict() for q in self.questions],
        }

    @classmethod
    def from_dict(cls, data: dict) -> "BenchmarkDataset":
        """Create from dictionary."""
        return cls(
            name=data["name"],
            description=data.get("description", ""),
            version=data.get("version", "1.0"),
            domain=data.get("domain", "ncs"),
            created_date=data.get("created_date", ""),
            questions=[BenchmarkQuestion.from_dict(q) for q in data.get("questions", [])],
        )


class DatasetLoader:
    """
    Loads benchmark datasets from files.

    Supports JSON format for benchmark definitions.
    """

    def load_from_file(self, path: Path) -> BenchmarkDataset:
        """
        Load dataset from JSON file.

        Args:
            path: Path to JSON file

        Returns:
            BenchmarkDataset
        """
        logger.info(f"Loading benchmark dataset from {path}")

        with open(path) as f:
            data = json.load(f)

        return BenchmarkDataset.from_dict(data)

    def save_to_file(self, dataset: BenchmarkDataset, path: Path) -> None:
        """
        Save dataset to JSON file.

        Args:
            dataset: Dataset to save
            path: Output path
        """
        logger.info(f"Saving benchmark dataset to {path}")

        with open(path, "w") as f:
            json.dump(dataset.to_dict(), f, indent=2)


def create_sample_ncs_dataset() -> BenchmarkDataset:
    """
    Create a sample NCS benchmark dataset.

    Returns a small dataset with representative questions
    for testing the evaluation system.
    """
    questions = [
        BenchmarkQuestion(
            id="ncs-001",
            question="What is the single-parameter subcritical mass limit for U-235 in water?",
            expected_answer="The single-parameter subcritical mass limit for U-235 is 760 g when reflected by an unlimited water reflector, as specified in ANSI/ANS-8.1.",
            expected_citations=["ANSI/ANS-8.1"],
            category=QuestionCategory.FACTUAL,
            difficulty=Difficulty.EASY,
            key_points=["760 g", "water reflector", "single-parameter"],
            source_documents=["ANSI/ANS-8.1-2014"],
            tags=["mass-limit", "u235", "standards"],
        ),
        BenchmarkQuestion(
            id="ncs-002",
            question="What are the double contingency principle requirements for criticality safety?",
            expected_answer="The double contingency principle requires that process designs shall incorporate sufficient factors of safety to require at least two unlikely, independent, and concurrent changes in process conditions before a criticality accident is possible.",
            expected_citations=["ANSI/ANS-8.1", "10 CFR 70.24"],
            category=QuestionCategory.REGULATORY,
            difficulty=Difficulty.MEDIUM,
            key_points=["two unlikely changes", "independent", "concurrent", "factors of safety"],
            source_documents=["ANSI/ANS-8.1-2014", "10 CFR 70"],
            tags=["double-contingency", "regulations", "standards"],
        ),
        BenchmarkQuestion(
            id="ncs-003",
            question="How does the critical mass of Pu-239 compare to U-235?",
            expected_answer="Pu-239 has a lower critical mass than U-235. The subcritical mass limit for Pu-239 is 450 g compared to 760 g for U-235 when water reflected, reflecting Pu-239's higher fission cross-section.",
            expected_citations=["ANSI/ANS-8.1", "ANSI/ANS-8.15"],
            category=QuestionCategory.COMPARATIVE,
            difficulty=Difficulty.MEDIUM,
            key_points=["450 g", "760 g", "lower critical mass", "fission cross-section"],
            source_documents=["ANSI/ANS-8.1-2014", "ANSI/ANS-8.15"],
            tags=["mass-limit", "pu239", "comparison"],
        ),
        BenchmarkQuestion(
            id="ncs-004",
            question="What validation requirements apply to criticality safety calculations?",
            expected_answer="Criticality calculations must be validated against applicable benchmark experiments. The validation should establish bias and bias uncertainty, and calculations should use an Upper Subcritical Limit (USL) that accounts for these factors. ICSBEP benchmarks are commonly used for validation.",
            expected_citations=["ANSI/ANS-8.24", "NUREG/CR-6698"],
            category=QuestionCategory.PROCEDURAL,
            difficulty=Difficulty.HARD,
            key_points=["benchmark experiments", "bias", "bias uncertainty", "USL", "ICSBEP"],
            source_documents=["ANSI/ANS-8.24", "NUREG/CR-6698"],
            tags=["validation", "calculations", "benchmarks"],
        ),
        BenchmarkQuestion(
            id="ncs-005",
            question="What is the minimum critical concentration of U-235 in aqueous solution?",
            expected_answer="The minimum critical concentration (single-parameter limit) for U-235 in uniform aqueous solution is 11.6 g U-235/L when optimally moderated and reflected by water.",
            expected_citations=["ANSI/ANS-8.1"],
            category=QuestionCategory.FACTUAL,
            difficulty=Difficulty.EASY,
            key_points=["11.6 g/L", "aqueous solution", "uniform", "optimally moderated"],
            source_documents=["ANSI/ANS-8.1-2014"],
            tags=["concentration-limit", "u235", "solution"],
        ),
    ]

    return BenchmarkDataset(
        name="NCS-Sample-v1",
        description="Sample NCS benchmark questions for RAG evaluation",
        version="1.0",
        domain="ncs",
        questions=questions,
    )


def load_or_create_sample_dataset(path: Path | None = None) -> BenchmarkDataset:
    """
    Load dataset from path or create sample if not found.

    Args:
        path: Optional path to dataset file

    Returns:
        BenchmarkDataset
    """
    if path and path.exists():
        loader = DatasetLoader()
        return loader.load_from_file(path)

    return create_sample_ncs_dataset()
