"""One-shot Colab training status over the HTTP runtime proxy (no ssh).

    python3 scripts/colab_status.py --host <instance.url netloc> --token <instance.token> \\
        [--lines 4] [--dir content/swift15/checkpoints/agent_persona_lora_v1]

Reads only `/api/contents`, so it works while the ssh slot is refused; see scripts/pull_colab_http.py
for why that happens. Credentials come from the caller because they must be re-listed from
`colab.instances()` each time -- both the url and the token expire in about an hour -- and the kernel's
plugin is not importable from a subprocess.
"""
import argparse
import collections
import json
import pathlib
import re
import sys
import urllib.parse
import urllib.request

WINDOW = re.compile(
    r"Epoch (\d)/2 \[(\d+)/322\] step=(\d+).*?pref_loss=([\d.]+) margin=([+-][\d.]+)"
)
# The preference pair for SFT window `idx` of epoch `e` is
#     pair = ftpo_train[(e * len(sft_train) + idx) % len(ftpo_train)]      (trainer:2343)
# so the window->category map SHIFTS every epoch. Hardcoding one epoch's range (as this script first
# did) mislabels a later epoch's windows: with 322 SFT windows and 290 pairs the shift is 32, which
# moved 6C from w198-w290 to w167-w258 in epoch 2 and put 6A single-token pairs at the tail -- values
# that look impossible (+40..+50) for a per-token 6C mean and are the tell that the map is wrong.
SFT_WINDOWS = 322
PAIRS = pathlib.Path(__file__).resolve().parents[1] / "persona" / "data" / "agent_persona_v1" / "ftpo_dpo_pairs.jsonl"


def pair_categories():
    rows = [json.loads(line) for line in PAIRS.read_text(encoding="utf-8").splitlines() if line.strip()]
    return [r["category"] for r in rows if r.get("split") == "train"]


def category_of(cats, epoch, window):
    """The category of the preference pair scored alongside SFT window `window` of `epoch` (1-based)."""
    if not cats:
        return ""
    return cats[((epoch - 1) * SFT_WINDOWS + (window - 1)) % len(cats)]


def contents(host, token, path, text=False, content=True):
    q = urllib.parse.urlencode({
        "colab-runtime-proxy-token": token,
        **({"format": "text"} if text else {}),
        "content": "1" if content else "0",
    })
    url = f"https://{host}/api/contents/{path.lstrip('/')}?{q}"
    req = urllib.request.Request(url, headers={"X-Colab-Client-Agent": "colab-cli"})
    with urllib.request.urlopen(req, timeout=40) as r:
        return json.loads(r.read())


def census(text):
    """Per-epoch preference-margin health, and the epoch-over-epoch comparison that matters.

    A positive margin alone is not evidence of learning: `+0.02` satisfies the sign and decides
    nothing. Track magnitude, and compare the *same* window across epochs, since the window order is
    deterministic (longest first).
    """
    cats = pair_categories()
    per = collections.defaultdict(list)
    by = collections.defaultdict(dict)
    bycat = collections.defaultdict(list)
    for line in text.splitlines():
        m = WINDOW.search(line)
        if not m:
            continue
        epoch, win, _step, _pl, margin = m.groups()
        per[int(epoch)].append(float(margin))
        by[int(win)][int(epoch)] = float(margin)
        bycat[(int(epoch), category_of(cats, int(epoch), int(win)))].append(float(margin))
    for epoch in sorted(per):
        v = sorted(per[epoch])
        near = sum(1 for x in v if x < 5.0)
        print(f"  epoch {epoch}: n={len(v):3d} median {v[len(v) // 2]:+7.2f} | "
              f"margin<5 {near / len(v):4.0%} | violated {sum(1 for x in v if x <= 0)}")
        # Spend windows are contiguous because the order is deterministic (longest first):
        # 6A w1-96, 6B w97-197, 6C w198-290, then the 6A/6B tail. 6A/6B score ONE logit pair,
        # 6C scores a per-token MEAN over ~96 positions (run_colab_a100_agent_persona.py:1986), so
        # pooling them compares values about 95x apart. Report by unit, never pooled.
        for (ep, cat), vals in sorted(bycat.items()):
            if ep != epoch or not cat:
                continue
            sv = sorted(vals)
            print(f"      {cat:26s} n={len(sv):3d} median {sv[len(sv) // 2]:+7.2f} "
                  f"min {sv[0]:+7.2f} max {sv[-1]:+7.2f} violated {sum(1 for x in sv if x <= 0)}")
    # A run of consecutive near-ties is a different finding from scattered ones: isolated lows are a
    # hard tail, a contiguous block is a family of pairs the objective is not separating. Runs are
    # measured within one category, because pooling 6A/6B single-token logits with 6C per-token means
    # compares values about 95x apart.
    for epoch in sorted(per):
        wins = sorted(w for w, v in by.items() if epoch in v)
        best = run = 0
        best_at = run_at = None
        for w in wins:
            if by[w][epoch] < 5.0 and category_of(cats, epoch, w).startswith(("6A", "6B")):
                if run == 0:
                    run_at = w
                run += 1
                if run > best:
                    best, best_at = run, run_at
            else:
                run = 0
        if best >= 3:
            print(f"      longest run of margin<5: {best} consecutive "
                  f"(w{best_at}-w{best_at + best - 1})")
    both = [(w, v[1], v[2]) for w, v in sorted(by.items())
            if 1 in v and 2 in v and category_of(cats, 1, w) == category_of(cats, 2, w)]
    if not both:
        print("  no window seen in both epochs yet")
        return
    drops = [(w, a, b) for w, a, b in both if b - a < -5.0]
    stretch = [(w, a, b) for w, a, b in both if category_of(cats, 2, w).startswith("6C")]
    print(f"  windows in both epochs: {len(both)} | margin fell >5: {len(drops)}"
          + (f" -> {[(f'w{w}', f'{a:+.1f}', f'{b:+.1f}') for w, a, b in drops]}" if drops else ""))
    if stretch:
        print(f"  collapsed stretch w{lo}-w{hi} revisited: {len(stretch)} window(s) -> "
              + ", ".join(f"w{w} {a:+.2f}->{b:+.2f}" for w, a, b in stretch))
    elif not stretch:
        e2_at = max(w for w, v in by.items() if 2 in v)
        print(f"  6C windows not yet revisited (epoch 2 is at w{e2_at}): "
              "answers whether the anti-slop DPO regime actually moved those pairs")


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--host", required=True)
    ap.add_argument("--token", required=True)
    ap.add_argument("--lines", type=int, default=4)
    ap.add_argument("--dir", default="content/swift15/checkpoints/agent_persona_lora_v1")
    ap.add_argument("--census", action="store_true", help="preference-margin health per epoch")
    a = ap.parse_args(argv)
    log = contents(a.host, a.token, f"{a.dir}/train.log", text=True)
    lines = (log.get("content") or "").splitlines()
    print(f"train.log {log.get('size', 0):,} B / {len(lines)} lines")
    for line in lines[-a.lines:]:
        print("  " + line)
    if a.census:
        census("\n".join(lines))
    for item in sorted((contents(a.host, a.token, a.dir).get("content") or []), key=lambda x: x["name"]):
        if item["name"].endswith((".pt", ".bin", ".pid")):
            print(f"  {item['name']:24s} {item.get('size', 0):>12,}  {item['last_modified'][11:19]}")
    return 0


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