#!/usr/bin/env python3
"""is_my_astra_fake.py — does your Codex account get the real gpt-6-astra?

Runs a fixed prompt through `codex exec` with your current login, in rounds of 3 parallel runs
(up to 3 rounds, 1–4 minutes, a few thousand tokens per run), and reads Codex's own session log to measure
how fast the answer text streams (after any reasoning, so network and startup delays don't count).
Real gpt-6-astra streams at a very steady ~33.4 tok/s; degraded accounts are served a different model, in bursts,
that streams clearly faster or slower. YES as soon as a round has 2+ anomalous runs.

Usage:  python3 is_my_astra_fake.py        (needs `codex` on PATH; Python 3.8+, stdlib only)
Exit code: 0 = NO (real astra), 1 = YES (degraded), 2 = UNKNOWN (could not measure)
"""
import base64, datetime, glob, json, math, os, shutil, subprocess, sys, tempfile, time
from concurrent.futures import ThreadPoolExecutor

RUNS = 3      # parallel runs per round
ROUNDS = 3    # stop early on a YES
PAUSE = 45    # seconds between rounds (degraded service comes in bursts)
NORMAL_STREAM = (30.5, 36.5)  # tok/s band for real gpt-6-astra text streaming (observed 32.5–33.8)
PROMPT = ("Create an SVG illustration of a pelican standing near water. Respond with only the SVG markup "
          "(roughly 60 elements) and nothing else. Don't use tools or write files.")
CODEX_HOME = os.path.expanduser(os.environ.get("CODEX_HOME", "~/.codex"))


def account():
    try:
        tok = json.load(open(os.path.join(CODEX_HOME, "auth.json")))["tokens"]["id_token"].split(".")[1]
        claims = json.loads(base64.urlsafe_b64decode(tok + "=" * (-len(tok) % 4)))
        plan = claims.get("https://api.openai.com/auth", {}).get("chatgpt_plan_type")
        return f"{claims.get('email')} ({plan})"
    except Exception:
        return "unknown (no ChatGPT login found)"


def run_once(workdir):
    """Run one benchmark prompt; returns the last lines of codex's stderr if it failed."""
    try:
        p = subprocess.run(["codex", "exec", "--skip-git-repo-check", "-m", "gpt-6-astra", "-c",
                            "model_reasoning_effort=medium", "--sandbox", "read-only", PROMPT], cwd=workdir,
                           stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE,
                           text=True, timeout=300)
        return None if p.returncode == 0 else (p.stderr.strip().splitlines() or ["exit %d" % p.returncode])[-1]
    except subprocess.TimeoutExpired:
        return "timed out after 5 minutes"


def ts(s):
    return datetime.datetime.fromisoformat(s.replace("Z", "+00:00")).timestamp()


def measure(workdir, since):
    """Find this run's session log (by working dir) and return the main response's stats."""
    for f in sorted(glob.glob(os.path.join(CODEX_HOME, "sessions", "*", "*", "*", "rollout-*.jsonl")),
                    key=os.path.getmtime, reverse=True):
        if os.path.getmtime(f) < since:
            break
        with open(f) as fh:
            try:
                meta = json.loads(fh.readline())["payload"]
            except (ValueError, KeyError):
                continue
            if os.path.realpath(meta.get("cwd", "")) != os.path.realpath(workdir):
                continue
            last, best, model, saw_record, think_done = None, None, None, False, None
            for line in fh:
                try:
                    o = json.loads(line)
                except ValueError:
                    continue
                p = o.get("payload") or {}
                if o.get("type") == "turn_context":
                    model = p.get("model", model)
                elif o.get("type") == "response_item" and p.get("type") == "reasoning" and last:
                    think_done = ts(o["timestamp"])
                elif o.get("type") == "response_item" and (p.get("role") == "user" or str(p.get("type", "")).endswith("_output")):
                    last, think_done = ts(o["timestamp"]), None
                elif last and (o.get("type") == "token_usage_record" or (
                        not saw_record and o.get("type") == "event_msg" and p.get("type") == "token_count"
                        and (p.get("info") or {}).get("last_token_usage"))):
                    saw_record |= o.get("type") == "token_usage_record"
                    u = p["usage"] if o.get("type") == "token_usage_record" else p["info"]["last_token_usage"]
                    u = {k: u.get(k) or 0 for k in ("output_tokens", "reasoning_output_tokens", "input_tokens")}
                    end = ts(o["timestamp"])
                    if end > last and (best is None or u["output_tokens"] > best["out"]):
                        text = u["output_tokens"] - u["reasoning_output_tokens"]
                        stream = text / (end - think_done) if think_done and end > think_done and text > 300 else None
                        best = dict(out=u["output_tokens"], rsn=u["reasoning_output_tokens"], inp=u["input_tokens"],
                                    tps=u["output_tokens"] / (end - last), stream=stream, model=model)
                    last, think_done = end, None
            return best
    return None


def verdict(r):
    """normal / anomalous / inconclusive for one run."""
    if r["stream"]:  # pure text-streaming speed: immune to network, queueing and startup delays
        lo, hi = NORMAL_STREAM
        r["speed"], r["expected"] = r["stream"], "30.5-36.5"
        r["state"] = "normal" if lo <= r["stream"] <= hi else "ANOMALOUS"
    else:  # no reasoning step to time from: only an unmistakably fast run counts
        exp = 21.292 + 2.58 * math.log(r["out"]) - 0.874 * math.log(max(r["inp"], 1))
        r["speed"], r["expected"] = r["tps"], f"~{exp:.1f}"
        r["state"] = "ANOMALOUS" if r["tps"] - exp > 4 else ("inconclusive" if r["tps"] - exp < -8 else "normal")
    r["bad"] = r["state"] == "ANOMALOUS"
    return r


def run_round():
    since = time.time() - 5
    dirs = [tempfile.mkdtemp(prefix="astra-check-") for _ in range(RUNS)]
    try:
        with ThreadPoolExecutor(RUNS) as ex:
            errors = [e for e in ex.map(run_once, dirs) if e]
        results = [verdict(m) for m in (measure(d, since) for d in dirs) if m and m["out"] >= 500]
    finally:
        for d in dirs:
            shutil.rmtree(d, ignore_errors=True)
    return results, errors


def unknown(msg):
    print(f"{msg}\nIs my gpt-6-astra fake?  UNKNOWN", file=sys.stderr)
    sys.exit(2)


def main():
    if not shutil.which("codex"):
        unknown("`codex` not found on PATH.")
    print(f"Account: {account()}")
    print(f"Benchmark: up to {ROUNDS} rounds of {RUNS} parallel gpt-6-astra runs (1-4 minutes)\n")
    print(f"{'round':<7}{'tok/s':>7}{'normal':>11}{'reasoning':>11}")
    all_runs, fake = [], False
    for rnd in range(1, ROUNDS + 1):
        if rnd > 1:
            time.sleep(PAUSE)
        rs, errors = run_round()
        if errors:
            print(f"{rnd:<7}codex error ({len(errors)}/{RUNS} runs): {errors[0][:200]}", flush=True)
        if not rs and errors:
            unknown("Codex runs failed (see error above; logged in? rate-limited?).")
        for r in rs:
            print(f"{rnd:<7}{r['speed']:>7.1f}{r['expected']:>11}{r['rsn']:>11}   {r['state']}", flush=True)
        all_runs += rs
        if sum(r["bad"] for r in rs) >= 2:
            fake = True
            break
    decided = [r for r in all_runs if r["state"] != "inconclusive"]
    if len(decided) < 3 and not fake:
        unknown(f"Only {len(decided)} run(s) gave a clear measurement, not enough for a verdict.")
    n_bad = sum(r["bad"] for r in all_runs)
    print(f"\nIs my gpt-6-astra fake?  {'YES' if fake else 'NO'}  ({n_bad}/{len(all_runs)} runs anomalous)")
    if fake:
        print("This account appears to be served a different model (most likely gpt-5.6-luna) under the gpt-6-astra label.")
    elif n_bad:
        print("One anomalous run is within normal noise; re-run later if results look off.")
    sys.exit(1 if fake else 0)


if __name__ == "__main__":
    main()