#!/usr/bin/env bash

spend_cap_status() {
  local account="$1" output
  SPEND_CAP_WINDOWS=""
  SPEND_CAP_RESET_AT=""
  output="$(SPEND_CAP_ACCOUNT="$account" \
    SPEND_CAP_STALE_AFTER_S="${SPEND_CAP_STALE_AFTER_S:-3600}" \
    SPEND_CAP_HEALTH_CACHE="${SPEND_CAP_HEALTH_CACHE:-$HOME/.local/state/overdeck/systray/runtime/health_cache.json}" \
    SPEND_CAP_ROUTING_RULES="${SPEND_CAP_ROUTING_RULES:-$HOME/.local/state/overdeck/systray/runtime/routing_rules.json}" \
    SPEND_CAP_CLAUDE_ROUTING_RULES="${SPEND_CAP_CLAUDE_ROUTING_RULES:-$HOME/.local/state/overdeck/systray/runtime/claude_routing_rules.json}" \
    python3 - <<'PY'
import json
import os
import sys
from datetime import UTC, datetime

account = os.environ.get("SPEND_CAP_ACCOUNT", "")
stale_after_s = int(os.environ.get("SPEND_CAP_STALE_AFTER_S", "3600"))
health_path = os.environ["SPEND_CAP_HEALTH_CACHE"]
rules_paths = [
    os.environ["SPEND_CAP_ROUTING_RULES"],
    os.environ["SPEND_CAP_CLAUDE_ROUTING_RULES"],
]

WINDOW_5H = "5h"
WINDOW_7D = "7d"
WINDOW_USED = {WINDOW_5H: "primary_used_pct", WINDOW_7D: "secondary_used_pct"}
WINDOW_RESET = {WINDOW_5H: "primary_reset_at", WINDOW_7D: "secondary_reset_at"}


def load_rules_caps():
    merged = {}
    for path in rules_paths:
        try:
            with open(path, encoding="utf-8") as handle:
                payload = json.load(handle)
        except (OSError, json.JSONDecodeError):
            return None
        if not isinstance(payload, dict):
            return None
        account_caps = payload.get("account_caps", {})
        if account_caps is None:
            account_caps = {}
        if not isinstance(account_caps, dict):
            return None
        for slug, windows in account_caps.items():
            if not isinstance(slug, str) or not slug or not isinstance(windows, dict):
                return None
            entry = merged.setdefault(slug, {})
            for window, cap_pct in windows.items():
                if window not in (WINDOW_5H, WINDOW_7D):
                    return None
                if not isinstance(cap_pct, int) or isinstance(cap_pct, bool) or not 1 <= cap_pct <= 100:
                    return None
                entry[window] = cap_pct
    return merged


def health_age_seconds():
    try:
        return datetime.now().timestamp() - os.path.getmtime(health_path)
    except OSError:
        return None


def load_snapshot(slug):
    try:
        with open(health_path, encoding="utf-8") as handle:
            payload = json.load(handle)
    except (OSError, json.JSONDecodeError):
        return None
    if not isinstance(payload, dict):
        return None
    value = payload.get(slug)
    return value if isinstance(value, dict) else None


def capped_windows(snapshot, caps):
    capped = []
    for window, cap_pct in caps.items():
        used_pct = snapshot.get(WINDOW_USED[window])
        if isinstance(used_pct, bool) or not isinstance(used_pct, int):
            used_pct = None
        if used_pct is not None and (100 - used_pct) <= cap_pct:
            capped.append(window)
    return capped


def cap_reset_at(snapshot, windows):
    reset_times = []
    for window in windows:
        reset_at = snapshot.get(WINDOW_RESET[window])
        if isinstance(reset_at, bool) or not isinstance(reset_at, (int, float)):
            continue
        reset_times.append(float(reset_at))
    if not reset_times:
        return ""
    dt = datetime.fromtimestamp(int(min(reset_times)), tz=UTC)
    return dt.strftime("%Y-%m-%dT%H:%M:%SZ")


caps_by_slug = load_rules_caps()
if caps_by_slug is None:
    print("unreadable")
    sys.exit(0)

account_caps = caps_by_slug.get(account, {})
if not account_caps:
    print("no-caps")
    sys.exit(0)

age = health_age_seconds()
if age is None or age > stale_after_s:
    print("stale")
    sys.exit(0)

snapshot = load_snapshot(account)
if snapshot is None:
    print("stale")
    sys.exit(0)

breached = capped_windows(snapshot, account_caps)
if not breached:
    print("clear")
    sys.exit(0)

ordered = [window for window in (WINDOW_5H, WINDOW_7D) if window in breached]
print("capped")
print(",".join(ordered))
print(cap_reset_at(snapshot, ordered))
PY
)" || {
    printf '%s\n' unreadable
    return 0
  }
  case "$output" in
    capped*)
      SPEND_CAP_WINDOWS="$(printf '%s\n' "$output" | sed -n '2p')"
      SPEND_CAP_RESET_AT="$(printf '%s\n' "$output" | sed -n '3p')"
      printf '%s\n' capped
      ;;
    *)
      printf '%s\n' "$output"
      ;;
  esac
}

spend_cap_default_account() {
  local output
  output="$(SPEND_CAP_ROUTING_RULES="${SPEND_CAP_ROUTING_RULES:-$HOME/.local/state/overdeck/systray/runtime/routing_rules.json}" \
    python3 - <<'PY'
import json
import os
import sys

path = os.environ["SPEND_CAP_ROUTING_RULES"]
try:
    with open(path, encoding="utf-8") as handle:
        payload = json.load(handle)
except (OSError, json.JSONDecodeError):
    sys.exit(0)

if not isinstance(payload, dict):
    sys.exit(0)

default = payload.get("default")
if isinstance(default, str) and default:
    print(default)
PY
)" || true
  [[ -n "$output" ]] && printf '%s\n' "$output"
}

spend_cap_refuse_if_capped() {
  local account="$1" detail="$2" thread_id="$3" protocol="$4" model="$5"
  local status resume_at status_json
  status="$(spend_cap_status "$account")"
  [[ "$status" == capped ]] || return 0
  resume_at="${SPEND_CAP_RESET_AT:-}"
  if [[ -z "$resume_at" ]]; then
    resume_at="$(python3 - <<'PY'
from datetime import UTC, datetime, timedelta
print((datetime.now(UTC) + timedelta(seconds=60)).strftime("%Y-%m-%dT%H:%M:%SZ"))
PY
)"
  fi
  status_json="$(THREAD_ID="$thread_id" RESUME_AT="$resume_at" python3 - <<'PY'
import json
import os

status = {"ok": False, "detail": "rate-limited", "resume_at": os.environ["RESUME_AT"], "capped": True}
thread_id = os.environ.get("THREAD_ID", "")
if thread_id:
    status["thread_id"] = thread_id
print(json.dumps(status, separators=(",", ":")))
PY
)"
  printf '%s\n' "$status_json"
  exit 75
}

test_zync2_remaining_cap_inversion_regression() {
  local tmp_dir health rules status
  tmp_dir="$(mktemp -d)"
  health="$tmp_dir/health_cache.json"
  rules="$tmp_dir/routing_rules.json"
  printf '%s\n' '{"version":"routing/v2","projects":{},"default":"zync2","fallback_chain":[],"account_caps":{"zync2":{"7d":10}}}' >"$rules"
  printf '%s\n' '{}' >"$tmp_dir/claude_routing_rules.json"

  printf '%s\n' '{"zync2":{"status":"ok","primary_used_pct":null,"secondary_used_pct":88,"secondary_reset_at":1786365653.0}}' >"$health"
  status="$(SPEND_CAP_HEALTH_CACHE="$health" SPEND_CAP_ROUTING_RULES="$rules" SPEND_CAP_CLAUDE_ROUTING_RULES="$tmp_dir/claude_routing_rules.json" SPEND_CAP_STALE_AFTER_S=3600 spend_cap_status zync2)"
  if [[ "$status" != clear ]]; then
    rm -rf "$tmp_dir"
    printf 'test_zync2_remaining_cap_inversion_regression: 88%% used expected clear, got %s\n' "$status" >&2
    return 1
  fi

  printf '%s\n' '{"zync2":{"status":"ok","primary_used_pct":null,"secondary_used_pct":90,"secondary_reset_at":1786365653.0}}' >"$health"
  status="$(SPEND_CAP_HEALTH_CACHE="$health" SPEND_CAP_ROUTING_RULES="$rules" SPEND_CAP_CLAUDE_ROUTING_RULES="$tmp_dir/claude_routing_rules.json" SPEND_CAP_STALE_AFTER_S=3600 spend_cap_status zync2)"
  if [[ "$status" != capped ]]; then
    rm -rf "$tmp_dir"
    printf 'test_zync2_remaining_cap_inversion_regression: 90%% used expected capped, got %s\n' "$status" >&2
    return 1
  fi

  printf '%s\n' '{"zync2":{"status":"ok","primary_used_pct":null,"secondary_used_pct":10,"secondary_reset_at":1786365653.0}}' >"$health"
  status="$(SPEND_CAP_HEALTH_CACHE="$health" SPEND_CAP_ROUTING_RULES="$rules" SPEND_CAP_CLAUDE_ROUTING_RULES="$tmp_dir/claude_routing_rules.json" SPEND_CAP_STALE_AFTER_S=3600 spend_cap_status zync2)"
  if [[ "$status" != clear ]]; then
    rm -rf "$tmp_dir"
    printf 'test_zync2_remaining_cap_inversion_regression: 10%% used expected clear, got %s\n' "$status" >&2
    return 1
  fi

  rm -rf "$tmp_dir"
  return 0
}

test_spend_cap_refuse_if_capped_emits_capped_flag() {
  local tmp_dir health rules output exit_code json
  tmp_dir="$(mktemp -d)"
  health="$tmp_dir/health_cache.json"
  rules="$tmp_dir/routing_rules.json"
  printf '%s\n' '{"version":"routing/v2","projects":{},"default":"zync2","fallback_chain":[],"account_caps":{"zync2":{"7d":10}}}' >"$rules"
  printf '%s\n' '{}' >"$tmp_dir/claude_routing_rules.json"
  printf '%s\n' '{"zync2":{"status":"ok","primary_used_pct":null,"secondary_used_pct":90,"secondary_reset_at":1786365653.0}}' >"$health"

  set +e
  output="$(SPEND_CAP_HEALTH_CACHE="$health" SPEND_CAP_ROUTING_RULES="$rules" SPEND_CAP_CLAUDE_ROUTING_RULES="$tmp_dir/claude_routing_rules.json" SPEND_CAP_STALE_AFTER_S=3600 bash -c "
    source \"\$1\"
    spend_cap_refuse_if_capped zync2 '' '' '' ''
  " bash "${BASH_SOURCE[0]}" 2>&1)"
  exit_code=$?
  set -e

  if [[ "$exit_code" -ne 75 ]]; then
    rm -rf "$tmp_dir"
    printf 'test_spend_cap_refuse_if_capped_emits_capped_flag: expected exit 75, got %s\n' "$exit_code" >&2
    return 1
  fi

  json="$(printf '%s\n' "$output" | tail -1)"
  if ! printf '%s' "$json" | python3 -c '
import json
import sys

status = json.loads(sys.stdin.read())
assert status.get("detail") == "rate-limited", status
assert status.get("capped") is True, status
assert isinstance(status.get("resume_at"), str) and status["resume_at"], status
'; then
    rm -rf "$tmp_dir"
    printf 'test_spend_cap_refuse_if_capped_emits_capped_flag: refusal JSON missing capped=true: %s\n' "$json" >&2
    return 1
  fi

  rm -rf "$tmp_dir"
  return 0
}

if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
  case "${1:-}" in
    --self-test)
      test_zync2_remaining_cap_inversion_regression && test_spend_cap_refuse_if_capped_emits_capped_flag
      exit $?
      ;;
    "")
      echo "usage: spend-cap.sh --self-test | spend-cap.sh <account>" >&2
      exit 2
      ;;
    *)
      spend_cap_status "$1"
      ;;
  esac
fi