#!/usr/bin/env python3
"""Show live subscription usage for each AI provider used by agent runtimes."""

import json
import os
import subprocess
import sys
import urllib.error
import urllib.request
from datetime import datetime, timezone
from pathlib import Path

ANTHROPIC_TOKEN_PATH = Path("~/dev/claude/secrets/anthropic/api_key").expanduser()
MINIMAX_SOPS_PATH = Path(
    "~/dev/claude/projects/agent-runtime-secrets/providers/minimax/v1/provider.sops.env"
).expanduser()
MINIMAX_SOPS_KEY = Path("~/dev/claude/secrets/sops/provider-age-key.txt").expanduser()
MINIMAX_DOTENV_KEY = "ANTHROPIC_AUTH_TOKEN"

# Reset field names. Anthropic emits Unix epoch seconds; MiniMax emits milliseconds.
ANTHROPIC_RESET_HEADERS = {
    "five_hour": "anthropic-ratelimit-unified-5h-reset",
    "seven_day": "anthropic-ratelimit-unified-7d-reset",
}
MINIMAX_RESET_FIELDS = {
    "five_hour": "end_time",          # ms since epoch
    "seven_day": "weekly_end_time",   # ms since epoch
}

# Nominal window lengths in seconds, used to compute elapsed_pct. These are
# fixed constants, not queried at runtime: Anthropic documents 5h/7d unified
# rate-limit windows (https://platform.claude.com/docs/en/api/rate-limits),
# and MiniMax v2's interval/weekly reset fields (end_time / weekly_end_time)
# are consistent with the same cadence. Neither API exposes an authoritative
# "window length" value to verify against, so these are assumed nominal
# lengths rather than confirmed at runtime.
WINDOW_SECONDS = {
    "five_hour": 18000,   # 5 * 3600
    "seven_day": 604800,  # 7 * 86400
}

# ANSI colour support
_use_colour = sys.stdout.isatty()
_GREEN  = "\033[32m" if _use_colour else ""
_YELLOW = "\033[33m" if _use_colour else ""
_RED    = "\033[31m" if _use_colour else ""
_RESET  = "\033[0m"  if _use_colour else ""
_BOLD   = "\033[1m"  if _use_colour else ""
_DIM    = "\033[2m"  if _use_colour else ""


def _colour_pct(pct: float) -> str:
    s = f"{pct:5.1f}%"
    if pct >= 80:
        return f"{_RED}{s}{_RESET}"
    if pct >= 60:
        return f"{_YELLOW}{s}{_RESET}"
    return f"{_GREEN}{s}{_RESET}"


def _parse_dotenv_key(content: str, key: str) -> str:
    for line in content.splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        if "=" in line:
            k, _, v = line.partition("=")
        elif ": " in line:
            k, _, v = line.partition(": ")
        else:
            continue
        if k.strip() == key:
            v = v.strip()
            if len(v) >= 2 and v[0] == v[-1] and v[0] in ('"', "'"):
                v = v[1:-1]
            return v
    raise KeyError(f"Key '{key}' not found in decrypted content")


def _read_anthropic_token() -> str:
    return ANTHROPIC_TOKEN_PATH.read_text().strip()


def _read_minimax_key() -> str:
    env = {**os.environ, "SOPS_AGE_KEY_FILE": str(MINIMAX_SOPS_KEY)}
    result = subprocess.run(
        ["sops", "--decrypt", "--output-type", "dotenv", str(MINIMAX_SOPS_PATH)],
        capture_output=True, text=True, check=True, env=env,
    )
    return _parse_dotenv_key(result.stdout, MINIMAX_DOTENV_KEY)


def _parse_anthropic_headers(headers) -> tuple[dict[str, float], dict[str, int]]:
    """Returns (utilization_pct_by_window, reset_epoch_seconds_by_window).

    Reset values are parsed as integer Unix epoch seconds; absent or malformed
    headers yield no entry for that window.
    """
    util: dict[str, float] = {}
    reset: dict[str, int] = {}
    fh_u = headers.get("anthropic-ratelimit-unified-5h-utilization")
    sd_u = headers.get("anthropic-ratelimit-unified-7d-utilization")
    for key, raw in [("five_hour", fh_u), ("seven_day", sd_u)]:
        if raw is not None:
            try:
                util[key] = float(raw) * 100.0
            except ValueError:
                pass
    for key, hdr in ANTHROPIC_RESET_HEADERS.items():
        raw = headers.get(hdr)
        if raw is None:
            continue
        try:
            reset[key] = int(raw)
        except ValueError:
            pass
    return util, reset


def probe_anthropic(token: str) -> tuple[dict[str, float], dict[str, int]] | None:
    """Returns (utilization_pct, reset_epoch_seconds) or None on total failure.

    Either dict may be empty; both are populated independently.
    """
    auth = {"Authorization": f"Bearer {token}", "anthropic-version": "2023-06-01"}

    try:
        req = urllib.request.Request("https://api.anthropic.com/v1/models", headers=auth)
        with urllib.request.urlopen(req, timeout=10) as resp:
            util, reset = _parse_anthropic_headers(resp.headers)
            if util or reset:
                return util, reset
    except Exception:
        pass

    try:
        body = json.dumps({
            "model": "claude-haiku-4-5-20251001",
            "max_tokens": 1,
            "messages": [{"role": "user", "content": "hi"}],
        }).encode()
        req = urllib.request.Request(
            "https://api.anthropic.com/v1/messages",
            data=body,
            headers={**auth, "content-type": "application/json"},
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=15) as resp:
            util, reset = _parse_anthropic_headers(resp.headers)
            if util or reset:
                return util, reset
            return None
    except urllib.error.HTTPError as e:
        util, reset = _parse_anthropic_headers(e.headers)
        if util or reset:
            return util, reset
        return None
    except Exception:
        pass

    return None


def probe_minimax(api_key: str) -> tuple[dict[str, float], dict[str, int]] | tuple[None, str]:
    """Returns ((utilization_pct, reset_epoch_seconds), None) on success,
    or (None, error_reason) on failure.

    API v2 shape (current): model_remains[].model_name with remaining_percent fields
    and end_time/weekly_end_time in **milliseconds**.
    API v1 shape (legacy): category_remains[].category == "text_generation" with count
    fields only — no reset timestamps.
    Tries v2 first; falls back to v1.
    """
    try:
        req = urllib.request.Request(
            "https://www.minimax.io/v1/token_plan/remains",
            headers={
                "Authorization": f"Bearer {api_key}",
                "User-Agent": "curl/7.88.1",
            },
        )
        with urllib.request.urlopen(req, timeout=10) as resp:
            data = json.loads(resp.read())
        status = data.get("base_resp", {}).get("status_code")
        if status != 0:
            msg = data.get("base_resp", {}).get("status_msg", "unknown")
            return None, f"API status_code={status} ({msg})"

        # v2: model_remains with remaining_percent fields (general = text)
        for entry in data.get("model_remains", []):
            if entry.get("model_name") == "general":
                util: dict[str, float] = {}
                reset: dict[str, int] = {}
                ih = entry.get("current_interval_remaining_percent")
                wh = entry.get("current_weekly_remaining_percent")
                if ih is not None:
                    util["five_hour"] = 100.0 - float(ih)
                if wh is not None:
                    util["seven_day"] = 100.0 - float(wh)
                # Reset fields are epoch milliseconds; normalise to seconds.
                end_ms  = entry.get("end_time")
                week_ms = entry.get("weekly_end_time")
                if isinstance(end_ms, (int, float)) and end_ms > 0:
                    reset["five_hour"] = int(end_ms) // 1000
                if isinstance(week_ms, (int, float)) and week_ms > 0:
                    reset["seven_day"] = int(week_ms) // 1000
                if util or reset:
                    return util, reset
                return None, "general model entry has no usable fields"

        # v1 fallback: category_remains with count fields — no reset timestamps.
        for cat in data.get("category_remains", []):
            if cat.get("category") == "text_generation":
                util = {}
                interval_total = cat.get("current_interval_total_count", 0)
                interval_used  = cat.get("current_interval_usage_count", 0)
                weekly_total   = cat.get("current_weekly_total_count", 0)
                weekly_used    = cat.get("current_weekly_usage_count", 0)
                if interval_total > 0:
                    util["five_hour"] = float(interval_used) / float(interval_total) * 100.0
                if weekly_total > 0:
                    util["seven_day"] = float(weekly_used) / float(weekly_total) * 100.0
                if util:
                    return util, {}  # no reset fields in v1
                return None, "category_remains has zero counts"

        return None, "no usable plan data in response (keys: " + ", ".join(data.keys()) + ")"
    except urllib.error.HTTPError as e:
        return None, f"HTTP {e.code} {e.reason}"
    except urllib.error.URLError as e:
        return None, f"network error: {e.reason}"
    except Exception as e:
        return None, str(e)


def _format_relative(seconds: int | None) -> str:
    """Human-readable relative time, e.g. '6d 8h', '4h 37m', '12m', '45s', 'UNKNOWN'."""
    if seconds is None:
        return "UNKNOWN"
    if seconds < 0:
        seconds = 0
    days, rem = divmod(seconds, 86400)
    hours, rem = divmod(rem, 3600)
    minutes, secs = divmod(rem, 60)
    if days > 0:
        return f"{days}d {hours}h"
    if hours > 0:
        return f"{hours}h {minutes}m"
    if minutes > 0:
        return f"{minutes}m"
    return f"{secs}s"


def _iso_from_epoch(seconds: int | None) -> str | None:
    if seconds is None:
        return None
    return datetime.fromtimestamp(seconds, tz=timezone.utc).isoformat()


def _iso_now(now: datetime | None = None) -> str:
    return (now or datetime.now(timezone.utc)).isoformat()


def _reset_in_seconds(reset_epoch: int | None, now_epoch: int | None = None) -> int | None:
    """Seconds from now until reset_epoch, clamped to 0 (never negative).

    A negative raw delta means the reset already fired (clock skew between
    probe time and the provider's clock, or a probe that lands right on the
    boundary) -- clamp at the source so every downstream consumer sees a
    consistent floor instead of re-implementing the clamp themselves. This
    also pins elapsed_pct (below) at 1.0 in that case.
    """
    if reset_epoch is None:
        return None
    now = now_epoch if now_epoch is not None else int(datetime.now(timezone.utc).timestamp())
    return max(0, reset_epoch - now)


def _elapsed_pct(reset_in_seconds: int | None, window_key: str) -> float | None:
    """Fraction of the window already elapsed: 1 - reset_in_seconds/window_seconds.

    Clamped to [0.0, 1.0]. Returns None (fail-closed) when reset_in_seconds is
    None -- i.e. no reset timestamp is available for this window (the MiniMax
    v1 fallback path returns utilization but no reset fields). Consumers must
    treat None as "cannot pace this window", never as 0.0 or 1.0.
    """
    if reset_in_seconds is None:
        return None
    window_seconds = WINDOW_SECONDS[window_key]
    pct = 1.0 - (reset_in_seconds / window_seconds)
    return max(0.0, min(1.0, pct))


def print_dryrun(json_mode: bool = False) -> None:
    if json_mode:
        # Sample reset_in_seconds; elapsed_pct is derived from these via the
        # real _elapsed_pct() so the sample stays consistent with live output.
        fh_reset_in_seconds = 16620
        sd_reset_in_seconds = 538200
        sample = {
            "probed_at": "2026-08-01T14:23:00+00:00",
            "providers": [
                {
                    "provider": "Anthropic",
                    "available": True,
                    "windows": {
                        "five_hour": {
                            "utilization_pct": 23.4,
                            "reset_at": "2026-08-01T19:00:00+00:00",
                            "reset_in_seconds": fh_reset_in_seconds,
                            "window_seconds": WINDOW_SECONDS["five_hour"],
                            "elapsed_pct": _elapsed_pct(fh_reset_in_seconds, "five_hour"),
                        },
                        "seven_day": {
                            "utilization_pct": 41.2,
                            "reset_at": "2026-08-07T23:00:00+00:00",
                            "reset_in_seconds": sd_reset_in_seconds,
                            "window_seconds": WINDOW_SECONDS["seven_day"],
                            "elapsed_pct": _elapsed_pct(sd_reset_in_seconds, "seven_day"),
                        },
                    },
                },
                {
                    "provider": "MiniMax",
                    "available": True,
                    "windows": {
                        "five_hour": {
                            "utilization_pct": 12.1,
                            "reset_at": "2026-08-01T19:00:00+00:00",
                            "reset_in_seconds": fh_reset_in_seconds,
                            "window_seconds": WINDOW_SECONDS["five_hour"],
                            "elapsed_pct": _elapsed_pct(fh_reset_in_seconds, "five_hour"),
                        },
                        "seven_day": {
                            "utilization_pct": 8.3,
                            "reset_at": "2026-08-07T23:00:00+00:00",
                            "reset_in_seconds": sd_reset_in_seconds,
                            "window_seconds": WINDOW_SECONDS["seven_day"],
                            "elapsed_pct": _elapsed_pct(sd_reset_in_seconds, "seven_day"),
                        },
                    },
                },
            ],
        }
        print(json.dumps(sample, indent=2))
        return

    print(f"{_BOLD}[dryrun] Would probe:{_RESET}")
    print(f"  {_BOLD}Anthropic{_RESET}  — {ANTHROPIC_TOKEN_PATH} (Bearer OAuth token)")
    print(f"             GET https://api.anthropic.com/v1/models")
    print(f"  {_BOLD}MiniMax{_RESET}    — {MINIMAX_SOPS_PATH} (SOPS, key: {MINIMAX_DOTENV_KEY})")
    print(f"             GET https://www.minimax.io/v1/token_plan/remains")


def print_table(
    rows: list[tuple[str, str, str, str]],
    col_widths: tuple[int, int, int, int] = (12, 10, 8, 12),
) -> None:
    print(f"\n{_BOLD}Agent Subscription Usage{_RESET}")
    print("========================\n")
    c1, c2, c3, c4 = col_widths
    header = f"{'Provider':<{c1}}  {'Window':<{c2}}  {'Usage':<{c3}}  {'Resets':<{c4}}"
    sep    = f"{'-'*c1}  {'-'*c2}  {'-'*c3}  {'-'*c4}"
    print(f"{_BOLD}{header}{_RESET}")
    print(sep)
    for provider, window, usage, resets in rows:
        print(f"{provider:<{c1}}  {window:<{c2}}  {usage:<{c3}}  {resets:<{c4}}")
    print()


def _window_record(util: dict[str, float], reset: dict[str, int], key: str, now_epoch: int) -> dict:
    pct = util.get(key)
    reset_epoch = reset.get(key)
    reset_in_seconds = _reset_in_seconds(reset_epoch, now_epoch)
    return {
        "utilization_pct": pct,
        "reset_at": _iso_from_epoch(reset_epoch),
        "reset_in_seconds": reset_in_seconds,
        "window_seconds": WINDOW_SECONDS[key],
        "elapsed_pct": _elapsed_pct(reset_in_seconds, key),
    }


def build_json_report(providers: list[dict]) -> dict:
    return {
        "probed_at": _iso_now(),
        "providers": providers,
    }


def _provider_record(
    name: str,
    util: dict[str, float] | None,
    reset: dict[str, int] | None,
    error: str | None,
    now_epoch: int,
) -> dict:
    if error is not None:
        return {"provider": name, "available": False, "error": error}
    if util is None and reset is None:
        return {"provider": name, "available": False, "error": "no data returned"}
    windows: dict[str, dict] = {}
    for key in ("five_hour", "seven_day"):
        windows[key] = _window_record(util or {}, reset or {}, key, now_epoch)
    return {"provider": name, "available": True, "windows": windows}


def main() -> None:
    dryrun = False
    json_mode = False
    args = sys.argv[1:]
    i = 0
    while i < len(args):
        arg = args[i]
        if arg in ("-h", "--help"):
            print(__doc__)
            print("Usage: agent-subscriptions [OPTIONS]")
            print()
            print("Options:")
            print("  -n, --dryrun         Preview probes without making API calls")
            print("  -j, --output json    Emit structured JSON (plain text, no ANSI)")
            print("  -h, --help           Show this help message and exit")
            sys.exit(0)
        elif arg in ("-n", "--dryrun"):
            dryrun = True
            i += 1
        elif arg == "-j":
            json_mode = True
            i += 1
        elif arg == "--output":
            if i + 1 >= len(args):
                print("Error: --output requires a value (json)", file=sys.stderr)
                sys.exit(1)
            value = args[i + 1]
            if value != "json":
                print(f"Error: unsupported --output value: {value}", file=sys.stderr)
                sys.exit(1)
            json_mode = True
            i += 2
        elif arg.startswith("--output="):
            value = arg.split("=", 1)[1]
            if value != "json":
                print(f"Error: unsupported --output value: {value}", file=sys.stderr)
                sys.exit(1)
            json_mode = True
            i += 1
        else:
            print(f"Error: unknown option: {arg}", file=sys.stderr)
            sys.exit(1)

    if dryrun:
        print_dryrun(json_mode=json_mode)
        return

    now_epoch = int(datetime.now(timezone.utc).timestamp())

    # Probe state: (utilization_pct dict, reset_epoch_seconds dict, error string)
    anth: tuple[dict | None, dict | None, str | None] = (None, None, None)
    mini: tuple[dict | None, dict | None, str | None] = (None, None, None)

    # Anthropic
    try:
        token = _read_anthropic_token()
        result = probe_anthropic(token)
        if result is None:
            anth = (None, None, "probe returned no data")
        else:
            anth = (result[0], result[1], None)
    except Exception as e:
        anth = (None, None, str(e))
        print(f"{_DIM}Warning: Anthropic probe failed: {e}{_RESET}", file=sys.stderr)

    # MiniMax
    try:
        key = _read_minimax_key()
        result = probe_minimax(key)
        if isinstance(result[1], str) and result[0] is None and not result[1]:
            mini = (None, None, "probe returned no data")
        elif result[1] is not None and not isinstance(result[1], dict):
            # error branch: (None, error_str)
            mini = (None, None, result[1])
            print(f"{_DIM}Warning: MiniMax probe: {result[1]}{_RESET}", file=sys.stderr)
        else:
            util, reset = result
            mini = (util, reset, None)
    except Exception as e:
        mini = (None, None, str(e))
        print(f"{_DIM}Warning: MiniMax probe failed: {e}{_RESET}", file=sys.stderr)

    if json_mode:
        report = build_json_report([
            _provider_record("Anthropic", anth[0], anth[1], anth[2], now_epoch),
            _provider_record("MiniMax",   mini[0], mini[1], mini[2], now_epoch),
        ])
        print(json.dumps(report, indent=2))
        return

    rows: list[tuple[str, str, str, str]] = []

    def add_rows(name: str, util: dict | None, reset: dict | None, error: str | None) -> None:
        if error is not None or (util is None and reset is None):
            rows.append((name, "5-hour", f"{_DIM}UNAVAILABLE{_RESET}", f"{_DIM}UNKNOWN{_RESET}"))
            rows.append((name, "7-day",  f"{_DIM}UNAVAILABLE{_RESET}", f"{_DIM}UNKNOWN{_RESET}"))
            return
        util = util or {}
        reset = reset or {}
        for window_label, key in (("5-hour", "five_hour"), ("7-day", "seven_day")):
            pct = util.get(key)
            usage_cell = _colour_pct(pct) if pct is not None else f"{_DIM}N/A{_RESET}"
            rel = _reset_in_seconds(reset.get(key), now_epoch)
            rows.append((
                name,
                window_label,
                usage_cell,
                _format_relative(rel) if rel is not None else f"{_DIM}UNKNOWN{_RESET}",
            ))

    add_rows("Anthropic", anth[0], anth[1], anth[2])
    add_rows("MiniMax",   mini[0], mini[1], mini[2])

    print_table(rows)


if __name__ == "__main__":
    main()