"""Shared codex account-auth audit — the ONE place that decides whether a codex
account's login still works, so `expired`, `relogin`, `list` and `status` can never
disagree with the shim's selection rule (bin/codex: expired_marked / auth_dead).

Codex auth material is a single $CODEX_HOME/auth.json written by `codex login`:
  {"auth_mode": "chatgpt", "OPENAI_API_KEY": null,
   "tokens": {"id_token": JWT, "access_token": JWT, "refresh_token": str,
              "account_id": str},
   "last_refresh": ISO}
The access token is a JWT whose `exp` claim is readable offline; the refresh token
is opaque (no readable expiry), so a dead refresh grant is only ever known from a
marker written by a failed refresh/verify — exactly like the shim assumes.

States:
  ok       usable right now (live access token, or expired-but-refreshable)
  expired  has auth material that CANNOT authenticate — needs `codex-accounts relogin`
  blocked  authenticates, but Codex access is disabled for it by a workspace admin —
           a re-login usually re-issues the grant; if it stays blocked, an admin must act
  missing  no auth on this machine, and this machine is supposed to own the grant
  remote   no auth here, but the manifest says another machine owns it (informational)

Only 'ok' accounts are selectable; everything else is excluded by the shim.

Run directly for a TSV dump:  python3 lib/codex_audit.py <acc-root> [mac|linux]
Columns: id, email, home, state, label, reason, fix
"""

import base64
import json
import os
import re
import sys
import time

VALID_ID = re.compile(r'acct-\d{2}')


def _mtime(path):
    try:
        return os.path.getmtime(path)
    except OSError:
        return 0.0


def _nonempty(path):
    try:
        return os.path.getsize(path) > 0
    except OSError:
        return False


def marker_reason(mpath):
    """Second line of a marker file, as a human string ('' when absent/unreadable)."""
    try:
        lines = open(mpath, errors='replace').read().splitlines()
    except OSError:
        return ''
    return lines[1].strip() if len(lines) > 1 else ''


def marker_slug(detail):
    """The reason=<slug> field of a marker line ('' when absent)."""
    m = re.search(r'reason=([A-Za-z0-9._-]+)', detail or '')
    return m.group(1) if m else ''


def expired_marked(d, now=None):
    """True when <d>/.expired is still in force. Mirrors the shim exactly:
    a CREDENTIAL-scoped park clears as soon as a newer auth.json lands (re-login, or a
    refresh by another process); a POLICY park (org-blocked) does not, because a token
    refresh says nothing about whether the workspace re-enabled Codex. Either kind also
    expires at its own soft_until when the shim wrote it from a single failed run."""
    now = time.time() if now is None else now
    mpath = os.path.join(d, '.expired')
    if not os.path.isfile(mpath):
        return False
    detail = marker_reason(mpath)
    if marker_slug(detail) != 'org-blocked':
        mt = _mtime(mpath)
        p = os.path.join(d, 'auth.json')
        if os.path.isfile(p) and _mtime(p) > mt:
            return False
    m = re.search(r'soft_until=(\d+)', detail)
    if m and int(m.group(1)) <= now:
        return False
    return True


def jwt_claims(token):
    """Decode a JWT payload without verification (identity/expiry live client-side
    anyway). Returns {} for anything that is not a well-formed JWT."""
    try:
        payload = str(token).split('.')[1]
        payload += '=' * (-len(payload) % 4)
        claims = json.loads(base64.urlsafe_b64decode(payload))
        return claims if isinstance(claims, dict) else {}
    except Exception:
        return {}


def auth_email(cpath):
    """The signed-in identity, read offline from the id_token ('' when unknown)."""
    try:
        doc = json.load(open(cpath))
        tokens = doc.get('tokens') or {}
        return str(jwt_claims(tokens.get('id_token')).get('email') or '')
    except Exception:
        return ''


def creds_state(cpath, now):
    """(state, reason) for an on-disk auth.json."""
    try:
        doc = json.load(open(cpath))
        if not isinstance(doc, dict):
            raise ValueError('auth.json is not an object')
    except Exception as e:
        # Unparseable (truncated / mid-write): the shim still selects this account
        # (it only checks for a non-empty access_token string), so fall back to the
        # same lenient rule — a scrape-able access token keeps it in the pool.
        try:
            txt = open(cpath, errors='replace').read()
        except OSError:
            return 'expired', f'auth.json unreadable ({str(e)[:60]})'
        if re.search(r'"access_token"\s*:\s*"[^"]', txt):
            return 'ok', 'auth.json unreadable but carries a token — left in the pool'
        return 'expired', f'auth.json unreadable ({str(e)[:60]})'
    tokens = doc.get('tokens')
    if not isinstance(tokens, dict) or not tokens.get('access_token'):
        if doc.get('OPENAI_API_KEY'):
            # Subscription-only by design (same rule as the claude pool: no API keys).
            return 'expired', ('signed in with an API key — not supported; '
                               'use the ChatGPT subscription login')
        return 'expired', 'auth.json has no ChatGPT tokens'
    exp = jwt_claims(tokens['access_token']).get('exp')
    try:
        exp = float(exp)
    except (TypeError, ValueError):
        exp = 0.0
    if exp > now:
        return 'ok', 'chatgpt access token valid'
    if not tokens.get('refresh_token'):
        return 'expired', 'access token expired and there is no refresh token'
    # The refresh token is opaque (no readable expiry): assume refreshable. A dead
    # grant is detected by the limits refresher / verify, which write `.expired`.
    return 'ok', 'access token stale but auto-refreshes'


LABELS = {'ok': 'OK', 'expired': 'EXPIRED', 'blocked': 'BLOCKED',
          'missing': 'NO LOGIN', 'remote': 'ELSEWHERE'}

# `import --home` speaks mac|server; machine_kind() speaks mac|linux. Same two boxes,
# two vocabularies — normalize, or an un-authenticated account on the very machine that
# owns it would be filed as "lives elsewhere" and silently drop off the worklist.
_MACHINE_ALIASES = {'mac': 'mac', 'macos': 'mac', 'darwin': 'mac', 'osx': 'mac',
                    'server': 'linux', 'linux': 'linux', 'ubuntu': 'linux',
                    'debian': 'linux', 'remote': 'linux'}


def _norm_machine(name):
    return _MACHINE_ALIASES.get(str(name or '').strip().lower(), '')


def _elsewhere(home, machine):
    """True only when the manifest names a DIFFERENT machine as the grant's owner.
    An unknown/blank home is never treated as 'elsewhere' — that would hide a real gap."""
    h, m = _norm_machine(home), _norm_machine(machine)
    return bool(h) and bool(m) and h != m


def _fix_for(state, aid):
    if state == 'blocked':
        return (f'codex-accounts relogin {aid}   '
                f'(if it stays BLOCKED, a workspace admin must enable Codex for it)')
    if state in ('expired', 'missing'):
        return f'codex-accounts relogin {aid}'
    return ''


def audit_account(root, acct, now=None, machine=None):
    now = time.time() if now is None else now
    aid = acct.get('id', '')
    d = os.path.join(root, aid)
    cpath = os.path.join(d, 'auth.json')
    has_creds = _nonempty(cpath)
    home = acct.get('home', '?')
    row = {'id': aid, 'email': acct.get('email', ''), 'home': home,
           'state': 'ok', 'reason': ''}

    def done(r):
        r['label'] = LABELS.get(r['state'], r['state'].upper())
        r['fix'] = _fix_for(r['state'], aid)
        return r

    if expired_marked(d, now=now):
        detail = marker_reason(os.path.join(d, '.expired'))
        slug = marker_slug(detail)
        if slug == 'org-blocked':
            row['state'] = 'blocked'
            row['reason'] = ('Codex access is disabled for this account by a '
                             'workspace admin')
        else:
            row['state'] = 'expired'
            row['reason'] = f'marked dead by the pool ({detail or "authentication failed"})'
        return done(row)
    if has_creds:
        state, reason = creds_state(cpath, now)
        row['state'] = state
        row['reason'] = reason
        return done(row)
    if machine and _elsewhere(home, machine):
        row['state'] = 'remote'
        row['reason'] = f'no auth here — the grant lives on the {home} machine'
    else:
        row['state'] = 'missing'
        row['reason'] = 'no credentials on this machine'
    return done(row)


def audit_all(root, machine=None, now=None):
    try:
        doc = json.load(open(os.path.join(root, 'accounts.json')))
    except Exception:
        return []
    out = []
    for a in doc.get('accounts', []):
        if isinstance(a, dict) and VALID_ID.fullmatch(str(a.get('id', ''))):
            out.append(audit_account(root, a, now=now, machine=machine))
    return out


if __name__ == '__main__':
    root = sys.argv[1]
    machine = sys.argv[2] if len(sys.argv) > 2 else None
    for r in audit_all(root, machine=machine):
        # Tabs are the field separator, so no field may contain one.
        print('\t'.join(str(r[k]).replace('\t', ' ')
                        for k in ('id', 'email', 'home', 'state', 'label', 'reason', 'fix')))