# SPDX-FileCopyrightText: 2026 Ashlesha T
# SPDX-License-Identifier: MIT
#
# This file is part of CogniRepo — https://github.com/ashlesh-t/cognirepo
# Licensed under MIT. See LICENSE file in repository root.

"""
tools/benchmark.py — Quantitative value metrics for CogniRepo.

Measures the concrete advantage CogniRepo provides to an AI agent:

  1. token_reduction_pct   — context_pack vs reading source files raw
  2. symbol_lookup_ms      — ASTIndexer.lookup_symbol latency (target < 5ms)
  3. grep_equivalent_ms    — subprocess grep timing for comparison
  4. cache_speedup_x       — warm hybrid_retrieve vs cold (target ≥ 10x)
  5. memory_recall_at_1    — top-1 accuracy for stored memories  [0..1]
  6. memory_recall_at_3    — recall@3                            [0..1]
  7. context_relevance_pct — % of context_pack sections that mention query terms

Run standalone:
    python -m tools.benchmark
    python -m tools.benchmark --json           # machine-readable JSON
    python -m tools.benchmark --compare        # compare with last run

Saved to .cognirepo/benchmark_history.jsonl for trend tracking.
"""
from __future__ import annotations

import json
import os
import subprocess
import time
from datetime import datetime, timezone
from pathlib import Path

REPO_ROOT = Path(__file__).parent.parent


# ─── helpers ─────────────────────────────────────────────────────────────────

def _count_tokens(text: str) -> int:
    try:
        import tiktoken
        return len(tiktoken.get_encoding("cl100k_base").encode(text))
    except Exception:  # pylint: disable=broad-except
        return max(1, len(text) // 4)


def _read_files_for_query(query_keyword: str, repo_root: Path) -> int:
    """
    Naive baseline: sum of token counts of ALL .py files containing the keyword.
    This inflates savings vs. real Claude Code usage (see _targeted_baseline).
    """
    total = 0
    for py_file in repo_root.rglob("*.py"):
        if any(part in {".git", "venv", "__pycache__", ".cognirepo"}
               for part in py_file.parts):
            continue
        try:
            text = py_file.read_text(encoding="utf-8", errors="replace")
            if query_keyword.lower() in text.lower():
                total += _count_tokens(text)
        except OSError:
            pass
    return total


def _targeted_baseline(query_keyword: str, repo_root: Path) -> int:
    """
    Realistic baseline: grep for exact matches, read top-2 files.
    Approximates what Claude Code would actually do with targeted grep + read_file.
    """
    try:
        result = subprocess.run(
            ["grep", "-rl", query_keyword, str(repo_root),
             "--include=*.py", "--exclude-dir=.git",
             "--exclude-dir=venv", "--exclude-dir=__pycache__",
             "--exclude-dir=.cognirepo"],
            capture_output=True, text=True, timeout=10,
        )
        matched_files = [Path(p) for p in result.stdout.strip().splitlines() if p][:2]
    except (subprocess.TimeoutExpired, OSError):
        return 0
    total = 0
    for py_file in matched_files:
        try:
            total += _count_tokens(py_file.read_text(encoding="utf-8", errors="replace"))
        except OSError:
            pass
    return total


# ─── individual benchmarks ────────────────────────────────────────────────────

def measure_token_reduction(queries: list[str]) -> dict:
    """
    For each query, compare:
      raw_tokens  = tokens in all source files containing the keyword
      pack_tokens = context_pack output tokens
    Returns average savings %.
    """
    from tools.context_pack import context_pack
    from memory.circuit_breaker import CircuitOpenError

    savings_naive = []
    savings_targeted = []
    details = []
    skipped = []
    for q in queries:
        keyword = q.split()[0].lower()
        naive_raw = _read_files_for_query(keyword, REPO_ROOT)
        targeted_raw = _targeted_baseline(keyword, REPO_ROOT)
        if naive_raw == 0:
            continue
        try:
            result = context_pack(q, max_tokens=2000)
        except CircuitOpenError as _cb_err:
            skipped.append({"query": q, "reason": "circuit_open", "detail": str(_cb_err)})
            continue
        packed = result.get("token_count", 0)
        if packed == 0:
            packed = 1  # context was empty — treat as 1 token for ratio purposes
        pct_naive = max(0.0, (naive_raw - packed) / naive_raw * 100)
        pct_targeted = (
            max(0.0, (targeted_raw - packed) / targeted_raw * 100)
            if targeted_raw > 0 else 0.0
        )
        savings_naive.append(pct_naive)
        savings_targeted.append(pct_targeted)
        details.append({
            "query": q,
            "naive_baseline_tokens": naive_raw,
            "targeted_baseline_tokens": targeted_raw,
            "packed_tokens": packed,
            "savings_vs_naive_pct": round(pct_naive, 1),
            "savings_vs_targeted_pct": round(pct_targeted, 1),
        })

    avg_naive = round(sum(savings_naive) / len(savings_naive), 1) if savings_naive else 0.0
    avg_targeted = round(sum(savings_targeted) / len(savings_targeted), 1) if savings_targeted else 0.0
    if skipped:
        print(f"  ⚠  {len(skipped)} quer{'y' if len(skipped)==1 else 'ies'} skipped "
              f"(CogniRepo server under memory pressure). Run: cognirepo server restart",
              flush=True)
    return {
        "token_reduction_pct": avg_naive,
        "token_reduction_vs_targeted_pct": avg_targeted,
        "details": details,
        "skipped": skipped,
    }


def measure_symbol_lookup(symbols: list[str]) -> dict:
    """
    Time ASTIndexer.lookup_symbol for known symbols.
    Returns mean latency in ms.
    """
    from indexer.ast_indexer import ASTIndexer
    from graph.knowledge_graph import KnowledgeGraph

    idx = ASTIndexer(graph=KnowledgeGraph())
    idx.load()

    timings = []
    hit_count = 0
    for sym in symbols:
        t0 = time.perf_counter()
        results = idx.lookup_symbol(sym)
        elapsed_ms = (time.perf_counter() - t0) * 1000
        timings.append(elapsed_ms)
        if results:
            hit_count += 1

    mean_ms = round(sum(timings) / len(timings), 3) if timings else 0.0
    return {
        "symbol_lookup_ms": mean_ms,
        "symbol_hit_rate": round(hit_count / len(symbols), 2) if symbols else 0.0,
        "symbols_tested": len(symbols),
    }


def measure_grep_equivalent(symbols: list[str]) -> dict:
    """
    Time `grep -rn --include="*.py" <symbol>` as a baseline comparison.
    """
    timings = []
    for sym in symbols:
        t0 = time.perf_counter()
        try:
            subprocess.run(
                ["grep", "-rn", "--include=*.py", sym, str(REPO_ROOT)],
                capture_output=True,
                timeout=10,
            )
        except (subprocess.TimeoutExpired, FileNotFoundError):
            pass
        elapsed_ms = (time.perf_counter() - t0) * 1000
        timings.append(elapsed_ms)

    mean_ms = round(sum(timings) / len(timings), 1) if timings else 0.0
    return {"grep_ms": mean_ms}


def measure_cache_speedup(queries: list[str]) -> dict:
    """
    Cold vs warm hybrid_retrieve — cache speedup factor.
    """
    from retrieval.hybrid import hybrid_retrieve, invalidate_hybrid_cache

    speedups = []
    for q in queries:
        invalidate_hybrid_cache()
        t0 = time.perf_counter()
        hybrid_retrieve(q, top_k=5)
        cold_ms = (time.perf_counter() - t0) * 1000

        t1 = time.perf_counter()
        hybrid_retrieve(q, top_k=5)
        warm_ms = (time.perf_counter() - t1) * 1000

        if warm_ms > 0:
            speedups.append(cold_ms / warm_ms)

    mean_speedup = round(sum(speedups) / len(speedups), 1) if speedups else 0.0
    return {"cache_speedup_x": mean_speedup}


def measure_memory_recall(test_memories: list[str]) -> dict:
    """
    Store memories and verify they surface in FAISS top-k results.
    Uses direct vector search (retrieve_memory) rather than hybrid_retrieve
    because hybrid may rank AST symbols above freshly-stored memories.
    Recall@k = fraction where stored text appears in the top-k results.
    """
    from tools.store_memory import store_memory
    from tools.retrieve_memory import retrieve_memory
    from retrieval.hybrid import invalidate_hybrid_cache

    stored = []
    for text in test_memories:
        store_memory(text, source="benchmark")
        stored.append(text)

    recall_at_1 = 0
    recall_at_3 = 0
    for text in stored:
        # Use the unique prefix (first token contains timestamp) as anchor query
        unique_prefix = text.split()[0]  # e.g. "cognirepo_bench_1775420000_a:"
        invalidate_hybrid_cache()  # must flush — new memories aren't in the cache
        results = retrieve_memory(text, top_k=5)
        retrieved_texts = [r.get("text", "").lower() for r in results]

        # The unique prefix guarantees we only match the exact stored memory
        def _is_hit(candidate_texts, _prefix=unique_prefix):
            return any(_prefix.lower() in t for t in candidate_texts)

        if _is_hit(retrieved_texts[:1]):
            recall_at_1 += 1
        if _is_hit(retrieved_texts[:3]):
            recall_at_3 += 1

    n = len(stored)
    return {
        "memory_recall_at_1": round(recall_at_1 / n, 2) if n else 0.0,
        "memory_recall_at_3": round(recall_at_3 / n, 2) if n else 0.0,
        "memories_tested": n,
    }


def measure_precision_at_k(golden: list[dict] | None = None, k: int = 3) -> dict:
    """
    precision@k = fraction of queries where expected_file_pattern appears in top-k sections.
    Loads golden set from tests/fixtures/benchmark_golden.json if not provided.
    """
    from tools.context_pack import context_pack

    if golden is None:
        golden_path = REPO_ROOT / "tests" / "fixtures" / "benchmark_golden.json"
        if not golden_path.exists():
            return {"precision_at_1": 0.0, "precision_at_3": 0.0, "queries_tested": 0,
                    "error": "golden file not found"}
        import json as _json
        golden = _json.loads(golden_path.read_text(encoding="utf-8"))

    hits_at_1 = 0
    hits_at_3 = 0
    tested = 0

    for entry in golden:
        query = entry.get("query", "")
        pattern = entry.get("expected_file_pattern", "")
        if not query or not pattern:
            continue
        try:
            result = context_pack(query, max_tokens=2000)
        except Exception:  # pylint: disable=broad-except
            continue
        sections = result.get("sections", [])
        if not sections:
            continue
        tested += 1
        sources_k1 = [s.get("file", s.get("source", "")) for s in sections[:1]]
        sources_k3 = [s.get("file", s.get("source", "")) for s in sections[:k]]
        if any(pattern in src for src in sources_k1):
            hits_at_1 += 1
        if any(pattern in src for src in sources_k3):
            hits_at_3 += 1

    return {
        "precision_at_1": round(hits_at_1 / tested, 3) if tested else 0.0,
        "precision_at_3": round(hits_at_3 / tested, 3) if tested else 0.0,
        "queries_tested": tested,
    }


def measure_latency(golden: list[dict] | None = None, repeats: int = 3) -> dict:
    """
    Run each golden query `repeats` times, record wall-clock time.
    Returns p50, p95, p99 in ms.
    """
    from tools.context_pack import context_pack
    import statistics

    if golden is None:
        golden_path = REPO_ROOT / "tests" / "fixtures" / "benchmark_golden.json"
        if not golden_path.exists():
            return {"latency_p50_ms": 0.0, "latency_p95_ms": 0.0, "latency_p99_ms": 0.0,
                    "error": "golden file not found"}
        import json as _json
        golden = _json.loads(golden_path.read_text(encoding="utf-8"))

    timings_ms: list[float] = []
    for entry in golden:
        query = entry.get("query", "")
        if not query:
            continue
        for _ in range(repeats):
            t0 = time.perf_counter()
            try:
                context_pack(query, max_tokens=2000)
            except Exception:  # pylint: disable=broad-except
                pass
            timings_ms.append((time.perf_counter() - t0) * 1000)

    if not timings_ms:
        return {"latency_p50_ms": 0.0, "latency_p95_ms": 0.0, "latency_p99_ms": 0.0,
                "latency_samples": 0}

    timings_ms.sort()
    n = len(timings_ms)

    def _percentile(pct: float) -> float:
        idx = int(pct / 100 * n)
        idx = max(0, min(idx, n - 1))
        return round(timings_ms[idx], 1)

    return {
        "latency_p50_ms": _percentile(50),
        "latency_p95_ms": _percentile(95),
        "latency_p99_ms": _percentile(99),
        "latency_samples": n,
    }


def measure_context_relevance(queries: list[str]) -> dict:
    """
    For each query, check what fraction of context_pack sections
    contain at least one query keyword.
    """
    from tools.context_pack import context_pack

    relevance_scores = []
    for q in queries:
        keywords = [w.lower() for w in q.split() if len(w) > 3]
        result = context_pack(q, max_tokens=2000)
        sections = result.get("sections", [])
        if not sections:
            continue
        relevant = sum(
            1 for s in sections
            if any(kw in s.get("content", "").lower() for kw in keywords)
        )
        relevance_scores.append(relevant / len(sections))

    avg = round(sum(relevance_scores) / len(relevance_scores) * 100, 1) if relevance_scores else 0.0
    return {"context_relevance_pct": avg}


# ─── full benchmark run ───────────────────────────────────────────────────────

_BENCHMARK_QUERIES = [
    "store_memory implementation",
    "context_pack token budget",
    "hybrid_retrieve BM25 vector",
    "EpisodicMemory search episodes",
    "knowledge graph node edges",
]

_BENCHMARK_SYMBOLS = [
    "context_pack",
    "hybrid_retrieve",
    "store_memory",
    "log_event",
    "lookup_symbol",
]

def _benchmark_memories() -> list[str]:
    """Generate unique benchmark memories using timestamp to avoid FAISS pollution."""
    ts = int(time.time())
    return [
        f"cognirepo_bench_{ts}_a: token reduction packing context code snippets",
        f"cognirepo_bench_{ts}_b: symbol lookup ASTIndexer reverse index file line",
        f"cognirepo_bench_{ts}_c: hybrid retrieval vector graph behaviour signals",
    ]


def run_benchmark() -> dict:
    """Run all benchmarks and return a consolidated metrics dict."""
    from memory.circuit_breaker import CircuitOpenError

    print("Running CogniRepo benchmark...", flush=True)

    print("  [1/7] Token reduction...", flush=True)
    try:
        token_metrics = measure_token_reduction(_BENCHMARK_QUERIES)
    except CircuitOpenError as _cb_err:
        print(
            f"\n⚠  Benchmark aborted at step [1/7]: CogniRepo server is under memory pressure.\n"
            f"   Error: {_cb_err}\n"
            f"   Fix:   cognirepo server restart\n",
            flush=True,
        )
        return {"aborted": True, "reason": "circuit_open", "detail": str(_cb_err)}

    # Steps 2–7 share the same memory-pressure failure mode as step 1:
    # any embedding/store call can raise CircuitOpenError on large repos.
    # Previously only step 1 was guarded — a breaker trip at step 4
    # (memory recall → store_memory) escaped as a raw traceback.
    try:
        print("  [2/7] Symbol lookup latency...", flush=True)
        lookup_metrics = measure_symbol_lookup(_BENCHMARK_SYMBOLS)
        grep_metrics = measure_grep_equivalent(_BENCHMARK_SYMBOLS[:2])  # grep is slow; test 2

        print("  [3/7] Cache speedup...", flush=True)
        cache_metrics = measure_cache_speedup(_BENCHMARK_QUERIES[:3])

        print("  [4/7] Memory recall...", flush=True)
        recall_metrics = measure_memory_recall(_benchmark_memories())

        print("  [5/7] Context relevance...", flush=True)
        relevance_metrics = measure_context_relevance(_BENCHMARK_QUERIES)

        print("  [6/7] Precision@k (golden set)...", flush=True)
        precision_metrics = measure_precision_at_k()

        print("  [7/7] Latency histogram...", flush=True)
        # Use first 5 golden queries × 3 repeats — enough for p50/p95 without being slow
        golden_path = REPO_ROOT / "tests" / "fixtures" / "benchmark_golden.json"
        if golden_path.exists():
            import json as _json
            golden_subset = _json.loads(golden_path.read_text(encoding="utf-8"))[:5]
        else:
            golden_subset = None
        latency_metrics = measure_latency(golden=golden_subset, repeats=3)
    except CircuitOpenError as _cb_err:
        print(
            f"\n⚠  Benchmark aborted: CogniRepo server is under memory pressure.\n"
            f"   Partial result: token_reduction_pct={token_metrics.get('token_reduction_pct')}\n"
            f"   Error: {_cb_err}\n"
            f"   Fix:   cognirepo server restart\n",
            flush=True,
        )
        return {
            "aborted": True,
            "reason": "circuit_open",
            "detail": str(_cb_err),
            "token_reduction_pct": token_metrics.get("token_reduction_pct"),
        }

    # Composite speedup: grep vs lookup
    lookup_ms = lookup_metrics["symbol_lookup_ms"]
    grep_ms = grep_metrics["grep_ms"]
    lookup_speedup = round(grep_ms / lookup_ms, 1) if lookup_ms > 0 else 0.0

    result = {
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "token_reduction_pct": token_metrics["token_reduction_pct"],
        "symbol_lookup_ms": lookup_metrics["symbol_lookup_ms"],
        "symbol_hit_rate": lookup_metrics["symbol_hit_rate"],
        "grep_ms": grep_metrics["grep_ms"],
        "lookup_speedup_vs_grep_x": lookup_speedup,
        "cache_speedup_x": cache_metrics["cache_speedup_x"],
        "memory_recall_at_1": recall_metrics["memory_recall_at_1"],
        "memory_recall_at_3": recall_metrics["memory_recall_at_3"],
        "context_relevance_pct": relevance_metrics["context_relevance_pct"],
        "precision_at_1": precision_metrics["precision_at_1"],
        "precision_at_3": precision_metrics["precision_at_3"],
        "precision_queries_tested": precision_metrics["queries_tested"],
        "latency_p50_ms": latency_metrics["latency_p50_ms"],
        "latency_p95_ms": latency_metrics["latency_p95_ms"],
        "latency_p99_ms": latency_metrics["latency_p99_ms"],
        "_details": token_metrics.get("details", []),
    }

    # Save to history
    _save_to_history(result)

    return result


def _save_to_history(metrics: dict) -> None:
    """Append metrics to .cognirepo/benchmark_history.jsonl."""
    try:
        from config.paths import get_path
        hist_path = get_path("benchmark_history.jsonl")
        with open(hist_path, "a", encoding="utf-8") as f:
            f.write(json.dumps({k: v for k, v in metrics.items() if k != "_details"}) + "\n")
    except Exception:  # pylint: disable=broad-except
        pass  # non-fatal


def load_last_run() -> dict | None:
    """Load the most recent benchmark result from history."""
    try:
        from config.paths import get_path
        hist_path = get_path("benchmark_history.jsonl")
        if not os.path.exists(hist_path):
            return None
        lines = Path(hist_path).read_text(encoding="utf-8").splitlines()
        valid = [l for l in lines if l.strip()]
        return json.loads(valid[-1]) if valid else None
    except Exception:  # pylint: disable=broad-except
        return None


def print_report(metrics: dict, compare: dict | None = None) -> None:
    """Print a human-readable benchmark report."""
    print("\n" + "═" * 62)
    print("  CogniRepo Benchmark Report")
    print("═" * 62)

    def _row(label: str, value, unit: str = "", prev=None, higher_is_better: bool = True):
        arrow = ""
        if prev is not None and isinstance(value, (int, float)) and isinstance(prev, (int, float)):
            diff = value - prev
            if abs(diff) > 0.01:
                symbol = "▲" if diff > 0 else "▼"
                color = symbol if (diff > 0) == higher_is_better else symbol
                arrow = f"  {color} {abs(diff):.1f}{unit} vs last run"
        print(f"  {label:<38} {value!s:>8}{unit}{arrow}")

    prev = compare or {}
    _row("Token reduction (vs raw file read)", f"{metrics['token_reduction_pct']:.1f}", "%",
         prev.get("token_reduction_pct"))
    _row("Symbol lookup latency (mean)",       f"{metrics['symbol_lookup_ms']:.2f}", " ms",
         prev.get("symbol_lookup_ms"), higher_is_better=False)
    _row("grep equivalent latency (mean)",     f"{metrics['grep_ms']:.0f}", " ms",
         prev.get("grep_ms"), higher_is_better=False)
    _row("Lookup speedup vs grep",             f"{metrics['lookup_speedup_vs_grep_x']:.0f}", "x",
         prev.get("lookup_speedup_vs_grep_x"))
    _row("Cache speedup (warm vs cold)",       f"{metrics['cache_speedup_x']:.0f}", "x",
         prev.get("cache_speedup_x"))
    _row("Memory recall@1",                   f"{metrics['memory_recall_at_1']:.0%}", "",
         prev.get("memory_recall_at_1"))
    _row("Memory recall@3",                   f"{metrics['memory_recall_at_3']:.0%}", "",
         prev.get("memory_recall_at_3"))
    _row("Context relevance",                 f"{metrics['context_relevance_pct']:.1f}", "%",
         prev.get("context_relevance_pct"))
    _row("Symbol hit rate",                   f"{metrics['symbol_hit_rate']:.0%}", "",
         prev.get("symbol_hit_rate"))
    if "precision_at_1" in metrics:
        _row("Precision@1 (golden set)",          f"{metrics['precision_at_1']:.0%}", "",
             prev.get("precision_at_1"))
        _row("Precision@3 (golden set)",          f"{metrics['precision_at_3']:.0%}", "",
             prev.get("precision_at_3"))
    if "latency_p50_ms" in metrics:
        _row("Latency p50",                       f"{metrics['latency_p50_ms']:.0f}", " ms",
             prev.get("latency_p50_ms"), higher_is_better=False)
        _row("Latency p95",                       f"{metrics['latency_p95_ms']:.0f}", " ms",
             prev.get("latency_p95_ms"), higher_is_better=False)
        _row("Latency p99",                       f"{metrics['latency_p99_ms']:.0f}", " ms",
             prev.get("latency_p99_ms"), higher_is_better=False)

    print("─" * 62)
    print(f"  Timestamp: {metrics['timestamp']}")
    print("═" * 62)


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="CogniRepo benchmark")
    parser.add_argument("--json", action="store_true", help="Output raw JSON")
    parser.add_argument("--compare", action="store_true", help="Compare with last run")
    args = parser.parse_args()

    _results = run_benchmark()

    if args.json:
        print(json.dumps({k: v for k, v in _results.items() if k != "_details"}, indent=2))
    else:
        _prev = load_last_run() if args.compare else None
        print_report(_results, compare=_prev)
