#!/usr/bin/env python3
"""Pick a load-balanced cswap account for an Omnigent session.

Sticky per session (OMNIGENT_SESSION_ID else cwd) so relaunches reuse the same account-keyed
cswap-run profile. Among healthy accounts in this machine's pool, pick the emptiest.
Multiple sessions may share an account — cswap-run's account-keyed dir + Claude's refresh lock
coordinate their token refreshes, so that's safe on ONE machine. Cross-machine isolation is the
job of the per-machine pool allowlist (~/.omni-cswap/pool, one account number per line): give
each machine disjoint accounts so two machines never refresh the same OAuth grant.

Prints the chosen account number, or nothing (shim then uses the default login)."""
import json, os, subprocess, sys, time
from pathlib import Path

SELF = Path(__file__).resolve().parent
STATE = SELF / "assignments.json"
LOCK = SELF / ".assignments.lock"
POOL_FILE = SELF / "pool"
THRESHOLD = float(os.environ.get("OMNI_CSWAP_THRESHOLD", "90"))
# Accounts to exclude for this pick (comma-separated) — used by the rescue watchdog when an
# account hit a cap the usage numbers don't show (e.g. a scoped per-model limit).
EXCLUDE = {x.strip() for x in os.environ.get("OMNI_CSWAP_EXCLUDE", "").split(",") if x.strip()}
# Model-aware pick: when set (e.g. "Fable"), also drop accounts whose scoped usage for that
# model is >= THRESHOLD, so a model-capped session re-picks an account with headroom for it.
MODEL = os.environ.get("OMNI_CSWAP_MODEL", "").strip()
STICKY_TTL = 6 * 3600

# Map a claude --model id to the cswap scoped-usage bucket name (only Fable is scoped today).
def _scoped_name(model):
    m = (model or "").lower()
    if "fable" in m: return "Fable"
    return None

def _key():
    return os.environ.get("OMNIGENT_SESSION_ID") or f"cwd:{os.path.realpath(os.getcwd())}"

def _pool():
    if POOL_FILE.exists():
        return {l.strip() for l in POOL_FILE.read_text().split() if l.strip()}
    return None  # None = all accounts allowed

def _accounts():
    out = subprocess.run(["cswap", "list", "--json"], capture_output=True, text=True, timeout=30).stdout
    data = json.loads(out)
    default = str(data.get("activeAccountNumber", ""))
    scoped_name = _scoped_name(MODEL)
    res = {}
    for a in data.get("accounts", []):
        if a.get("usageStatus") != "ok" or a.get("disabled") is True:
            continue
        u = a.get("usage") or {}
        five = (u.get("fiveHour") or {}).get("pct"); seven = (u.get("sevenDay") or {}).get("pct")
        load = max(v for v in (five, seven, 0) if v is not None)
        if scoped_name:  # fold the model's scoped usage into the account's load
            for s in (u.get("scoped") or []):
                if s.get("name") == scoped_name and s.get("pct") is not None:
                    load = max(load, s["pct"])
        res[str(a["number"])] = load
    return res, default

def main():
    import fcntl
    key = _key(); now = time.time()
    with open(LOCK, "w") as lf:
        fcntl.flock(lf, fcntl.LOCK_EX)
        try: st = json.loads(STATE.read_text())
        except Exception: st = {}
        st = {k: v for k, v in st.items() if now - v.get("ts", 0) < STICKY_TTL or k == key}

        acc, default = _accounts()
        pool = _pool()
        if pool is not None:
            acc = {n: l for n, l in acc.items() if n in pool}
        acc = {n: l for n, l in acc.items() if n not in EXCLUDE}
        elig = {n: l for n, l in acc.items() if l < THRESHOLD}

        prev = st.get(key, {}).get("account")
        if prev in elig:
            st[key] = {"account": prev, "ts": now}; STATE.write_text(json.dumps(st, indent=1))
            print(prev); return
        if not elig:
            sys.exit(1)
        # Prefer NON-DEFAULT accounts: cswap-auto rotates the default, and a pinned session
        # sharing the rotation target's grant risks a refresh collision. Only fall back to the
        # default account when it's the sole eligible one.
        non_default = {n: l for n, l in elig.items() if n != default}
        cands = non_default or elig
        best = min(cands, key=cands.get)
        st[key] = {"account": best, "ts": now}; STATE.write_text(json.dumps(st, indent=1))
        print(best)

if __name__ == "__main__":
    try:
        main()
    except SystemExit:
        raise
    except Exception as e:
        sys.stderr.write(f"pick_account error: {e}\n"); sys.exit(1)