"""
benchmark.py
-------------
Runs three REAL systems against the labeled eval set in
data/eval_questions.json and reports REAL, measured metrics for each.
No numbers here are invented; run this script yourself to reproduce them
(`python scripts/evaluate.py`).

Systems compared:
  1. Basic Vector RAG -- single semantic-similarity retrieval pass, top-5,
     no reranking, no agents, no verification loop.
  2. Hybrid RAG        -- hybrid (BM25 + semantic, RRF-fused) retrieval +
     cross-encoder/MMR rerank, single pass, still no agentic loop.
  3. Agentic RAG        -- the full ResearchPilot pipeline: Planner,
     adaptive multi-cycle Retrieval, Evidence grading + contradiction
     detection, Reasoning, Verification, Final Answer.
"""
from __future__ import annotations

import json
import time
from collections import defaultdict
from typing import Dict, List

from src.agents.evidence import evaluate_evidence
from src.agents.graph import run_research
from src.agents.planner import plan
from src.agents.reasoning import reason
from src.agents.state import EvidenceItem, ResearchState
from src.agents.verification import verify
from src.evaluation.metrics import (
    citation_correctness,
    context_relevance,
    evidence_coverage,
    hallucination_rate_proxy,
    recall_at_k,
)
from src.rag.reranker import Reranker
from src.rag.retriever import HybridRetriever
from src.rag.vector_store import VectorStore


def _basic_vector_rag(store: VectorStore, retriever: HybridRetriever, question: str, top_k: int = 5) -> List[EvidenceItem]:
    passages = retriever.retrieve([question], strategy="semantic", top_k=top_k)
    return [
        EvidenceItem(
            chunk_id=p.chunk.chunk_id, doc_id=p.chunk.doc_id, title=p.chunk.title,
            authors=p.chunk.authors, year=p.chunk.year, url=p.chunk.url, text=p.chunk.text,
            retrieval_score=float(p.score), sub_question=question, strategy="semantic",
            relevance_label="SUPPORTED", relevance_score=float(p.score),
        )
        for p in passages
    ]


def _hybrid_rag(retriever: HybridRetriever, reranker: Reranker, question: str, top_k: int = 5) -> List[EvidenceItem]:
    passages = retriever.retrieve([question], strategy="hybrid", top_k=top_k * 2)
    passages = reranker.rerank(question, passages, top_k=top_k)
    return [
        EvidenceItem(
            chunk_id=p.chunk.chunk_id, doc_id=p.chunk.doc_id, title=p.chunk.title,
            authors=p.chunk.authors, year=p.chunk.year, url=p.chunk.url, text=p.chunk.text,
            retrieval_score=float(p.score), sub_question=question, strategy="hybrid",
            relevance_label="SUPPORTED", relevance_score=float(p.score),
        )
        for p in passages
    ]


def run_benchmark(corpus_path: str, eval_path: str) -> Dict[str, List[dict]]:
    store = VectorStore(corpus_path)
    retriever = HybridRetriever(store)
    reranker = Reranker()

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

    results: Dict[str, List[dict]] = defaultdict(list)

    for item in eval_questions:
        qid, question, gold = item["id"], item["question"], set(item["gold_doc_ids"])

        # --- Basic Vector RAG ---
        t0 = time.perf_counter()
        ev = _basic_vector_rag(store, retriever, question)
        latency = time.perf_counter() - t0
        retrieved_ids = [e["doc_id"] for e in ev]
        cited = set(retrieved_ids)
        results["Basic Vector RAG"].append({
            "question_id": qid,
            "recall_at_k": recall_at_k(retrieved_ids, gold),
            "context_relevance": context_relevance(ev),
            "evidence_coverage": evidence_coverage(ev, [question]),
            "citation_correctness": citation_correctness(cited, gold),
            "hallucination_rate_proxy": hallucination_rate_proxy(cited, gold),
            "retrieval_iterations": 1,
            "latency_seconds": round(latency, 4),
        })

        # --- Hybrid RAG ---
        t0 = time.perf_counter()
        ev = _hybrid_rag(retriever, reranker, question)
        latency = time.perf_counter() - t0
        retrieved_ids = [e["doc_id"] for e in ev]
        cited = set(retrieved_ids)
        results["Hybrid RAG"].append({
            "question_id": qid,
            "recall_at_k": recall_at_k(retrieved_ids, gold),
            "context_relevance": context_relevance(ev),
            "evidence_coverage": evidence_coverage(ev, [question]),
            "citation_correctness": citation_correctness(cited, gold),
            "hallucination_rate_proxy": hallucination_rate_proxy(cited, gold),
            "retrieval_iterations": 1,
            "latency_seconds": round(latency, 4),
        })

        # --- Full Agentic RAG ---
        t0 = time.perf_counter()
        state = run_research(question, max_iterations=3)
        latency = time.perf_counter() - t0
        used_evidence = [e for e in state.get("evidence", []) if e["relevance_label"] in ("SUPPORTED", "WEAKLY_SUPPORTED")]
        retrieved_ids = [e["doc_id"] for e in used_evidence]
        cited = set(retrieved_ids)
        results["Agentic RAG"].append({
            "question_id": qid,
            "recall_at_k": recall_at_k(retrieved_ids, gold),
            "context_relevance": context_relevance(used_evidence),
            "evidence_coverage": evidence_coverage(state.get("evidence", []), state.get("sub_questions", [question])),
            "citation_correctness": citation_correctness(cited, gold),
            "hallucination_rate_proxy": hallucination_rate_proxy(cited, gold),
            "retrieval_iterations": state.get("iteration", 0) + 1,
            "latency_seconds": round(latency, 4),
        })

    return results


def _ablation_run(question: str, disable: set) -> ResearchState:
    """Runs a hand-assembled pipeline with specific agentic components removed,
    reusing the exact same underlying agent functions as production -- this is
    NOT a separate fake implementation, just fewer nodes wired together.
    """
    state: ResearchState = {"question": question, "max_iterations": 1 if "adaptive_retrieval" in disable else 3,
                             "iteration": 0, "trace": []}

    if "planner" in disable:
        state["sub_questions"] = [question]
        state["question_type"] = "single_fact"
    else:
        state = plan(state)

    from src.agents.retrieval import retrieve
    state = retrieve(state)

    if "evidence_verification" in disable:
        for e in state["evidence"]:
            e["relevance_label"] = "SUPPORTED"
            e["relevance_score"] = 1.0
        state["contradictions"] = []
        state["sufficient"] = True
    else:
        if "contradiction_detection" in disable:
            state = evaluate_evidence(state)
            state["contradictions"] = []
        else:
            state = evaluate_evidence(state)

    state = reason(state)

    if "final_verification" in disable:
        state["verification_passed"] = True
    else:
        state = verify(state)
        if not state["verification_passed"] and "adaptive_retrieval" not in disable:
            state = retrieve(state)
            state = evaluate_evidence(state)
            state = reason(state)
            state = verify(state)

    from src.agents.final_answer import finalize
    state = finalize(state)
    return state


def run_ablation(corpus_path: str, eval_path: str) -> Dict[str, List[dict]]:
    from src.agents import retrieval as retrieval_module

    if retrieval_module._STORE is None:
        retrieval_module.init_rag(corpus_path)

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

    configs = {
        "Full Agentic RAG": set(),
        "Without Planner": {"planner"},
        "Without Evidence Verification": {"evidence_verification"},
        "Without Adaptive Retrieval": {"adaptive_retrieval"},
        "Without Contradiction Detection": {"contradiction_detection"},
        "Without Final Verification": {"final_verification"},
    }

    results: Dict[str, List[dict]] = defaultdict(list)
    for name, disable in configs.items():
        for item in eval_questions:
            qid, question, gold = item["id"], item["question"], set(item["gold_doc_ids"])
            t0 = time.perf_counter()
            state = _ablation_run(question, disable)
            latency = time.perf_counter() - t0
            used_evidence = [e for e in state.get("evidence", []) if e["relevance_label"] in ("SUPPORTED", "WEAKLY_SUPPORTED")]
            retrieved_ids = [e["doc_id"] for e in used_evidence]
            cited = set(retrieved_ids)
            results[name].append({
                "question_id": qid,
                "recall_at_k": recall_at_k(retrieved_ids, gold),
                "context_relevance": context_relevance(used_evidence),
                "evidence_coverage": evidence_coverage(state.get("evidence", []), state.get("sub_questions", [question])),
                "citation_correctness": citation_correctness(cited, gold),
                "hallucination_rate_proxy": hallucination_rate_proxy(cited, gold),
                "retrieval_iterations": state.get("iteration", 0) + 1,
                "latency_seconds": round(latency, 4),
            })
    return results


def summarize(results: Dict[str, List[dict]]) -> Dict[str, Dict[str, float]]:
    summary = {}
    metric_keys = ["recall_at_k", "context_relevance", "evidence_coverage",
                   "citation_correctness", "hallucination_rate_proxy",
                   "retrieval_iterations", "latency_seconds"]
    for system, rows in results.items():
        summary[system] = {
            k: round(sum(r[k] for r in rows) / len(rows), 4) for k in metric_keys
        }
    return summary