"""
External benchmark loaders (optional dependencies).

Supports:
- HotpotQA (multi-hop reasoning)
- Natural Questions (single-hop baseline)
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Optional, Iterator
import warnings


@dataclass
class BenchmarkExample:
    """A benchmark example from external dataset."""
    query: str
    supporting_facts: list[str]  # Document titles that contain the answer
    answer: str
    context_docs: list[tuple[str, str]]  # (title, content) pairs
    question_type: str = "unknown"


def load_hotpotqa(
    split: str = "validation",
    num_samples: Optional[int] = 100,
    question_type: Optional[str] = None,
) -> Iterator[BenchmarkExample]:
    """
    Load HotpotQA dataset (requires 'datasets' library).

    Args:
        split: Dataset split ('train', 'validation')
        num_samples: Max samples to load (None for all)
        question_type: Filter by type ('bridge', 'comparison', None for all)

    Yields:
        BenchmarkExample objects

    Example:
        >>> for ex in load_hotpotqa(num_samples=10):
        ...     print(ex.query, ex.answer)
    """
    try:
        from datasets import load_dataset
    except ImportError:
        raise ImportError(
            "HotpotQA requires the 'datasets' library. "
            "Install with: pip install datasets"
        )

    # Load distractor setting (more challenging)
    dataset = load_dataset("hotpot_qa", "distractor", split=split)  # nosec B615

    count = 0
    for item in dataset:
        # Filter by question type if specified
        if question_type and item.get("type") != question_type:
            continue

        # Get supporting facts (titles of documents with the answer)
        supporting_titles = list(set(
            item.get("supporting_facts", {}).get("title", [])
        ))

        # Get context documents
        context = item.get("context", {})
        titles = context.get("title", [])
        sentences = context.get("sentences", [])

        docs = []
        for title, sents in zip(titles, sentences):
            content = " ".join(sents)
            docs.append((title, content))

        yield BenchmarkExample(
            query=item["question"],
            supporting_facts=supporting_titles,
            answer=item["answer"],
            context_docs=docs,
            question_type=item.get("type", "unknown"),
        )

        count += 1
        if num_samples and count >= num_samples:
            break


def load_natural_questions(
    split: str = "validation",
    num_samples: Optional[int] = 100,
) -> Iterator[BenchmarkExample]:
    """
    Load Natural Questions dataset (requires 'datasets' library).

    This is a single-hop baseline - graph shouldn't help much here.

    Args:
        split: Dataset split
        num_samples: Max samples to load

    Yields:
        BenchmarkExample objects
    """
    try:
        from datasets import load_dataset
    except ImportError:
        raise ImportError(
            "Natural Questions requires the 'datasets' library. "
            "Install with: pip install datasets"
        )

    dataset = load_dataset("natural_questions", "default", split=split)  # nosec B615

    count = 0
    for item in dataset:
        # NQ has complex structure, simplify
        question = item.get("question", {}).get("text", "")
        if not question:
            continue

        # Get short answer if available
        annotations = item.get("annotations", {})
        short_answers = annotations.get("short_answers", [])

        if short_answers and short_answers[0]:
            answer = short_answers[0].get("text", "")
        else:
            # Fall back to yes/no answer
            answer = annotations.get("yes_no_answer", "unknown")

        # Document is the Wikipedia page
        doc_text = item.get("document", {}).get("text", "")[:2000]  # Truncate
        doc_title = item.get("document", {}).get("title", "wiki_doc")

        yield BenchmarkExample(
            query=question,
            supporting_facts=[doc_title],
            answer=str(answer),
            context_docs=[(doc_title, doc_text)],
            question_type="single_hop",
        )

        count += 1
        if num_samples and count >= num_samples:
            break


class BenchmarkLoader:
    """
    Unified interface for loading external benchmarks.

    Example:
        >>> loader = BenchmarkLoader()
        >>> examples = loader.load("hotpotqa", num_samples=50)
        >>> for ex in examples:
        ...     print(ex.query)
    """

    BENCHMARKS = {
        "hotpotqa": load_hotpotqa,
        "natural_questions": load_natural_questions,
    }

    def __init__(self):
        self._check_dependencies()

    def _check_dependencies(self):
        """Check if datasets library is available."""
        try:
            import datasets  # noqa: F401
            self.datasets_available = True
        except ImportError:
            self.datasets_available = False
            warnings.warn(
                "External benchmarks require 'datasets' library. "
                "Install with: pip install datasets"
            )

    def list_benchmarks(self) -> list[str]:
        """List available benchmarks."""
        return list(self.BENCHMARKS.keys())

    def load(
        self,
        name: str,
        num_samples: Optional[int] = 100,
        **kwargs,
    ) -> list[BenchmarkExample]:
        """
        Load a benchmark by name.

        Args:
            name: Benchmark name ('hotpotqa', 'natural_questions')
            num_samples: Max samples to load
            **kwargs: Additional arguments for the loader

        Returns:
            List of BenchmarkExample objects
        """
        if name not in self.BENCHMARKS:
            available = ", ".join(self.BENCHMARKS.keys())
            raise ValueError(f"Unknown benchmark: {name}. Available: {available}")

        loader = self.BENCHMARKS[name]
        return list(loader(num_samples=num_samples, **kwargs))


def benchmark_to_evaluation_examples(
    examples: list[BenchmarkExample],
    memory,
) -> tuple[list, dict[str, str]]:
    """
    Convert benchmark examples to EvaluationExample format.

    This inserts documents into memory and creates evaluation examples
    that reference the inserted episode IDs.

    Args:
        examples: Benchmark examples to convert
        memory: Memory instance to insert documents into

    Returns:
        Tuple of (evaluation_examples, id_mapping)
    """
    from hieromem_core.evaluation import create_example

    id_mapping = {}  # title -> episode_id
    eval_examples = []

    # Insert all documents first
    for ex in examples:
        for title, content in ex.context_docs:
            if title in id_mapping:
                continue  # Already inserted

            ep = memory.add_text_episode(content, doc_id=title)
            id_mapping[title] = ep.id

    # Create evaluation examples
    for ex in examples:
        relevant_ids = [
            id_mapping[title]
            for title in ex.supporting_facts
            if title in id_mapping
        ]

        if relevant_ids:  # Only if we have relevant docs
            eval_ex = create_example(
                query=ex.query,
                relevant_ids=relevant_ids,
                answer=ex.answer,
            )
            eval_examples.append(eval_ex)

    return eval_examples, id_mapping
