#!/usr/bin/env python3
"""Grok token broker — hold a refreshable xAI session credential WITHOUT touching the CLI's.

⭐ THE PROBLEM. The Grok Build coding plan (`https://cli-chat-proxy.grok.com/v1`, serving
grok-4.6) authenticates with an OIDC **session** token, not an API key: `api_key: null`,
`env_key: null` in the CLI's own model cache. It expires roughly every six hours. LiteLLM and
every other gateway on this fleet hold a STATIC credential, so a naive alias works until the
session lapses and then fails closed.

⚠⚠ THE HAZARD THAT SHAPES THIS ENTIRE DESIGN, AND WHY IT IS NOT A SIMPLE REFRESH LOOP.
`https://auth.x.ai/oauth2/token` supports the `refresh_token` grant. But **a refresh may ROTATE
the refresh token**, and many OIDC providers invalidate the old one server-side the instant a
new one is issued. The operator's interactive `grok` CLI holds its refresh token in
`~/.grok/auth.json`. If xAI rotates-and-invalidates, then a broker refreshing against the SAME
credential silently kills the CLI login — and it does so **even though the broker never writes
to that file.** Not-writing is necessary and NOT sufficient; the damage would be server-side.

heid's `groa_http_dispatch.py` refuses to refresh at all for this reason, which is the correct
call absent an answer. This broker's job is to get the answer first and then act on it.

⛔ SO THE REFRESH LOOP IS NOT ARMED UNTIL `probe-rotation` HAS RUN AND SAID IT IS SAFE.
`serve` and `refresh` both refuse without a stored probe verdict. That is deliberate: this is a
tool whose whole value is contingent on one unknown, and running it before measuring the unknown
is how you find out by breaking the operator's login.

## The probe, and what each outcome means

`probe-rotation` performs ONE refresh using a COPY of the credential, then re-presents the
ORIGINAL refresh token:

  original still valid   -> NON-ROTATING (or rotation without invalidation). The broker can
                            coexist with the CLI on one credential. Loop may be armed.
  original now rejected  -> ROTATING-AND-INVALIDATING. The broker MUST NOT share the CLI's
                            credential. It needs its own login, or the idea is dead. Refuses
                            to arm, and says so.

⚠ The probe itself consumes one refresh. If the answer is "rotating", the CLI's token is
already dead by the time you read the verdict — there is no way to ask this question without
spending the credential, and pretending otherwise would be worse. The remedy is a `grok` CLI
re-login, which is the same remedy heid's exit-3 path already names. **Run it when a broken CLI
login is a two-minute annoyance, not mid-panel.**

## What it never does

- Never writes `~/.grok/auth.json`. Reads it once, on `seed`, and never again.
- Never prints a token to stdout except under `token --raw`, which is for a consumer to capture.
- Never logs a token value. Expiry, scope and subject only.
"""
from __future__ import annotations

import argparse
import base64
import json
import os
import pathlib
import sys
import time
import urllib.error
import urllib.parse
import urllib.request

CLI_AUTH = pathlib.Path.home() / ".grok" / "auth.json"
STORE = pathlib.Path(os.environ.get("GROK_BROKER_HOME",
                                    pathlib.Path.home() / ".config" / "grok-token-broker"))
CRED = STORE / "credential.json"
VERDICT = STORE / "rotation-verdict.json"
TOKEN_URL = "https://auth.x.ai/oauth2/token"
API_BASE = "https://cli-chat-proxy.grok.com/v1"
REFRESH_MARGIN_S = 1800          # refresh when under 30 min remain


def _claims(jwt: str) -> dict:
    pl = jwt.split(".")[1]
    pl += "=" * (-len(pl) % 4)
    return json.loads(base64.urlsafe_b64decode(pl))


def _remaining(jwt: str) -> int:
    try:
        return int(_claims(jwt)["exp"] - time.time())
    except Exception:
        return -1


def _post_form(url: str, fields: dict, timeout: int = 30):
    data = urllib.parse.urlencode(fields).encode()
    req = urllib.request.Request(url, data=data,
                                 headers={"Content-Type": "application/x-www-form-urlencoded"})
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            return r.status, json.loads(r.read())
    except urllib.error.HTTPError as e:
        body = e.read().decode("utf-8", "replace")[:400]
        return e.code, {"_error": body}
    except Exception as e:                                   # noqa: BLE001
        return 0, {"_error": f"{type(e).__name__}: {e}"}


def _read_cli_credential() -> dict:
    """Pull access token, refresh token and client id out of the CLI's auth file.

    Shape-tolerant on purpose: this file belongs to another tool and its layout is not ours to
    depend on. Anything not found is reported, never guessed.
    """
    if not CLI_AUTH.exists():
        raise SystemExit(f"REFUSING: no CLI credential at {CLI_AUTH}")
    blob = json.loads(CLI_AUTH.read_text(encoding="utf-8"))
    found: dict = {}

    def walk(o, key_hint=""):
        if isinstance(o, dict):
            for k, v in o.items():
                if isinstance(v, str):
                    if k == "key" and v.count(".") == 2 and "access" not in found:
                        found["access"] = v
                    elif "refresh" in k.lower() and len(v) > 20:
                        found.setdefault("refresh", v)
                walk(v, k)
            if "auth.x.ai::" in key_hint:
                found.setdefault("client_id", key_hint.split("::", 1)[1])
        elif isinstance(o, list):
            for v in o:
                walk(v, key_hint)

    for k, v in blob.items():
        walk(v, k)
    missing = [f for f in ("access", "refresh") if f not in found]
    if missing:
        raise SystemExit(f"REFUSING: could not locate {missing} in {CLI_AUTH}. "
                         "Its layout changed; inspect it by hand rather than letting this guess.")
    return found


def cmd_seed(_a) -> int:
    STORE.mkdir(parents=True, exist_ok=True)
    os.chmod(STORE, 0o700)
    c = _read_cli_credential()
    CRED.write_text(json.dumps({**c, "seeded_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
                                "source": str(CLI_AUTH)}, indent=2), encoding="utf-8")
    os.chmod(CRED, 0o600)
    cl = _claims(c["access"])
    print(f"  seeded -> {CRED} (0600)")
    print(f"  scope   : {cl.get('scope')}")
    print(f"  expires : {time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime(cl['exp']))} "
          f"({_remaining(c['access']) // 60} min left)")
    print(f"  client  : {c.get('client_id', '<not found>')}")
    print("\n  ⛔ The refresh loop is NOT armed. Run `probe-rotation` next, and read its warning "
          "first — the probe spends one refresh and may end the CLI's session.")
    return 0


def cmd_probe(a) -> int:
    if not CRED.exists():
        raise SystemExit("REFUSING: run `seed` first")
    c = json.loads(CRED.read_text(encoding="utf-8"))
    if not a.i_accept_this_may_end_the_cli_session:
        print("REFUSING without --i-accept-this-may-end-the-cli-session.\n\n"
              "  This probe performs a real refresh. If xAI rotates AND invalidates, the\n"
              "  operator's `grok` CLI login dies at that moment and needs an interactive\n"
              "  re-login. There is no way to answer the question without spending the\n"
              "  credential. Run it when that is a two-minute annoyance, not mid-panel.",
              file=sys.stderr)
        return 2

    old_refresh = c["refresh"]
    fields = {"grant_type": "refresh_token", "refresh_token": old_refresh}
    if c.get("client_id"):
        fields["client_id"] = c["client_id"]

    print("  [1/2] refreshing with the stored refresh token ...")
    st, body = _post_form(TOKEN_URL, fields)
    if st != 200:
        print(f"  refresh FAILED http={st}: {str(body.get('_error'))[:200]}", file=sys.stderr)
        print("  Verdict NOT recorded. The loop stays disarmed.", file=sys.stderr)
        return 4
    new_refresh = body.get("refresh_token")
    rotated = bool(new_refresh and new_refresh != old_refresh)
    print(f"      ok. refresh_token {'ROTATED' if rotated else 'unchanged'}")

    if not rotated:
        verdict, safe = "non-rotating", True
        note = "Provider returned the same refresh token; the CLI's copy is untouched."
    else:
        print("  [2/2] re-presenting the ORIGINAL refresh token to see if it still works ...")
        st2, _ = _post_form(TOKEN_URL, fields)
        if st2 == 200:
            verdict, safe = "rotating-old-still-valid", True
            note = ("Rotation happens but the old refresh token still works, so the CLI's copy "
                    "survives. Coexistence is safe.")
        else:
            verdict, safe = "rotating-and-invalidating", False
            note = (f"The original refresh token is now rejected (http={st2}). The CLI's login "
                    "is ALREADY DEAD and needs an interactive `grok` re-login. A broker MUST "
                    "NOT share this credential — it needs its own login.")
        print(f"      original token -> http={st2}")

    c["access"] = body.get("access_token", c["access"])
    if new_refresh:
        c["refresh"] = new_refresh
    CRED.write_text(json.dumps(c, indent=2), encoding="utf-8")
    os.chmod(CRED, 0o600)

    VERDICT.write_text(json.dumps({
        "verdict": verdict, "safe_to_arm": safe, "note": note,
        "probed_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
        "token_url": TOKEN_URL, "api_base": API_BASE,
    }, indent=2), encoding="utf-8")
    os.chmod(VERDICT, 0o600)
    print(f"\n  VERDICT: {verdict}  (safe_to_arm={safe})\n  {note}")
    if not safe:
        print("\n  ⚠ TELL THE OPERATOR: the `grok` CLI on this box needs a re-login.", file=sys.stderr)
    return 0 if safe else 5


def _ensure_armed() -> dict:
    if not VERDICT.exists():
        raise SystemExit("REFUSING: no rotation verdict. Run `probe-rotation` first — this tool "
                         "will not refresh a credential whose rotation behaviour is unmeasured.")
    v = json.loads(VERDICT.read_text(encoding="utf-8"))
    if not v.get("safe_to_arm"):
        raise SystemExit(f"REFUSING: verdict is '{v['verdict']}'. {v['note']}")
    return v


def cmd_refresh(_a) -> int:
    _ensure_armed()
    c = json.loads(CRED.read_text(encoding="utf-8"))
    fields = {"grant_type": "refresh_token", "refresh_token": c["refresh"]}
    if c.get("client_id"):
        fields["client_id"] = c["client_id"]
    st, body = _post_form(TOKEN_URL, fields)
    if st != 200:
        print(f"refresh failed http={st}: {str(body.get('_error'))[:200]}", file=sys.stderr)
        return 4
    c["access"] = body["access_token"]
    if body.get("refresh_token"):
        c["refresh"] = body["refresh_token"]
    c["refreshed_at"] = time.strftime("%Y-%m-%dT%H:%M:%S%z")
    CRED.write_text(json.dumps(c, indent=2), encoding="utf-8")
    os.chmod(CRED, 0o600)
    print(f"  refreshed; {_remaining(c['access']) // 60} min remaining")
    return 0


def cmd_token(a) -> int:
    """Print a currently-valid access token, refreshing first if it is close to expiry."""
    if not CRED.exists():
        raise SystemExit("REFUSING: run `seed` first")
    c = json.loads(CRED.read_text(encoding="utf-8"))
    if _remaining(c["access"]) < REFRESH_MARGIN_S:
        _ensure_armed()
        if cmd_refresh(a) != 0:
            return 4
        c = json.loads(CRED.read_text(encoding="utf-8"))
    if a.raw:
        sys.stdout.write(c["access"])
    else:
        print(f"  valid for {_remaining(c['access']) // 60} more minutes "
              f"(use --raw to emit the token itself)")
    return 0


def cmd_status(_a) -> int:
    print(f"  store   : {STORE}")
    if not CRED.exists():
        print("  state   : NOT SEEDED")
        return 0
    c = json.loads(CRED.read_text(encoding="utf-8"))
    print(f"  seeded  : {c.get('seeded_at')} from {c.get('source')}")
    print(f"  access  : {_remaining(c['access']) // 60} min remaining")
    if VERDICT.exists():
        v = json.loads(VERDICT.read_text(encoding="utf-8"))
        print(f"  rotation: {v['verdict']}  safe_to_arm={v['safe_to_arm']}  ({v['probed_at']})")
        print(f"  armed   : {'YES' if v['safe_to_arm'] else 'NO — ' + v['verdict']}")
    else:
        print("  rotation: UNMEASURED — loop disarmed, `token` will refuse to refresh")
    print(f"  api_base: {API_BASE}")
    return 0


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    sub = ap.add_subparsers(dest="cmd", required=True)
    sub.add_parser("seed", help="copy the CLI credential into the broker's own store")
    p = sub.add_parser("probe-rotation", help="measure whether refresh rotates AND invalidates")
    p.add_argument("--i-accept-this-may-end-the-cli-session", action="store_true")
    sub.add_parser("refresh", help="refresh now (requires a safe verdict)")
    t = sub.add_parser("token", help="emit a valid access token, refreshing if near expiry")
    t.add_argument("--raw", action="store_true")
    sub.add_parser("status", help="report state without touching anything")
    a = ap.parse_args()
    return {"seed": cmd_seed, "probe-rotation": cmd_probe, "refresh": cmd_refresh,
            "token": cmd_token, "status": cmd_status}[a.cmd](a)


if __name__ == "__main__":
    raise SystemExit(main())