#!/usr/bin/env bash
# vshell-ai-usage — compact, span-free JSON of Claude Code / OpenAI Codex subscription
# usage, for the VGS "aiUsage" bar widget.
#
# Wraps the vendored claudebar/codexbar engines (which do the OAuth refresh and
# call the provider usage API) and re-emits a small JSON object the QML widget
# can parse without dealing with Pango markup.
#
# Multi-account: people run several subscriptions side by side by pointing
# CLAUDE_CONFIG_DIR (or CODEX_HOME) at different directories, one per wrapper
# script. Every such directory is discovered and reported as its own account,
# labelled by the signed-in email. Single-account machines produce a one-entry
# list, and the top-level fields still describe that account, so the widget
# looks and behaves exactly as it did before.
#
# Usage: vshell-ai-usage [claude|codex]
#
# Output (one line of JSON):
#   {"ok":true,"provider":"claude","plan":"Max 20x","class":"critical",
#    "session":{"pct":16,"reset":"3h 53m","resetAt":1785265401},
#    "weekly":{"pct":85,"reset":"5d 12h","resetAt":1785409199},
#    "third":{"label":"Fable","pct":89,"reset":"5d 12h"},
#    "accounts":[{"id":"claude","label":"you@example.com","plan":"Max 20x",
#                 "ok":true,"class":"high","session":{...},"weekly":{...},
#                 "models":[{"label":"Fable","pct":89,"reset":"5d 12h"}]}],
#    "aggregate":{"count":1,"session":16,"weekly":85,"class":"high"}}
# On failure:
#   {"ok":false,"provider":"codex","error":"..."}
set -euo pipefail

# Ensure the sibling engines (claudebar/codexbar) are found even when invoked
# from an environment without ~/.local/bin on PATH (e.g. the shell systemd unit).
SELF_DIR="$(cd "$(dirname "$(readlink -f "$0" 2>/dev/null || echo "$0")")" && pwd)"
case ":$PATH:" in
    *":$SELF_DIR:"*) : ;;
    *) PATH="$SELF_DIR:$PATH" ;;
esac
export PATH

# The shell (and anything it spawns) may itself be running under one account's
# wrapper, which would otherwise silently make every lookup report that account.
# Account selection here is always explicit via --config-dir.
unset CLAUDE_CONFIG_DIR CODEX_HOME

provider="${1:-claude}"

case "$provider" in
    claude) engine=claudebar ;;
    codex)  engine=codexbar ;;
    *)
        jq -nc --arg p "$provider" '{ok:false,provider:$p,error:("unknown provider: "+$p)}'
        exit 0
        ;;
esac

if ! command -v "$engine" >/dev/null 2>&1; then
    jq -nc --arg p "$provider" --arg e "$engine" '{ok:false,provider:$p,error:($e+" not found in PATH")}'
    exit 0
fi

# --- Account discovery -------------------------------------------------------

# Claude: any directory holding a .credentials.json with a claudeAiOauth block.
# Profiles whose tokens live in the desktop keyring instead (mcpOAuth only) are
# not pollable over this API and are left out rather than shown permanently broken.
discover_claude() {
    local dir
    for dir in "$HOME/.claude" "$HOME"/.*claude*; do
        [[ -d "$dir" ]] || continue
        [[ -f "$dir/.credentials.json" ]] || continue
        jq -e '.claudeAiOauth.accessToken // empty' "$dir/.credentials.json" >/dev/null 2>&1 || continue
        printf '%s\n' "$dir"
    done | awk '!seen[$0]++'
}

discover_codex() {
    local dir
    for dir in "$HOME/.codex" "$HOME"/.codex-*; do
        [[ -d "$dir" ]] || continue
        [[ -f "$dir/auth.json" ]] || continue
        printf '%s\n' "$dir"
    done | awk '!seen[$0]++'
}

# The signed-in email is the only label that means anything when several
# accounts are on the same plan. The default profile keeps its profile JSON at
# ~/.claude.json rather than inside the config dir, so both layouts are tried.
label_for() {
    local dir="$1" email=""
    if [[ "$provider" == "claude" ]]; then
        for candidate in "$dir/.claude.json" "$HOME/.claude.json"; do
            [[ -f "$candidate" ]] || continue
            email=$(jq -r '.oauthAccount.emailAddress // empty' "$candidate" 2>/dev/null || true)
            [[ -n "$email" ]] && break
            # Only fall back to the home-root profile for the default config dir;
            # an alt dir without its own profile must not borrow the default's identity.
            [[ "$dir" == "$HOME/.claude" ]] || break
        done
        [[ -z "$email" && -f "$dir/.auth-email-cache" ]] && email=$(head -n1 "$dir/.auth-email-cache" 2>/dev/null || true)
    fi
    printf '%s' "$email"
}

case "$provider" in
    claude) mapfile -t dirs < <(discover_claude) ;;
    codex)  mapfile -t dirs < <(discover_codex) ;;
esac

if ((${#dirs[@]} == 0)); then
    jq -nc --arg p "$provider" '{ok:false,provider:$p,error:"no signed-in accounts found"}'
    exit 0
fi

# --- Per-account collection --------------------------------------------------
# Sequential on purpose: each engine keeps its own short-lived cache and fetch
# lock per account, so walking them in turn spreads any real API calls out
# instead of firing every account at once.

accounts_json="[]"
for dir in "${dirs[@]}"; do
    id="$(basename "$dir")"
    id="${id#.}"
    raw="$("$engine" --config-dir "$dir" --emit-json 2>/dev/null || true)"
    if [[ -z "$raw" ]] || ! jq -e . >/dev/null 2>&1 <<<"$raw"; then
        raw='{"ok":false,"error":"no output from engine"}'
    fi

    label="$(label_for "$dir")"
    # Codex reports the account email in its usage payload; prefer it.
    [[ -z "$label" ]] && label="$(jq -r '.email // empty' <<<"$raw" 2>/dev/null || true)"
    [[ -z "$label" ]] && label="$id"

    # Not fatal. Under `set -e` a jq failure here would abort the script with no
    # output at all, so one unreadable account would cost the whole payload and
    # the widget would show a blanket error instead of the accounts that are
    # fine. One account degrades to an unavailable entry, and the reason is said
    # out loud on stderr rather than swallowed.
    normalized=""
    if ! normalized="$(jq -c \
        --argjson acc "$accounts_json" \
        --argjson entry "$raw" \
        --arg id "$id" \
        --arg label "$label" \
        --arg dir "$dir" \
        '
        # Engine messages are written for Pango tooltip markup. The widget
        # renders every string through StyledText, which is Text.PlainText by
        # design (RichText would make provider-supplied strings an injection
        # surface), so the markup has to be resolved here — one owner, same
        # reason the "\n" flattening below lives here rather than in QML.
        #
        # Only the Pango tag names are stripped, so a message containing "<3"
        # or "a < b" survives intact. Tags come out before the entities are
        # unescaped, so text the engine deliberately escaped ("&lt;b&gt;")
        # stays literal instead of being stripped on a second pass; "&amp;"
        # is unescaped last for the same reason.
        #
        # The "i" flag needs three-argument gsub, which is NOT a newer builtin
        # than the two-argument form: jq 1.5 (2015-08-16) defines them side by
        # side in builtin.c —
        #   def gsub($re; s; flags): sub($re; s; flags + "g");
        #   def gsub($re; s): sub($re; s; "g");
        # — and jq 1.4 has no regex builtins at all, so a jq that cannot run
        # gsub/3 cannot run the gsub/2 calls this script has always used. The
        # flag therefore adds no version exposure of its own; jq >= 1.5 is the
        # floor for this script as a whole, and it is declared as such in
        # config/vshell/dependencies.json.
        def pango:
          if . == null then null
          else gsub("</?(b|i|u|s|tt|big|small|sub|sup|span|markup)([ \t\r\n][^>]*)?>"; ""; "i")
               | gsub("&lt;"; "<") | gsub("&gt;"; ">")
               | gsub("&quot;"; "\"")
               # An apostrophe cannot appear literally in this jq program —
               # the whole program is a single-quoted bash string — so the
               # replacement is built from its code point.
               | gsub("&(apos|#0*39);"; ([39] | implode); "i")
               | gsub("&amp;"; "&")
          end;
        $acc + [{
            id: $id,
            label: $label,
            configDir: $dir,
            ok: ($entry.ok == true),
            # Also carries literal "\n" sequences; flatten them so the widget
            # never prints them raw.
            error: (($entry.error // null) | if . == null then null
                    else gsub("\\\\n"; " ") | gsub("\\\\t"; " ") | pango
                         | gsub("[[:space:]]+"; " ") | sub("^ +"; "") | sub(" +$"; "") end),
            plan: (($entry.plan // "") | pango),
            # Credit-billed Claude seats identify as Enterprise and expose a
            # spend pool instead of rate-limit windows. Keep the explicit flag
            # in the shared payload so every UI can group account types without
            # inferring them from an email domain.
            enterprise: (((($entry.plan // "") | ascii_downcase) | startswith("enterprise"))
                         or ($entry.spend != null)),
            class: ($entry.class // "low"),
            session: ($entry.session // null),
            weekly: ($entry.weekly // null),
            # Every remaining engine-populated string the widget renders as
            # text goes through the same normalization — a lane label and a
            # spend detail reach StyledText exactly as an error message does.
            models: (($entry.models // [])
                     | map(if type == "object" and (.label | type) == "string"
                           then .label = (.label | pango) else . end)),
            spend: (($entry.spend // null)
                    | if type == "object" and (.detail | type) == "string"
                      then .detail = (.detail | pango) else . end)
        }]' <<<'null')"; then
        printf 'vshell-ai-usage: jq failed normalizing account %s (%s); reporting it as unavailable\n' \
            "$id" "$(jq --version 2>/dev/null || echo 'jq version unknown')" >&2
        normalized="$(jq -c \
            --argjson acc "$accounts_json" \
            --arg id "$id" \
            --arg label "$label" \
            --arg dir "$dir" \
            '$acc + [{id:$id, label:$label, configDir:$dir, ok:false,
                      error:"account payload could not be normalized"}]' <<<'null')"
    fi
    accounts_json="$normalized"
done

# --- Aggregate + backward-compatible top level -------------------------------
# The pill shows consumption against the whole pool, so each account counts for
# one 100% allowance and the aggregate is consumed/available across them — five
# accounts at 50% read as 50%, never 250%.

jq -nc \
    --arg provider "$provider" \
    --argjson accounts "$accounts_json" \
    '
    ($accounts | map(select(.ok))) as $live
    | ($live | map(select(.session != null))) as $sess
    | ($live | map(select(.weekly  != null))) as $week
    | (if ($sess | length) > 0 then (($sess | map(.session.pct) | add) / ($sess | length) | round) else null end) as $aggSession
    | (if ($week | length) > 0 then (($week | map(.weekly.pct)  | add) / ($week | length) | round) else null end) as $aggWeekly
    # Headline: each account counts for its TIGHTEST window, because that is what
    # actually blocks you — an account with a 3% session and a 96% weekly is 96%
    # spent, not 3%. Averaging one window type across accounts (which is what
    # $aggSession/$aggWeekly do) reported 8% on a pool that was nearly exhausted.
    | ($live | map([ (.session.pct // empty), (.weekly.pct // empty),
                     (.models[]?.pct // empty), (.spend.pct // empty) ]
                   | map(select(. != null)) | max // 0)) as $tightest
    | (if ($tightest | length) > 0
       then (($tightest | add) / ($tightest | length) | round) else 0 end) as $aggPct
    | ([ $aggSession, $aggWeekly,
         ($live | map(.models[]?.pct) | max),
         ($live | map(.spend?.pct // empty) | max) ] | map(select(. != null)) | max // 0) as $peak
    | (if   $peak >= 90 then "critical"
       elif $peak >= 75 then "high"
       elif $peak >= 50 then "mid"
       else "low" end) as $aggClass
    # Primary account keeps the single-account contract intact: same top-level
    # shape the widget has always parsed, so nothing regresses with one account.
    | ($live[0] // $accounts[0]) as $primary
    | ($primary.models // []) as $pm
    | if ($live | length) == 0 then
        {ok:false, provider:$provider,
         error: ($accounts[0].error // "usage unavailable"),
         accounts:$accounts}
      else
        {ok:true, provider:$provider,
         plan: ($primary.plan // ""),
         class: $aggClass,
         session: $primary.session,
         weekly: $primary.weekly,
         third: (if ($pm | length) > 0
                 then ($pm | max_by(.pct) | {label:.label, pct:.pct, reset:.reset, resetAt:(.resetAt // 0)})
                 else null end),
         accounts: $accounts,
         aggregate: {count: ($live | length),
                     total: ($accounts | length),
                     pct: $aggPct,
                     session: $aggSession,
                     weekly: $aggWeekly,
                     class: $aggClass}}
      end
    '
