#!/usr/bin/env python3
"""Session-sticky routing imbalance simulation over a formatted ali-trace.

Routing policy:
  - turn==1 (parent_chat_id == -1): dispatch to the instance with the FEWEST
    total requests received in the trailing LOAD_WIN seconds (least-loaded).
  - turn>1: pinned to the same instance as the request's parent turn (sticky).

For every tumbling WINDOW-second time window we count the total requests landed
on each of the N instances and report imbalance = max / mean over all N
instances (idle instances counted in the mean). A CDF of the per-window
imbalance values is plotted.

Run on dash0 (trace lives there). Trace is timestamp-sorted and every parent
precedes its child, so a single in-order pass suffices.
"""
import argparse
import json
from collections import deque

import numpy as np
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt


def simulate(trace, n_inst, window, load_win, weight="req", route_by="req"):
    # `route_by` is the trailing-load_win quantity the turn-1 least-loaded
    # routing balances on; `weight` is what each window accumulates for the
    # imbalance statistic. Both take values:
    #   "req"    -> 1 per request (request count)
    #   "tokens" -> input_length + output_length per request (total tokens)
    assign = {}                       # chat_id -> instance
    cnt = np.zeros(n_inst, dtype=np.int64)   # trailing-load_win route load per inst
    dq = deque()                      # (ts, inst, rw) pending expiry from the window
    wincounts = {}                    # window_idx -> np.array(n_inst)

    with open(trace) as fh:
        for line in fh:
            d = json.loads(line)
            t = d["timestamp"]
            tok = d["input_length"] + d["output_length"]

            # expire dispatches older than the trailing load window
            thr = t - load_win
            while dq and dq[0][0] < thr:
                _, old, ow = dq.popleft()
                cnt[old] -= ow

            if d["parent_chat_id"] == -1:          # turn 1: least-loaded
                inst = int(cnt.argmin())
            else:                                  # turn>1: sticky to parent
                inst = assign[d["parent_chat_id"]]
            assign[d["chat_id"]] = inst            # chain so deeper turns follow

            rw = (tok if route_by == "tokens"
                  else d["input_length"] if route_by == "input"
                  else 1)
            dq.append((t, inst, rw))
            cnt[inst] += rw

            w = int(t // window)
            arr = wincounts.get(w)
            if arr is None:
                arr = np.zeros(n_inst, dtype=np.int64)
                wincounts[w] = arr
            arr[inst] += tok if weight == "tokens" else 1

    ws = sorted(wincounts)
    imb = np.empty(len(ws))
    for i, w in enumerate(ws):
        arr = wincounts[w]
        mean = arr.sum() / n_inst
        imb[i] = arr.max() / mean if mean > 0 else 0.0
    return imb, [wincounts[w] for w in ws]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--trace", required=True)
    ap.add_argument("--n-inst", type=int, default=380)
    ap.add_argument("--window", type=float, default=60.0)
    ap.add_argument("--load-win", type=float, default=60.0)
    ap.add_argument("--weight", choices=["req", "tokens"], default="req",
                    help="per-window quantity: request count or total tokens")
    ap.add_argument("--route-by", choices=["req", "tokens", "input"],
                    default="req",
                    help="turn-1 least-loaded routing balances on this quantity "
                         "(input = arrival-known input_length only)")
    ap.add_argument("--out-prefix", required=True)
    args = ap.parse_args()

    imb, _ = simulate(args.trace, args.n_inst, args.window, args.load_win,
                      args.weight, args.route_by)

    pcts = {p: float(np.percentile(imb, p)) for p in (50, 90, 95, 99)}
    stats = {
        "trace": args.trace,
        "n_inst": args.n_inst,
        "window_s": args.window,
        "load_win_s": args.load_win,
        "weight": args.weight,
        "route_by": args.route_by,
        "n_windows": len(imb),
        "imbalance_min": float(imb.min()),
        "imbalance_mean": float(imb.mean()),
        "imbalance_p50": pcts[50],
        "imbalance_p90": pcts[90],
        "imbalance_p95": pcts[95],
        "imbalance_p99": pcts[99],
        "imbalance_max": float(imb.max()),
    }
    with open(args.out_prefix + "_stats.json", "w") as f:
        json.dump(stats, f, indent=2)
    np.save(args.out_prefix + "_imbalance.npy", imb)

    # CDF
    x = np.sort(imb)
    y = np.arange(1, len(x) + 1) / len(x)
    fig, ax = plt.subplots(figsize=(6, 4))
    ax.plot(x, y, marker=".", ms=4, lw=1.4, color="#1f77b4")
    ax.axvline(1.0, color="gray", ls="--", lw=1, label="perfect balance (1.0)")
    ax.axvline(pcts[50], color="#2ca02c", ls=":", lw=1,
               label=f"p50={pcts[50]:.2f}")
    ax.axvline(pcts[90], color="#d62728", ls=":", lw=1,
               label=f"p90={pcts[90]:.2f}")
    qty = "total tokens" if args.weight == "tokens" else "num_req"
    ax.set_xlabel("per-window %s imbalance (max / mean over %d instances)"
                  % (qty, args.n_inst))
    ax.set_ylabel("CDF over %d one-minute windows" % len(imb))
    routeq = {"tokens": "tokens", "input": "input-only"}.get(
        args.route_by, "num_req")
    ax.set_title("Session-sticky imbalance: measure=%s, route-by=%s"
                 % (qty, routeq))
    ax.grid(True, alpha=0.3)
    ax.legend(fontsize=8)
    fig.tight_layout()
    fig.savefig(args.out_prefix + "_cdf.png", dpi=140)
    print(json.dumps(stats, indent=2))


if __name__ == "__main__":
    main()