"""
Evaluation harness for the DSA RAG chatbot.

Runs the predefined question set (evaluation/questions.json) through the
agent router + retriever (and optionally the full LLM pipeline) and reports:

    - routing accuracy        (predicted route_type vs expected_route)
    - retrieval topic recall  (expected topics found among retrieved chunks)
    - retrieval latency       (avg / p95, ms)
    - generation latency      (avg / p95, ms) -- only with --with-llm

Usage:
    python -m evaluation.evaluate
    python -m evaluation.evaluate --with-llm
    python -m evaluation.evaluate --questions evaluation/questions.json

Run from the project root (dsa-rag-chatbot/) so `config` and the app
packages resolve correctly.
"""

import argparse
import json
import os
import statistics
import time

import config
from agents.followup import gather_followup_context
from agents.router import RouteType, classify, gather_context
from logs.logger import get_logger

logger = get_logger(__name__)

DEFAULT_QUESTIONS_PATH = os.path.join(os.path.dirname(__file__), "questions.json")

# A tiny bit of fake prior conversation so the follow-up route can actually
# be exercised (is_followup_query requires has_conversation_history=True).
_FAKE_HISTORY_FOR_FOLLOWUPS = [
    {"role": "user", "content": "Explain merge sort"},
    {"role": "assistant", "content": "Merge sort is a divide-and-conquer sorting algorithm..."},
]


def _load_questions(path: str) -> list:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def _percentile(values: list, pct: float) -> float:
    if not values:
        return 0.0
    values = sorted(values)
    k = (len(values) - 1) * (pct / 100)
    f, c = int(k), min(int(k) + 1, len(values) - 1)
    if f == c:
        return values[f]
    return values[f] + (k - f) * (values[c] - values[f])


def _retrieved_topics(chunks: list) -> set:
    return {c.get("metadata", {}).get("topic") for c in chunks if c.get("metadata", {}).get("topic")}


def _topic_recall(expected_topics: list, retrieved: set) -> float:
    if not expected_topics:
        return 1.0  # nothing expected (e.g. out_of_scope) => trivially satisfied
    hits = sum(1 for t in expected_topics if t in retrieved)
    return hits / len(expected_topics)


def evaluate_one(item: dict, retriever, with_llm: bool) -> dict:
    query = item["query"]
    expected_route = item["expected_route"]
    expected_topics = item.get("expected_topics", [])
    recent_messages = _FAKE_HISTORY_FOR_FOLLOWUPS if item.get("requires_history") else []

    t0 = time.perf_counter()
    decision = classify(query, has_conversation_history=len(recent_messages) > 0)
    decision = gather_context(decision, retriever, recent_messages=recent_messages)
    retrieval_latency_ms = (time.perf_counter() - t0) * 1000

    if decision.route_type == RouteType.SINGLE_TOPIC:
        retrieved = _retrieved_topics(decision.single_chunks)
    elif decision.route_type == RouteType.COMPARISON:
        retrieved = set()
        for chunks in decision.comparison_context.values():
            retrieved |= _retrieved_topics(chunks)
    elif decision.route_type == RouteType.FOLLOWUP:
        retrieved = _retrieved_topics(decision.followup_context.get("additional_chunks", []))
    else:
        retrieved = set()

    route_correct = decision.route_type.value == expected_route
    topic_recall = _topic_recall(expected_topics, retrieved)

    result = {
        "id": item.get("id"),
        "query": query,
        "expected_route": expected_route,
        "predicted_route": decision.route_type.value,
        "route_correct": route_correct,
        "expected_topics": expected_topics,
        "retrieved_topics": sorted(t for t in retrieved if t),
        "topic_recall": round(topic_recall, 3),
        "retrieval_latency_ms": round(retrieval_latency_ms, 2),
    }

    if with_llm:
        from llm.generate import generate
        from llm.prompts import (
            build_comparison_prompt,
            build_followup_prompt,
            build_reject_prompt,
            build_single_prompt,
        )

        if decision.route_type == RouteType.SINGLE_TOPIC:
            prompt = build_single_prompt(decision.single_chunks, [], recent_messages, query)
        elif decision.route_type == RouteType.COMPARISON:
            prompt = build_comparison_prompt(decision.comparison_context, [], recent_messages, query)
        elif decision.route_type == RouteType.FOLLOWUP:
            prompt = build_followup_prompt(decision.followup_context, [], recent_messages, query)
        else:
            prompt = build_reject_prompt(query)

        t1 = time.perf_counter()
        try:
            response_text = generate(prompt)
            gen_error = None
        except Exception as exc:  # keep the run going even if the LLM call fails
            response_text = None
            gen_error = str(exc)
        generation_latency_ms = (time.perf_counter() - t1) * 1000

        result["generation_latency_ms"] = round(generation_latency_ms, 2)
        result["response_preview"] = (response_text or "")[:200]
        if gen_error:
            result["generation_error"] = gen_error

    return result


def run(questions_path: str = DEFAULT_QUESTIONS_PATH, with_llm: bool = False) -> dict:
    from rag.retriever import get_retriever

    questions = _load_questions(questions_path)
    retriever = get_retriever()

    results = [evaluate_one(item, retriever, with_llm) for item in questions]

    route_accuracy = sum(r["route_correct"] for r in results) / len(results)
    avg_topic_recall = statistics.mean(r["topic_recall"] for r in results)
    retrieval_latencies = [r["retrieval_latency_ms"] for r in results]

    summary = {
        "total_questions": len(results),
        "route_accuracy": round(route_accuracy, 3),
        "avg_topic_recall": round(avg_topic_recall, 3),
        "retrieval_latency_ms_avg": round(statistics.mean(retrieval_latencies), 2),
        "retrieval_latency_ms_p95": round(_percentile(retrieval_latencies, 95), 2),
    }

    if with_llm:
        gen_latencies = [r["generation_latency_ms"] for r in results if "generation_latency_ms" in r]
        if gen_latencies:
            summary["generation_latency_ms_avg"] = round(statistics.mean(gen_latencies), 2)
            summary["generation_latency_ms_p95"] = round(_percentile(gen_latencies, 95), 2)
        summary["generation_errors"] = sum(1 for r in results if r.get("generation_error"))

    report = {"summary": summary, "results": results}

    logger.info("Evaluation complete: %s", summary)
    return report


def _print_report(report: dict) -> None:
    summary = report["summary"]
    print("\n=== DSA RAG Chatbot — Evaluation Summary ===")
    for key, value in summary.items():
        print(f"  {key}: {value}")

    print("\n=== Per-question results ===")
    header = f"{'id':<5} {'route (exp->got)':<28} {'recall':<8} {'ret_ms':<8} query"
    print(header)
    print("-" * len(header))
    for r in report["results"]:
        route_str = f"{r['expected_route']} -> {r['predicted_route']}"
        mark = "OK" if r["route_correct"] else "MISS"
        print(
            f"{r['id']:<5} {route_str:<28} {r['topic_recall']:<8} "
            f"{r['retrieval_latency_ms']:<8} [{mark}] {r['query']}"
        )


def main():
    parser = argparse.ArgumentParser(description="Evaluate DSA RAG chatbot retrieval + routing.")
    parser.add_argument(
        "--questions", default=DEFAULT_QUESTIONS_PATH, help="Path to questions.json"
    )
    parser.add_argument(
        "--with-llm",
        action="store_true",
        help="Also call the configured LLM provider end-to-end (uses API quota).",
    )
    parser.add_argument(
        "--out",
        default=None,
        help="Path to write the full JSON report (default: evaluation/last_report.json)",
    )
    args = parser.parse_args()

    report = run(questions_path=args.questions, with_llm=args.with_llm)
    _print_report(report)

    out_path = args.out or os.path.join(os.path.dirname(__file__), "last_report.json")
    with open(out_path, "w", encoding="utf-8") as f:
        json.dump(report, f, indent=2)
    print(f"\nFull report written to: {out_path}")


if __name__ == "__main__":
    main()