#!/usr/bin/env python3
"""Summarize trufflepig-agent audit logs and flag sessions where the tool struggled.

Usage: trufflepig-audit [--harness NAME] [--session ID] [--since HOURS]
                        [--json] [--calls] [--adoption] [LOG_DIR]

Reads <LOG_DIR>/<harness>.jsonl written by `trufflepig-agent` (default
$XDG_STATE_HOME/trufflepig/agent-audit plus the sandbox fallback
configured disk-backed runtime directory) and prints one row per session with
call counts, error and empty-result rates, truncation, repeated queries,
latency, and follow-through (searches followed by show/ctx/more). A session
is flagged when any struggle threshold is crossed. `--calls` lists the
individual calls that carried signals. Exit code 3 when any session is flagged.
"""
from __future__ import annotations

import argparse
import json
import os
import statistics
import sys
import time
from collections import defaultdict
from pathlib import Path

from trufflepig_runtime import fallback_root

THRESHOLDS = {
    "error_rate": 0.25,
    "no_hit_rate": 0.5,
    "repeat_queries": 3,
    "budget_errors": 2,
    "stale_errors": 2,
    "usage_errors": 2,
    "p95_seconds": 5.0,
    "min_calls": 3,
}


def default_dir() -> Path:
    base = os.environ.get("XDG_STATE_HOME") or os.path.join(os.path.expanduser("~"), ".local", "state")
    return Path(os.environ.get("TRUFFLEPIG_AGENT_LOG_DIR") or Path(base) / "trufflepig" / "agent-audit")


def fallback_dir() -> Path:
    return fallback_root() / "agent-audit"


def load(directories: list[Path], harness: str | None, since_hours: float | None) -> list[dict]:
    records: list[dict] = []
    cutoff = time.time() - since_hours * 3600 if since_hours else None
    for path in sorted(path for directory in directories for path in directory.glob("*.jsonl")):
        if harness and path.stem != harness:
            continue
        with path.open(encoding="utf-8") as handle:
            for line in handle:
                line = line.strip()
                if not line:
                    continue
                try:
                    record = json.loads(line)
                except ValueError:
                    records.append({"harness": path.stem, "session": "?", "malformed": True})
                    continue
                if cutoff is not None:
                    try:
                        stamp = time.mktime(time.strptime(record["ts"][:19], "%Y-%m-%dT%H:%M:%S"))
                    except (KeyError, ValueError):
                        stamp = None
                    if stamp is not None and stamp < cutoff:
                        continue
                records.append(record)
    records.sort(key=lambda record: record.get("ts", ""))
    return records


def summarize(records: list[dict]) -> list[dict]:
    sessions: dict[tuple[str, str], list[dict]] = defaultdict(list)
    for record in records:
        sessions[(record.get("harness", "?"), record.get("session", "?"))].append(record)
    rows = []
    for (harness, session), calls in sessions.items():
        calls = [c for c in calls if not c.get("malformed")]
        if not calls:
            continue
        hooks = [c for c in calls if str(c.get("verb", "")).startswith("hook:")]
        calls = [c for c in calls if not str(c.get("verb", "")).startswith("hook:")]
        signals = [s for c in [*calls, *hooks] for s in c.get("signals", [])]
        if not calls:
            calls = hooks  # a session that only tripped the hook still shows up
        searches = [c for c in calls if c.get("verb") == "search"]
        latencies = [c.get("elapsed_ms", 0) / 1000 for c in calls]
        follow = 0
        for index, call in enumerate(calls):
            if call.get("verb") == "search" and (call.get("hits") or 0) > 0:
                if any(later.get("verb") in ("show", "ctx", "more") for later in calls[index + 1:index + 4]):
                    follow += 1
        productive = sum(1 for c in searches if (c.get("hits") or 0) > 0)
        row = {
            "harness": harness,
            "session": session,
            "first": calls[0].get("ts"),
            "last": calls[-1].get("ts"),
            "calls": len(calls),
            "searches": len(searches),
            "shows": sum(1 for c in calls if c.get("verb") in ("show", "more")),
            "ctx": sum(1 for c in calls if c.get("verb") == "ctx"),
            "errors": signals.count("error"),
            "no_hits": signals.count("no_hits"),
            "truncated": signals.count("truncated"),
            "repeats": sum(1 for c in calls if "repeat_query" in c.get("signals", [])),
            "budget": signals.count("budget"),
            "stale": signals.count("stale"),
            "usage_errors": signals.count("usage_error"),
            "semantic_degraded": signals.count("semantic_degraded"),
            "rerank_unavailable": signals.count("rerank_unavailable"),
            "slow": signals.count("slow"),
            "steer": signals.count("steer_block") + signals.count("steer_nudge"),
            "p50_s": round(statistics.median(latencies), 2),
            "p95_s": round(sorted(latencies)[max(0, int(len(latencies) * 0.95) - 1)], 2),
            "follow_through": f"{follow}/{productive}",
        }
        reasons = []
        if row["calls"] >= THRESHOLDS["min_calls"]:
            if row["errors"] / row["calls"] > THRESHOLDS["error_rate"]:
                reasons.append("error rate")
            if searches and row["no_hits"] / len(searches) > THRESHOLDS["no_hit_rate"]:
                reasons.append("empty searches")
            if row["repeats"] >= THRESHOLDS["repeat_queries"]:
                reasons.append("repeated queries")
            if row["budget"] >= THRESHOLDS["budget_errors"]:
                reasons.append("budget errors")
            if row["stale"] >= THRESHOLDS["stale_errors"]:
                reasons.append("stale handles")
            if row["usage_errors"] >= THRESHOLDS["usage_errors"]:
                reasons.append("usage errors")
            if row["p95_s"] > THRESHOLDS["p95_seconds"]:
                reasons.append("slow")
            if productive >= 3 and follow == 0:
                reasons.append("results never opened")
        row["struggling"] = bool(reasons)
        row["reasons"] = reasons
        rows.append(row)
    rows.sort(key=lambda r: (r["first"] or ""), reverse=True)
    return rows


NAVIGATION_VERBS = {"search", "refs", "map", "show", "ctx", "more"}


def adoption(records: list[dict]) -> list[dict]:
    """Per harness and agent scope: trufflepig navigation calls against ordinary searches
    the steering hook saw. Blocked searches are excluded from the denominator because
    the agent had to retry them some other way; `read`-only records (`sed -n`, `cat` of
    indexed source) are counted as `reads`, not escaped searches."""
    groups: dict[tuple[str, str], dict] = {}
    for record in records:
        if record.get("malformed"):
            continue
        harness = record.get("harness", "?")
        verb = str(record.get("verb", ""))
        if verb == "hook:steer":
            scope = "subagent" if record.get("agent") else "main"
        elif verb in NAVIGATION_VERBS:
            scope = "all"
        else:
            continue
        row = groups.setdefault((harness, scope), {
            "harness": harness, "scope": scope, "trufflepig": 0, "escaped": 0, "reads": 0, "blocked": 0,
            "fallback": 0, "classes": defaultdict(int)})
        classes = record.get("classes") or ["unclassified"]
        if verb in NAVIGATION_VERBS:
            row["trufflepig"] += 1
        elif record.get("status") == "block":
            row["blocked"] += 1
        else:
            row["reads" if set(classes) == {"read"} else "escaped"] += 1
            row["fallback"] += bool(record.get("fallback"))
            for kind in classes:
                row["classes"][kind] += 1
    rows = []
    for harness in sorted({h for h, _ in groups}):
        used = groups.get((harness, "all"), {}).get("trufflepig", 0)
        escaped = sum(r["escaped"] for (h, s), r in groups.items() if h == harness and s != "all")
        reads = sum(r["reads"] for (h, s), r in groups.items() if h == harness and s != "all")
        for (h, scope), row in sorted(groups.items()):
            if h != harness:
                continue
            row["classes"] = dict(sorted(row["classes"].items(), key=lambda item: -item[1]))
            rows.append(row)
        rows.append({"harness": harness, "scope": "total", "trufflepig": used, "escaped": escaped, "reads": reads,
                     "adoption": round(used / (used + escaped), 3) if used + escaped else None})
    return rows


def print_adoption(rows: list[dict]) -> None:
    print("harness\tscope\ttrufflepig\tescaped\treads\tblocked\tfallback\tadoption\tclasses")
    for row in rows:
        classes = ", ".join(f"{k}={v}" for k, v in (row.get("classes") or {}).items())
        adoption_rate = "" if row.get("adoption") is None else f"{row['adoption']:.0%}"
        print("\t".join(str(v) for v in (row["harness"], row["scope"], row["trufflepig"], row["escaped"],
                                         row.get("reads", ""), row.get("blocked", ""), row.get("fallback", ""),
                                         adoption_rate, classes)))


def print_table(rows: list[dict]) -> None:
    header = ("harness", "session", "calls", "srch", "show", "ctx", "err", "empty", "trunc",
              "rept", "steer", "p50s", "p95s", "follow", "flag")
    print("\t".join(header))
    for row in rows:
        print("\t".join(str(v) for v in (
            row["harness"], row["session"][:28], row["calls"], row["searches"], row["shows"],
            row["ctx"], row["errors"], row["no_hits"], row["truncated"], row["repeats"],
            row["steer"], row["p50_s"], row["p95_s"], row["follow_through"],
            ("STRUGGLING: " + ", ".join(row["reasons"])) if row["struggling"] else "ok",
        )))


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("log_dir", nargs="?", type=Path, help="audit log directory")
    parser.add_argument("--harness")
    parser.add_argument("--session")
    parser.add_argument("--since", type=float, help="only calls from the last N hours")
    parser.add_argument("--json", action="store_true")
    parser.add_argument("--calls", action="store_true", help="list calls carrying signals")
    parser.add_argument("--adoption", action="store_true",
                        help="trufflepig calls versus ordinary searches seen by the steering hook")
    args = parser.parse_args()
    directories = [args.log_dir] if args.log_dir else [default_dir(), fallback_dir()]
    directories = list(dict.fromkeys(directory.resolve() for directory in directories if directory.is_dir()))
    if not directories:
        print(f"no audit log directory at {args.log_dir or default_dir()}", file=sys.stderr)
        return 1
    records = load(directories, args.harness, args.since)
    if args.session:
        records = [r for r in records if r.get("session") == args.session]
    if args.adoption:
        rows = adoption(records)
        if args.json:
            print(json.dumps({"adoption": rows}, indent=1))
        else:
            print_adoption(rows)
        return 0
    rows = summarize(records)
    if args.json:
        print(json.dumps({"sessions": rows, "records": len(records)}, indent=1))
    else:
        print_table(rows)
        if args.calls:
            print()
            for record in records:
                if record.get("signals"):
                    print(f"{record.get('ts')} {record.get('harness')} {record.get('session', '')[:20]} "
                          f"{record.get('verb')} {' '.join(record.get('args', []))[:60]!r} "
                          f"{record.get('elapsed_ms')}ms hits={record.get('hits')} "
                          f"signals={','.join(record['signals'])} {record.get('error') or record.get('stderr') or ''}"[:220])
        print(f"\n{len(records)} calls, {len(rows)} sessions, "
              f"{sum(1 for r in rows if r['struggling'])} flagged. "
              f"Daemon-side ledger: `trufflepig audit` in the repository.")
    return 3 if any(r["struggling"] for r in rows) else 0


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