"""Reproducible retrieval evaluation for agentmemory.

Run with::

    python -m agentmemory.eval

It builds a small, seeded, self-contained corpus of agent observations and a
set of queries with known relevant memory ids, runs the bundled retrievers, and
prints a metrics report. The metrics are computed from this run and saved to a
JSON file; the numbers shown on the site's /evaluation/ page are the output of
exactly this command (with a matching seed), so they are reproducible.

This is intentionally *not* a proprietary benchmark and makes no claim about
commercial systems. It is a transparent, runnable reference evaluation.
"""

from __future__ import annotations

import argparse
import json
import random
import sys
from pathlib import Path
from typing import Dict, List, Tuple

from .mem import Entry, MemoryStore
from .retrieval import BM25Scorer, Retrieval, TagScorer


# ---------------------------------------------------------------------------
# reproducible corpus: agent observations about operating a fleet of servers
# ---------------------------------------------------------------------------
def build_corpus(seed: int = 7) -> List[Entry]:
    rng = random.Random(seed)
    base = [
        ("deployed v2.3 of the order service to production", "release-notes/2026-03-01.md", ["deploy", "order"]),
        ("order service latency spiked to 800ms after v2.3 rollout", "incidents/2026-03-02.md", ["incident", "order", "latency"]),
        ("rollback to v2.2 restored order service latency to 120ms", "incidents/2026-03-02-followup.md", ["rollback", "order"]),
        ("the payment webhook uses HMAC-SHA256 with a rotating key", "docs/payment-webhook.md", ["payment", "security"]),
        ("new relic key rotates every 30 days; stored in vault path secret/pay", "docs/payment-webhook.md", ["payment", "vault"]),
        ("database backup runs nightly at 02:00 UTC to s3://backups", "runbooks/backups.md", ["backup", "database"]),
        ("restore drills happen quarterly and take about 40 minutes", "runbooks/backups.md", ["backup", "drill"]),
        ("auth tokens expire after 15 minutes; refresh uses rotating cookies", "docs/auth.md", ["auth", "security"]),
        ("the auth service can be scaled to 6 replicas without state", "docs/auth.md", ["auth", "scaling"]),
        ("on-call number is +1-555-0100 and routes to the primary engineer", "runbooks/oncall.md", ["oncall", "ops"]),
        ("postmortem 42 cites the missing read-replica timeout as root cause", "postmortems/pm42.md", ["postmortem", "database"]),
        ("SLO for order availability is 99.9% monthly; current is 99.95%", "slo.md", ["slo", "order"]),
        ("feature flag 'checkout-v2' is 10% rolled out", "flags.md", ["flag", "checkout"]),
        ("checkout-v2 rollback key is stored under vault path secret/checkout", "flags.md", ["checkout", "vault"]),
    ]
    # duplicate some observations to give consolidation something to do
    entries = [Entry(text=t, source=s, tags=list(tags)) for (t, s, tags) in base]
    entries.append(Entry(text=base[0][0], source="release-notes/2026-03-01.md", tags=["deploy", "order"]))
    entries.append(Entry(text=base[2][0], source="incidents/2026-03-02-followup.md", tags=["rollback", "order"]))
    # shuffle deterministically to emulate out-of-order arrival
    entries = [entries[i] for i in rng.sample(range(len(entries)), len(entries))]
    return entries


def build_queries() -> List[Tuple[str, List[str]]]:
    """Return (query_text, list_of_relevant_memory_ids)."""
    return [
        ("why is the order service slow after the deploy", ["l1", "l2"]),
        ("how do we restore a database", ["l6", "l7"]),
        ("how long do auth tokens last", ["l8"]),
        ("who to call on nights and weekends", ["l10"]),
        ("where is the checkout rollout key", ["l13", "l14"]),
    ]


def _metrics_at_k(relevant_ids, returned_ids, k) -> Dict:
    rel = set(relevant_ids)
    returned = returned_ids[:k]
    hit = len(rel & set(returned))
    recall = hit / len(rel) if rel else 0.0
    precision = hit / k if k else 0.0
    # reciprocal rank over the returned window
    rr = 0.0
    for rank, rid in enumerate(returned, start=1):
        if rid in rel:
            rr = 1.0 / rank
            break
    return {"k": k, "recall@k": recall, "precision@k": precision, "rr": rr}


def evaluate(entries: List[Entry], queries: List[Tuple[str, List[str]]], k: int = 5) -> Dict:
    store = MemoryStore()
    for e in entries:
        store.add(e)
    results = {
        "seed": None,
        "k": k,
        "num_entries": len(entries),
        "num_queries": len(queries),
        "scorers": {},
        "macro": {},
    }
    scorers = {
        "bm25": BM25Scorer(),
        "tag": TagScorer(),
    }
    for name, scorer in scorers.items():
        r = Retrieval(scorer=scorer)
        per_query = []
        all_recall, all_prec, all_rr = [], [], []
        for q, rel in queries:
            rel_ids = []
            for rid in rel:
                # build a stable id lookup by matching the query-position labels
                rel_ids.append(_resolve_id(store, rid, queries, entries))
            out = r.search_terms(entries, q, k=k)
            returned_ids = [s.entry.id for s in out]
            m = _metrics_at_k(rel_ids, returned_ids, k)
            per_query.append({"query": q, "relevant": rel_ids, "returned": returned_ids[:k], "metrics": m})
            all_recall.append(m["recall@k"])
            all_prec.append(m["precision@k"])
            all_rr.append(m["rr"])
        results["scorers"][name] = {
            "per_query": per_query,
            "aggregate": {
                "mean_recall@k": sum(all_recall) / len(all_recall),
                "mean_precision@k": sum(all_prec) / len(all_prec),
                "mean_reciprocal_rank": sum(all_rr) / len(all_rr),
            },
        }
    return results


def _resolve_id(store: MemoryStore, label: str, queries, entries) -> str:
    """Map a query label like 'l1' to an actual entry id deterministically."""
    # Build label->entry index by matching text against the corpus base order.
    corpus_order = [
        "deployed v2.3 of the order service",
        "order service latency spiked",
        "rollback to v2.2 restored",
        "payment webhook uses HMAC",
        "new relic key rotates every 30 days",
        "database backup runs nightly",
        "restore drills happen quarterly",
        "auth tokens expire after 15 minutes",
        "auth service can be scaled to 6 replicas",
        "on-call number is +1-555-0100",
        "postmortem 42 cites the missing read-replica",
        "SLO for order availability is 99.9%",
        "feature flag 'checkout-v2' is 10%",
        "checkout-v2 rollback key is stored",
    ]
    n = int(label[1:]) - 1
    text_frag = corpus_order[n]
    for e in store:
        if text_frag in e.text:
            return e.id
    raise KeyError(f"could not resolve {label}")


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description="agentmemory retrieval evaluation")
    ap.add_argument("--seed", type=int, default=7)
    ap.add_argument("--k", type=int, default=5)
    ap.add_argument("--out", type=str, default=None, help="path to write results.json")
    args = ap.parse_args(argv)

    entries = build_corpus(args.seed)
    queries = build_queries()
    results = evaluate(entries, queries, k=args.k)
    results["seed"] = args.seed
    results["engine"] = "agentmemory-" + _version()

    print("=" * 64)
    print("agentmemory retrieval evaluation (reproducible, seeded)")
    print(f"  seed={args.seed}  entries={results['num_entries']}  queries={results['num_queries']}  k={args.k}")
    print("=" * 64)
    for name, blk in results["scorers"].items():
        agg = blk["aggregate"]
        print(f"\n[{name}]")
        print(f"  mean_recall@k      = {agg['mean_recall@k']:.4f}")
        print(f"  mean_precision@k   = {agg['mean_precision@k']:.4f}")
        print(f"  mean_reciprocal_rank = {agg['mean_reciprocal_rank']:.4f}")

    out_path = args.out
    if not out_path:
        here = Path(__file__).resolve().parent.parent.parent
        out_path = str(here / "eval-results.json")
    Path(out_path).parent.mkdir(parents=True, exist_ok=True)
    with open(out_path, "w") as fh:
        json.dump(results, fh, indent=2)
    print(f"\nwrote results -> {out_path}")
    return 0


def _version() -> str:
    from . import __version__
    return __version__


if __name__ == "__main__":
    sys.exit(main())