"""Evaluation runner: load corpus, run queries, sweep parameters.

Usage:
    # Run once with current config
    python -m domesday.eval.runner

    # Use a custom config (e.g. local embedder for offline eval)
    python -m domesday.eval.runner --config src/domesday/eval/eval_config.toml

    # Sweep parameters
    python -m domesday.eval.runner --sweep [--quick]

    # Interactive mode (inspect individual results)
    python -m domesday.eval.runner --interactive
"""

from __future__ import annotations

import itertools
import json
import logging
from collections.abc import Sequence
from dataclasses import dataclass
from pathlib import Path

from domesday import chunking, config
from domesday.core import models, pipeline
from domesday.eval import corpus, llm_judge
from domesday.eval import models as eval_models

logger = logging.getLogger(__name__)


@dataclass
class EvalRunner:
    """Runs evaluation queries against a loaded pipeline and scores results."""

    pipeline: pipeline.Pipeline
    snippet_id_map: dict[str, str]  # id_prefix → full snippet UUID

    @classmethod
    def create(
        cls,
        config_path: Path | None = None,
        snippets: Sequence[corpus.TestSnippet] | None = None,
    ) -> EvalRunner:
        """Build pipeline and ingest test corpus."""

        pl = config.build_pipeline(config_path)
        snippets = snippets or corpus.SNIPPETS
        id_map: dict[str, str] = {}

        # Check if already loaded (idempotent for repeated runs)
        existing_count = pl.snippet_store.count()
        if existing_count >= len(snippets):
            logger.info(
                "Corpus appears loaded (%d snippets), resolving IDs...", existing_count
            )
            all_active = pl.snippet_store.get_all_active()
            # Rebuild map from source_file field which stores id_prefix
            for s in all_active:
                if s.source_file and s.source_file.startswith("test:"):
                    prefix = s.source_file.removeprefix("test:")
                    id_map[prefix] = s.id
        else:
            logger.info("Loading %d test snippets...", len(snippets))
            for ts in snippets:
                snippet = pl.add_snippet(
                    ts.text,
                    project=pl.default_project,
                    author=ts.author,
                    tags=ts.tags,
                    source_file=f"test:{ts.id_prefix}",
                )
                id_map[ts.id_prefix] = snippet.id

        logger.info("Ready: %d snippets mapped", len(id_map))
        return cls(pipeline=pl, snippet_id_map=id_map)

    def _resolve_prefix(self, results: list[models.SearchResult]) -> list[str]:
        """Map retrieved snippet UUIDs back to test id_prefixes."""
        uuid_to_prefix = {v: k for k, v in self.snippet_id_map.items()}
        return [uuid_to_prefix.get(r.snippet.id, r.snippet.id[:8]) for r in results]

    def run_query(
        self,
        eq: corpus.EvalQuery,
        *,
        k: int = 10,
        min_score: float = 0.0,
    ) -> eval_models.QueryResult:
        """Evaluate a single query against ground truth."""
        results = self.pipeline.search(eq.query, k=k, min_score=min_score)

        retrieved_prefixes = self._resolve_prefix(results)
        scores = [r.score for r in results]

        return eval_models.QueryResult(
            query=eq.query,
            category=eq.category,
            relevant_ids=eq.relevant_ids,
            retrieved_ids=retrieved_prefixes,
            retrieved_scores=scores,
            irrelevant_ids=eq.irrelevant_ids,
        )

    def run_all(
        self,
        queries: Sequence[corpus.EvalQuery] | None = None,
        *,
        k: int = 10,
        min_score: float = 0.0,
        run_name: str = "default",
        params: dict[str, object] | None = None,
        judge: bool = False,
    ) -> eval_models.EvalRun:
        """Run all eval queries and return aggregated metrics.

        Args:
            judge: If True, also run the full RAG pipeline and have an LLM
                   judge score each generated answer. Slower and costs API
                   calls, but measures generation quality — not just retrieval.
        """
        queries = queries or corpus.EVAL_QUERIES
        eval_run = eval_models.EvalRun(run_name=run_name, params=params or {})

        for eq in queries:
            result = self.run_query(eq, k=k, min_score=min_score)
            eval_run.results.append(result)

        if judge:
            judge_instance = llm_judge.LLMJudge()
            logger.info("Running LLM judge on %d generated answers...", len(queries))

            # Run full RAG for each query and judge the response
            items: list[tuple[str, models.RAGResponse, bool]] = []
            for eq in queries:
                rag_response = self.pipeline.ask(eq.query, k=k, min_score=min_score)
                is_negative = eq.category == "negative"
                items.append((eq.query, rag_response, is_negative))

            judgments = judge_instance.judge_batch(items)

            # Attach judgments to the eval run
            eval_run.judgments = judgments  # type: ignore[attr-defined]

            # Print judgment summary
            avg_composite = sum(j.composite for j in judgments) / len(judgments)
            avg_faith = sum(j.faithfulness for j in judgments) / len(judgments)
            avg_abstain = sum(j.abstention for j in judgments) / len(judgments)
            logger.info(
                "Judge scores — composite=%.3f faithfulness=%.3f abstention=%.3f",
                avg_composite,
                avg_faith,
                avg_abstain,
            )

        return eval_run


# ---------------------------------------------------------------------------
# Parameter sweep
# ---------------------------------------------------------------------------


@dataclass(frozen=True, slots=True)
class SweepConfig:
    """Defines the parameter grid to sweep."""

    min_scores: list[float]
    k_values: list[int]
    chunk_max_tokens: list[int]
    chunk_overlaps: list[int]

    @classmethod
    def default(cls) -> SweepConfig:
        return cls(
            min_scores=[0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.5],
            k_values=[5, 10, 15],
            chunk_max_tokens=[200, 400, 600],
            chunk_overlaps=[0, 50, 100],
        )

    @classmethod
    def quick(cls) -> SweepConfig:
        """Smaller grid for fast iteration."""
        return cls(
            min_scores=[0.2, 0.3, 0.4],
            k_values=[5, 10],
            chunk_max_tokens=[400],
            chunk_overlaps=[50],
        )


def run_sweep(
    sweep_config: SweepConfig | None = None,
    config_path: Path | None = None,
    output_path: Path | None = None,
) -> list[eval_models.EvalRun]:
    """Sweep parameter combinations and report results.

    Note: chunk_max_tokens and chunk_overlaps require re-ingesting the corpus
    (different chunking → different embeddings). min_score and k are cheap to
    sweep because they only affect the retrieval filter.
    """
    sweep = sweep_config or SweepConfig.quick()
    all_runs: list[eval_models.EvalRun] = []

    for max_tokens, overlap in itertools.product(
        sweep.chunk_max_tokens, sweep.chunk_overlaps
    ):
        if overlap >= max_tokens:
            continue  # nonsensical combination

        logger.info("=== Chunk config: max=%d overlap=%d ===", max_tokens, overlap)

        runner = EvalRunner.create(config_path, corpus.SNIPPETS)
        runner.pipeline.chunker = chunking.SimpleChunker(
            max_tokens=max_tokens, overlap_tokens=overlap
        )

        for k_val, min_score in itertools.product(sweep.k_values, sweep.min_scores):
            run_name = f"chunk{max_tokens}_overlap{overlap}_k{k_val}_min{min_score}"
            params: dict[str, object] = {
                "chunk_max_tokens": max_tokens,
                "chunk_overlap": overlap,
                "k": k_val,
                "min_score": min_score,
            }

            eval_run = runner.run_all(
                k=k_val,
                min_score=min_score,
                run_name=run_name,
                params=params,
            )
            all_runs.append(eval_run)
            logger.info(
                "%s → P=%.3f R=%.3f MRR=%.3f Neg=%.3f",
                run_name,
                eval_run.mean_precision,
                eval_run.mean_recall,
                eval_run.mean_mrr,
                eval_run.negative_accuracy,
            )

    # Sort by composite score (weighted recall + negative accuracy)
    all_runs.sort(
        key=lambda r: (
            r.mean_recall * 0.5 + r.negative_accuracy * 0.3 + r.mean_mrr * 0.2
        ),
        reverse=True,
    )

    # Save results
    if output_path is None:
        output_path = Path("./data/eval_results.json")
    output_path.parent.mkdir(parents=True, exist_ok=True)

    serializable = [
        {
            "run_name": run.run_name,
            "params": run.params,
            "summary": run.summary(),
            "by_category": run.summary_by_category(),
        }
        for run in all_runs
    ]
    output_path.write_text(json.dumps(serializable, indent=2, default=str))
    logger.info("Results saved to %s", output_path)

    return all_runs


# ---------------------------------------------------------------------------
# Interactive mode
# ---------------------------------------------------------------------------


def interactive(config_path: Path | None = None) -> None:
    """Interactive REPL for inspecting retrieval results."""
    runner = EvalRunner.create(config_path)

    print("\n domesday eval — interactive mode")
    print("Commands: query <text>, eval <index>, judge <index>, sweep, report, quit\n")

    while True:
        try:
            raw = input("domesday> ").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            break

        if not raw:
            continue

        if raw.lower() in ("quit", "exit", "q"):
            break

        if raw.lower() == "report":
            eval_run = runner.run_all(run_name="interactive-full")
            print(eval_run.format_report())
            continue

        if raw.lower().startswith("eval"):
            parts = raw.split(maxsplit=1)
            idx = int(parts[1]) if len(parts) > 1 else 0
            if idx >= len(corpus.EVAL_QUERIES):
                print(f"Index out of range (0–{len(corpus.EVAL_QUERIES) - 1})")
                continue
            eq = corpus.EVAL_QUERIES[idx]
            result = runner.run_query(eq, k=10)
            print(f"\n  Query:    {eq.query}")
            print(f"  Category: {eq.category}")
            print(f"  Expected: {eq.relevant_ids}")
            print(f"  Got:      {result.retrieved_ids}")
            print(f"  Scores:   {[f'{s:.3f}' for s in result.retrieved_scores]}")
            print(
                f"  P={result.precision_at_k:.3f} R={result.recall:.3f} MRR={result.mrr:.3f}"
            )
            if result.has_irrelevant_intrusion:
                print("  ⚠ IRRELEVANT INTRUSION")
            print()
            continue

        if raw.lower().startswith("judge"):
            parts = raw.split(maxsplit=1)
            idx = int(parts[1]) if len(parts) > 1 else 0
            if idx >= len(corpus.EVAL_QUERIES):
                print(f"Index out of range (0–{len(corpus.EVAL_QUERIES) - 1})")
                continue
            eq = corpus.EVAL_QUERIES[idx]
            print(f"\n  Generating answer for: {eq.query}")
            rag_response = runner.pipeline.ask(eq.query, k=10)
            print(f"\n  Answer:\n  {rag_response.answer[:300]}...")
            print("\n  Judging...")
            judge = llm_judge.LLMJudge()
            scores = judge.judge(
                eq.query, rag_response, is_negative=(eq.category == "negative")
            )
            print(f"  Faithfulness:      {scores.faithfulness:.2f}")
            print(f"  Citation accuracy: {scores.citation_accuracy:.2f}")
            print(f"  Abstention:        {scores.abstention:.2f}")
            print(f"  Synthesis:         {scores.synthesis:.2f}")
            print(f"  Relevance:         {scores.relevance:.2f}")
            print(f"  Composite:         {scores.composite:.2f}")
            print(f"  Explanation:       {scores.explanation}")
            print()
            continue

        # Default: treat as a free query
        results = runner.pipeline.search(raw, k=5)
        if not results:
            print("  (no results above threshold)\n")
        else:
            for i, r in enumerate(results, 1):
                prefix = "?"
                uuid_to_prefix = {v: k for k, v in runner.snippet_id_map.items()}
                prefix = uuid_to_prefix.get(r.snippet.id, r.snippet.id[:8])
                print(
                    f"  {i}. [{prefix}] score={r.score:.3f} "
                    f"({r.snippet.author}) {r.snippet.raw_text[:80]}…"
                )
            print()


# ---------------------------------------------------------------------------
# CLI entry point
# ---------------------------------------------------------------------------


def main() -> None:
    import argparse

    parser = argparse.ArgumentParser(description="domesday evaluation")
    parser.add_argument("--sweep", action="store_true", help="Run parameter sweep")
    parser.add_argument(
        "--interactive", "-i", action="store_true", help="Interactive mode"
    )
    parser.add_argument(
        "--judge",
        "-j",
        action="store_true",
        help="Also judge generation quality with LLM",
    )
    parser.add_argument("--config", type=Path, default=None)
    parser.add_argument("--quick", action="store_true", help="Use quick sweep grid")
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")

    if args.interactive:
        interactive(args.config)
    elif args.sweep:
        sweep_cfg = SweepConfig.quick() if args.quick else SweepConfig.default()
        runs = run_sweep(sweep_cfg, args.config)
        print("\nTop 5 configurations:")
        for run in runs[:5]:
            print(run.format_report())
    else:
        # Default: single run with current config, print report
        runner = EvalRunner.create(args.config)
        eval_run = runner.run_all(run_name="default", judge=args.judge)
        print(eval_run.format_report())


if __name__ == "__main__":
    main()
