"""Evaluate the copilot on the golden set (lesson 10), optionally as a Langfuse experiment (lesson 18).

    python -m copilot.evaluate              # local report: retrieval + answer metrics
    python -m copilot.evaluate --langfuse   # also upload the dataset and run an experiment in Langfuse
                                            # (needs `pip install langfuse` and LANGFUSE_* variables)
"""

from __future__ import annotations

import json
import sys
from dataclasses import dataclass
from pathlib import Path
from statistics import mean

GOLDEN_PATH = Path(__file__).resolve().parents[1] / "eval" / "golden.jsonl"


def load_golden(path: Path = GOLDEN_PATH) -> list[dict]:
    return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]


def ranked_lessons(copilot, question: str, k: int = 10) -> list[str]:
    seen: list[str] = []
    for hit in copilot.index.search(question, k=40):
        lesson = hit.chunk.url.split("#")[0]
        if lesson not in seen:
            seen.append(lesson)
        if len(seen) == k:
            break
    return seen


# ---------------------------------------------------------------- evaluators, in Langfuse's signature
@dataclass
class Evaluation:                                   # same fields as langfuse.Evaluation
    name: str
    value: float
    comment: str | None = None


def hit_at_5(*, input, output, expected_output, **kwargs) -> Evaluation:
    hit = bool(set(expected_output) & set(output["lessons"][:5]))
    return Evaluation("hit@5", float(hit), None if hit else f"top: {output['lessons'][:1]}")


def reciprocal_rank(*, input, output, expected_output, **kwargs) -> Evaluation:
    rank = next((i for i, url in enumerate(output["lessons"], 1) if url in expected_output), None)
    return Evaluation("reciprocal_rank", 1 / rank if rank else 0.0)


def grounded(*, input, output, expected_output, **kwargs) -> Evaluation:
    return Evaluation("grounded", float(output["grounded"]))


def cites_expected_lesson(*, input, output, expected_output, **kwargs) -> Evaluation:
    cited = {url.split("#")[0] for url in output["sources"]}
    return Evaluation("cites_expected", float(bool(cited & set(expected_output))))


EVALUATORS = [hit_at_5, reciprocal_rank, grounded, cites_expected_lesson]


def make_task(copilot):
    def task(*, item, **kwargs) -> dict:
        question = item["input"] if isinstance(item, dict) else item.input
        answer = copilot.ask(question, user_id="eval")
        return {"answer": answer.text, "grounded": answer.grounded, "sources": [c.url for c in answer.sources],
                "lessons": ranked_lessons(copilot, question), "cost_usd": answer.cost_usd,
                "latency_s": answer.latency_s}
    return task


def run_local(copilot, golden: list[dict] | None = None) -> dict:
    golden = golden or load_golden()
    task = make_task(copilot)
    rows = []
    for g in golden:
        item = {"input": g["question"], "expected_output": g["lessons"]}
        output = task(item=item)
        scores = {e.name: e for e in (ev(input=item["input"], output=output, expected_output=g["lessons"])
                                      for ev in EVALUATORS)}
        rows.append({"question": g["question"], "output": output, "scores": scores})
    summary = {name: mean(r["scores"][name].value for r in rows) for name in rows[0]["scores"]}
    summary["items"] = len(rows)
    summary["avg_latency_ms"] = 1000 * mean(r["output"]["latency_s"] for r in rows)
    summary["total_cost_usd"] = sum(r["output"]["cost_usd"] for r in rows)
    summary["misses"] = [r["question"] for r in rows if r["scores"]["hit@5"].value == 0]
    return summary


def run_langfuse(copilot, dataset_name: str = "docs-copilot-golden") -> None:
    """Upload the golden set and run an experiment in Langfuse (lesson 18). Not exercised offline."""
    from langfuse import Evaluation as LfEvaluation, get_client
    langfuse = get_client()
    langfuse.create_dataset(name=dataset_name)
    for g in load_golden():
        langfuse.create_dataset_item(dataset_name=dataset_name, input=g["question"], expected_output=g["lessons"])

    def wrap(ev):
        def lf_evaluator(**kwargs):
            e = ev(**kwargs)
            return LfEvaluation(name=e.name, value=e.value, comment=e.comment)
        return lf_evaluator

    result = langfuse.get_dataset(dataset_name).run_experiment(
        name=f"{copilot.model.model}-{copilot.index.embedder.name}", task=make_task(copilot),
        evaluators=[wrap(ev) for ev in EVALUATORS], max_concurrency=4)
    print(result.format())
    langfuse.flush()


def main(argv: list[str]) -> int:
    from copilot import build
    copilot = build()
    s = run_local(copilot)
    print(f"{s['items']} questions · model={copilot.model.model} · embedder={copilot.index.embedder.name}")
    for name in ("hit@5", "reciprocal_rank", "grounded", "cites_expected"):
        print(f"  {name:16} {s[name]:.2f}")
    print(f"  avg latency      {s['avg_latency_ms']:.0f} ms   total cost ${s['total_cost_usd']:.4f}")
    for q in s["misses"]:
        print(f"  miss: {q}")
    if "--langfuse" in argv:
        run_langfuse(copilot)
    copilot.tel.shutdown()
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))