#!/usr/bin/env python3
"""Cache-aware results summarizer for the cache-aware-context experiment.

Extends trials/osworld-image-resize-experiment/summarize_results.py, fixing two
accounting flaws in the original:
  1. It summed usage once per traj.jsonl ROW, but batched actions write one row
     per action sharing a single API call's usage -> multi-counting. Here usage
     is deduped per step_num (one API call per predict step).
  2. It ignored cache fields. usage.input_tokens on cached runs is only the
     uncached remainder. We report:
       raw prompt tokens  = input + cache_creation + cache_read
       cost-weighted      = input + 1.25*cache_creation + 0.1*cache_read
     plus $ estimates at Sonnet list price and OpenRouter upstream costs.

Also reports per-step cache hit rate, cache-invalidation events (steps where
cache_read falls well below the previous step's prompt size), policy firing
events, and summarizer-call usage (context_mgmt_usage) as full-price tokens.
"""
import csv
import json
import sys
from datetime import datetime
from pathlib import Path

ROOT = Path("/srv/controller-data")
EXPERIMENT = sys.argv[1] if len(sys.argv) > 1 else "focuschain-ab"
BASE = ROOT / "results" / EXPERIMENT
OUT = ROOT / "trials" / "focuschain-ab" / "summary" / EXPERIMENT

# Sonnet list pricing per token
P_IN = 3e-6
P_OUT = 15e-6
W_CREATE = 1.25
W_READ = 0.1


def parse_ts(ts: str):
    for fmt in ("%Y%m%d@%H%M%S%f", "%Y%m%d@%H%M%S"):
        try:
            return datetime.strptime(ts, fmt)
        except ValueError:
            continue
    return None


def task_metrics(traj_path: Path):
    """Parse traj.jsonl -> per-step usage (deduped by step_num) and task metrics."""
    steps = {}  # step_num -> dict
    timestamps = []
    if traj_path.exists():
        for line in traj_path.read_text(errors="replace").splitlines():
            try:
                row = json.loads(line)
            except Exception:
                continue
            ts = parse_ts(row.get("action_timestamp", "") or "")
            if ts:
                timestamps.append(ts)
            action = row.get("action") or {}
            usage = action.get("usage") or {}
            step_num = row.get("step_num")
            if usage and step_num is not None and step_num not in steps:
                steps[step_num] = usage

    per_step = []
    for step_num in sorted(steps):
        u = steps[step_num]
        inp = int(u.get("input_tokens") or 0)
        out = int(u.get("output_tokens") or 0)
        creation = int(u.get("cache_creation_input_tokens") or 0)
        read = int(u.get("cache_read_input_tokens") or 0)
        cost_details = u.get("cost_details") or {}
        upstream = float(cost_details.get("upstream_inference_cost") or 0.0)
        cm = u.get("context_mgmt_usage") or {}
        cm_in = int(cm.get("input_tokens") or 0) + int(cm.get("cache_creation_input_tokens") or 0) + int(cm.get("cache_read_input_tokens") or 0)
        cm_out = int(cm.get("output_tokens") or 0)
        cm_cost_details = cm.get("cost_details") or {}
        cm_upstream = float(cm_cost_details.get("upstream_inference_cost") or 0.0)
        per_step.append({
            "step": step_num,
            "input": inp,
            "output": out,
            "creation": creation,
            "read": read,
            "prompt_total": inp + creation + read,
            "upstream_cost": upstream,
            "context_event": u.get("context_event") or "",
            "cm_input": cm_in,
            "cm_output": cm_out,
            "cm_upstream_cost": cm_upstream,
        })

    m = {
        "steps": len(per_step),
        "input_tokens": sum(s["input"] for s in per_step),
        "output_tokens": sum(s["output"] for s in per_step),
        "cache_creation_tokens": sum(s["creation"] for s in per_step),
        "cache_read_tokens": sum(s["read"] for s in per_step),
        "cm_input_tokens": sum(s["cm_input"] for s in per_step),
        "cm_output_tokens": sum(s["cm_output"] for s in per_step),
        "upstream_cost": sum(s["upstream_cost"] + s["cm_upstream_cost"] for s in per_step),
        "compactions": sum(1 for s in per_step if s["context_event"].startswith("compaction_fired")),
        "image_trims": sum(1 for s in per_step if s["context_event"].startswith("image_trim_fired")),
    }
    m["raw_prompt_tokens"] = m["input_tokens"] + m["cache_creation_tokens"] + m["cache_read_tokens"]
    m["raw_total_tokens"] = m["raw_prompt_tokens"] + m["output_tokens"] + m["cm_input_tokens"] + m["cm_output_tokens"]
    # cost-weighted input-equivalents (mission accounting rule); summarizer
    # calls are uncached full-price input
    m["cost_weighted_input"] = (
        m["input_tokens"]
        + W_CREATE * m["cache_creation_tokens"]
        + W_READ * m["cache_read_tokens"]
        + m["cm_input_tokens"]
    )
    m["est_cost_usd"] = (
        m["cost_weighted_input"] * P_IN
        + (m["output_tokens"] + m["cm_output_tokens"]) * P_OUT
    )
    # cache hit rate: fraction of prompt tokens served from cache (steps >= 2;
    # step 1 is always a cold write)
    tail = [s for s in per_step if s["step"] >= 2]
    denom = sum(s["prompt_total"] for s in tail)
    m["cache_hit_rate"] = (sum(s["read"] for s in tail) / denom) if denom else 0.0
    # invalidation events: cache_read falls below half of the previous step's
    # total prompt (an append-only cached agent should re-read ~everything)
    inval = 0
    for prev, cur in zip(per_step, per_step[1:]):
        if prev["prompt_total"] > 2000 and cur["read"] < 0.5 * prev["prompt_total"]:
            inval += 1
    m["invalidation_events"] = inval
    m["wall_clock_s"] = (
        (max(timestamps) - min(timestamps)).total_seconds() if len(timestamps) >= 2 else 0.0
    )
    return m, per_step


def read_result(path: Path):
    try:
        return float(path.read_text().strip())
    except Exception:
        return 0.0


def main():
    OUT.mkdir(parents=True, exist_ok=True)
    rows, step_rows = [], []
    for variant_dir in sorted(p for p in BASE.glob("*") if p.is_dir()):
        for result_file in variant_dir.glob("claude_computer_use/screenshot/*/*/*/result.txt"):
            task_dir = result_file.parent
            domain = task_dir.parent.name
            task_id = task_dir.name
            m, per_step = task_metrics(task_dir / "traj.jsonl")
            row = {
                "variant": variant_dir.name,
                "domain": domain,
                "task_id": task_id,
                "score": read_result(result_file),
                **{k: (round(v, 6) if isinstance(v, float) else v) for k, v in m.items()},
            }
            rows.append(row)
            for s in per_step:
                step_rows.append({
                    "variant": variant_dir.name,
                    "domain": domain,
                    "task_id": task_id,
                    **s,
                })

    if not rows:
        print(f"No results found under {BASE}")
        return

    results_csv = OUT / "results.csv"
    with results_csv.open("w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
        writer.writeheader()
        writer.writerows(rows)

    steps_csv = OUT / "steps.csv"
    with steps_csv.open("w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=list(step_rows[0].keys()))
        writer.writeheader()
        writer.writerows(step_rows)

    md = OUT / "summary.md"
    with md.open("w") as f:
        f.write(f"# {EXPERIMENT} summary\n\n")
        f.write("| Variant | Tasks | Score | Steps | Raw prompt tok | Cache read | Cache write | Uncached in | Output | Cost-weighted in | Est $ | Upstream $ | Hit rate | Invalidations | Compactions | Trims | Wall (s) |\n")
        f.write("|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|\n")
        for variant in sorted(set(r["variant"] for r in rows)):
            sub = [r for r in rows if r["variant"] == variant]
            n = len(sub)
            agg = lambda k: sum(r[k] for r in sub)
            hit = sum(r["cache_hit_rate"] * r["raw_prompt_tokens"] for r in sub) / max(1, agg("raw_prompt_tokens"))
            f.write(
                f"| {variant} | {n} | {agg('score'):g} | {agg('steps')} "
                f"| {agg('raw_prompt_tokens'):,} | {agg('cache_read_tokens'):,} | {agg('cache_creation_tokens'):,} "
                f"| {agg('input_tokens'):,} | {agg('output_tokens'):,} "
                f"| {round(agg('cost_weighted_input')):,} | {agg('est_cost_usd'):.4f} | {agg('upstream_cost'):.4f} "
                f"| {hit:.3f} | {agg('invalidation_events')} | {agg('compactions')} | {agg('image_trims')} "
                f"| {agg('wall_clock_s'):.0f} |\n"
            )
    print(md)
    print(results_csv)
    print(steps_csv)


if __name__ == "__main__":
    main()
