#!/usr/bin/env python3
"""Generate all characterization CSV/JSON data for the qwenapp trace.

Reads the formatted qwenapp trace (release format with hash_ids) and produces:
  1. token_length_cdf_qwenapp.csv + .stats.json
  2. f2a_block_lifetime_cdf_qwenapp.csv + .stats.json
  3. f2a_reuse_gap_cdf_qwenapp.csv
  4. f2a_reuse_timeline_qwenapp.csv
  5. f2b_session_skew_cdf_qwenapp.json

Run on dash1:
    python paper/data/analyze_qwenapp_trace.py \
        ~/ali-trace/trace-qwenapp-formatted/0529-1000-1200-ttp.jsonl
"""
from __future__ import annotations

import argparse
import csv
import json
import math
import sys
from collections import Counter, defaultdict
from pathlib import Path

import numpy as np


def _percentile(sorted_arr, p):
    idx = int(len(sorted_arr) * p / 100.0)
    return sorted_arr[min(idx, len(sorted_arr) - 1)]


def _log_bins(lo, hi, n=200):
    return np.logspace(np.log10(max(lo, 1e-4)), np.log10(hi), n)


# ---------------------------------------------------------------------------
# 1. Token length CDF
# ---------------------------------------------------------------------------
def build_token_length(rows, sessions, outdir):
    input_toks = sorted(r["input_length"] for r in rows)
    output_toks = sorted(r["output_length"] for r in rows)

    append_toks = []
    for sid, sess_rows in sessions.items():
        if len(sess_rows) < 2:
            continue
        for i in range(1, len(sess_rows)):
            delta = sess_rows[i]["input_length"] - sess_rows[i - 1]["input_length"]
            append_toks.append(delta)
    append_toks.sort()

    n = max(len(input_toks), len(output_toks), len(append_toks))
    cdf_points = np.linspace(0, 1, min(n, 2000))

    with (outdir / "token_length_cdf_qwenapp.csv").open("w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["cdf", "input_tokens", "append_tokens", "output_tokens"])
        for c in cdf_points:
            ii = min(int(c * len(input_toks)), len(input_toks) - 1)
            ai = min(int(c * len(append_toks)), len(append_toks) - 1) if append_toks else 0
            oi = min(int(c * len(output_toks)), len(output_toks) - 1)
            w.writerow([
                f"{c:.3f}",
                input_toks[ii],
                append_toks[ai] if append_toks else 0,
                output_toks[oi],
            ])

    n_root = sum(1 for sid, sr in sessions.items() for i, r in enumerate(sr) if i == 0)
    n_neg = sum(1 for a in append_toks if a < 0)
    stats = {
        "trace": "trace-qwenapp-formatted/0529-1000-1200.jsonl",
        "n_records": len(rows),
        "input": {
            "n": len(input_toks), "p50": _percentile(input_toks, 50),
            "p90": _percentile(input_toks, 90), "p99": _percentile(input_toks, 99),
            "mean": sum(input_toks) / len(input_toks),
            "min": input_toks[0], "max": input_toks[-1],
        },
        "append": {
            "n": len(append_toks),
            "p50": _percentile(append_toks, 50) if append_toks else 0,
            "p90": _percentile(append_toks, 90) if append_toks else 0,
            "p99": _percentile(append_toks, 99) if append_toks else 0,
            "mean": sum(append_toks) / len(append_toks) if append_toks else 0,
            "min": append_toks[0] if append_toks else 0,
            "max": append_toks[-1] if append_toks else 0,
            "n_root": n_root, "n_negative": n_neg,
        },
        "output": {
            "n": len(output_toks), "p50": _percentile(output_toks, 50),
            "p90": _percentile(output_toks, 90), "p99": _percentile(output_toks, 99),
            "mean": sum(output_toks) / len(output_toks),
            "min": output_toks[0], "max": output_toks[-1],
        },
    }
    (outdir / "token_length_cdf_qwenapp.stats.json").write_text(
        json.dumps(stats, indent=4, ensure_ascii=False))
    print(f"  token_length: {len(rows)} rows, append={len(append_toks)}")


# ---------------------------------------------------------------------------
# 2. Block lifetime CDF
# ---------------------------------------------------------------------------
def build_block_lifetime(rows, outdir):
    block_first: dict[int, float] = {}
    block_last: dict[int, float] = {}
    for r in rows:
        ts = r["timestamp"]
        for h in r["hash_ids"]:
            if h not in block_first:
                block_first[h] = ts
            block_last[h] = ts

    lifetimes = []
    for h in block_first:
        lt = block_last[h] - block_first[h]
        if lt > 0:
            lifetimes.append(lt)
    lifetimes.sort()

    if not lifetimes:
        print("  block_lifetime: no reused blocks!")
        return

    bins = np.concatenate([
        _log_bins(0.01, lifetimes[-1] * 1.01, 300),
        [float("inf")],
    ])
    counts, _ = np.histogram(lifetimes, bins=bins)

    with (outdir / "f2a_block_lifetime_cdf_qwenapp.csv").open("w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["bucket_upper_s", "count"])
        for edge, cnt in zip(bins[1:], counts):
            w.writerow([edge, int(cnt)])

    stats = {
        "n_records": len(rows),
        "n_distinct_blocks": len(block_first),
        "n_reused_blocks": len(lifetimes),
        "reused_fraction": len(lifetimes) / len(block_first) if block_first else 0,
        "p50_s": _percentile(lifetimes, 50),
        "p90_s": _percentile(lifetimes, 90),
        "p99_s": _percentile(lifetimes, 99),
        "mean_s": sum(lifetimes) / len(lifetimes),
        "max_s": lifetimes[-1],
    }
    (outdir / "f2a_block_lifetime_cdf_qwenapp.stats.json").write_text(
        json.dumps(stats, indent=4, ensure_ascii=False))
    print(f"  block_lifetime: {len(lifetimes)} reused blocks / {len(block_first)} total")


# ---------------------------------------------------------------------------
# 3. Reuse gap CDF
# ---------------------------------------------------------------------------
def build_reuse_gap(rows, sessions, outdir):
    block_last_touch: dict[int, float] = {}
    block_session: dict[int, str] = {}

    # For qwenapp: no system/content split via tools. Instead:
    #   intra-session = reuse within the same session
    #   cross-session = reuse from a different session
    # We don't have a system-prompt block identifier, so approximate:
    #   system-prompt blocks = the first N blocks of each request that match
    #   across sessions (shared system prompt prefix)
    gaps_intra = []
    gaps_cross = []

    for r in rows:
        ts = r["timestamp"]
        sid = r["session_id"]
        for h in r["hash_ids"]:
            if h in block_last_touch:
                gap = ts - block_last_touch[h]
                if gap > 0:
                    prev_sid = block_session[h]
                    if prev_sid == sid:
                        gaps_intra.append(gap)
                    else:
                        gaps_cross.append(gap)
            block_last_touch[h] = ts
            block_session[h] = sid

    max_gap = max(
        max(gaps_intra) if gaps_intra else 1,
        max(gaps_cross) if gaps_cross else 1,
    )
    bins = np.concatenate([
        _log_bins(0.01, max_gap * 1.01, 300),
        [float("inf")],
    ])
    h_intra, _ = np.histogram(gaps_intra, bins=bins) if gaps_intra else (np.zeros(len(bins) - 1), None)
    h_cross, _ = np.histogram(gaps_cross, bins=bins) if gaps_cross else (np.zeros(len(bins) - 1), None)

    with (outdir / "f2a_reuse_gap_cdf_qwenapp.csv").open("w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["bucket_upper_s", "intra", "cross_session"])
        for edge, ci, cc in zip(bins[1:], h_intra, h_cross):
            w.writerow([edge, int(ci), int(cc)])

    print(f"  reuse_gap: intra={len(gaps_intra)}, cross={len(gaps_cross)}")


# ---------------------------------------------------------------------------
# 4. Reuse timeline
# ---------------------------------------------------------------------------
def build_reuse_timeline(rows, sessions, outdir, bin_s=30.0):
    max_ts = max(r["timestamp"] for r in rows)
    n_bins = int(math.ceil(max_ts / bin_s))

    bin_data = [{"n_req": 0, "intra_tok": 0, "cross_tok": 0,
                 "new_tok": 0, "total_tok": 0} for _ in range(n_bins)]

    block_last_touch: dict[int, float] = {}
    block_session: dict[int, str] = {}

    for r in rows:
        ts = r["timestamp"]
        b = min(int(ts / bin_s), n_bins - 1)
        bin_data[b]["n_req"] += 1

        n_blocks = len(r["hash_ids"])
        reuse_intra = 0
        reuse_cross = 0
        new_blocks = 0
        sid = r["session_id"]

        for h in r["hash_ids"]:
            if h in block_last_touch:
                prev_sid = block_session[h]
                if prev_sid == sid:
                    reuse_intra += 1
                else:
                    reuse_cross += 1
            else:
                new_blocks += 1
            block_last_touch[h] = ts
            block_session[h] = sid

        tok_per_block = r["input_length"] / n_blocks if n_blocks > 0 else 0
        bin_data[b]["intra_tok"] += reuse_intra * tok_per_block
        bin_data[b]["cross_tok"] += reuse_cross * tok_per_block
        bin_data[b]["new_tok"] += new_blocks * tok_per_block
        bin_data[b]["total_tok"] += r["input_length"]

    with (outdir / "f2a_reuse_timeline_qwenapp.csv").open("w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["bin_start_s", "n_req", "qps", "intra_tok_s",
                     "cross_tok_s", "new_tok_s", "reuse_tok_s", "total_tok_s", "reuse_frac"])
        for i, d in enumerate(bin_data):
            t0 = i * bin_s
            qps = d["n_req"] / bin_s
            reuse = d["intra_tok"] + d["cross_tok"]
            total = d["total_tok"]
            frac = reuse / total if total > 0 else 0
            w.writerow([
                f"{t0:.1f}", d["n_req"], f"{qps:.3f}",
                f"{d['intra_tok']:.1f}", f"{d['cross_tok']:.1f}",
                f"{d['new_tok']:.1f}", f"{reuse:.1f}", f"{total:.1f}",
                f"{frac:.4f}",
            ])

    print(f"  reuse_timeline: {n_bins} bins")


# ---------------------------------------------------------------------------
# 5. Session skew
# ---------------------------------------------------------------------------
def build_session_skew(rows, sessions, outdir):
    session_tokens = {}
    for sid, sess_rows in sessions.items():
        session_tokens[sid] = sum(r["input_length"] for r in sess_rows)

    total = sum(session_tokens.values())
    ranked = sorted(session_tokens.values(), reverse=True)
    cumsum = np.cumsum(ranked) / total * 100.0
    n = len(ranked)
    rank_pct = np.arange(1, n + 1) / n * 100.0

    step = max(1, n // 500)
    samples = [{"rank_pct": float(rank_pct[i]), "cum_pct": float(cumsum[i])}
               for i in range(0, n, step)]
    if samples[-1]["rank_pct"] < 100.0:
        samples.append({"rank_pct": 100.0, "cum_pct": 100.0})

    anchors = {}
    for pct in [1, 5, 10, 25, 50]:
        idx = max(0, int(n * pct / 100.0) - 1)
        anchors[f"top{pct}pct"] = float(cumsum[idx])

    data = {
        "n_sessions": n,
        "total_tokens": int(total),
        "anchors_check": anchors,
        "source_trace": "trace-qwenapp-formatted/0529-1000-1200.jsonl",
        "samples": samples,
    }
    (outdir / "f2b_session_skew_cdf_qwenapp.json").write_text(
        json.dumps(data, indent=2, ensure_ascii=False))
    print(f"  session_skew: {n} sessions, top1%={anchors['top1pct']:.1f}%")


# ---------------------------------------------------------------------------
def main():
    p = argparse.ArgumentParser()
    p.add_argument("trace", type=Path)
    p.add_argument("--outdir", type=Path, default=None)
    args = p.parse_args()

    outdir = args.outdir or Path(__file__).resolve().parent
    outdir.mkdir(parents=True, exist_ok=True)

    print(f"Reading {args.trace} ...")
    rows = []
    sessions: dict[str, list[dict]] = defaultdict(list)
    with args.trace.open() as f:
        for line in f:
            r = json.loads(line)
            rows.append(r)
            sessions[r["session_id"]].append(r)

    for sess_rows in sessions.values():
        sess_rows.sort(key=lambda r: r["turn"])

    print(f"Loaded {len(rows)} rows, {len(sessions)} sessions")

    build_token_length(rows, sessions, outdir)
    build_block_lifetime(rows, outdir)
    build_reuse_gap(rows, sessions, outdir)
    build_reuse_timeline(rows, sessions, outdir)
    build_session_skew(rows, sessions, outdir)

    print("Done.")


if __name__ == "__main__":
    main()