#!/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 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"

# 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) -> dict[str, float]:
    result = {}
    fh = headers.get("anthropic-ratelimit-unified-5h-utilization")
    sd = headers.get("anthropic-ratelimit-unified-7d-utilization")
    for key, raw in [("five_hour", fh), ("seven_day", sd)]:
        if raw is not None:
            try:
                result[key] = float(raw) * 100.0
            except ValueError:
                pass
    return result


def probe_anthropic(token: str) -> dict[str, float] | None:
    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:
            result = _parse_anthropic_headers(resp.headers)
            if result:
                return result
    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:
            return _parse_anthropic_headers(resp.headers) or None
    except urllib.error.HTTPError as e:
        result = _parse_anthropic_headers(e.headers)
        return result or None
    except Exception:
        pass

    return None


def probe_minimax(api_key: str) -> tuple[dict[str, float] | None, str | None]:
    """Returns (usage_dict, error_reason). On success error_reason is None."""
    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})"
        for cat in data.get("category_remains", []):
            if cat.get("category") == "text_generation":
                result = {}
                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:
                    result["five_hour"] = float(interval_used) / float(interval_total) * 100.0
                if weekly_total > 0:
                    result["seven_day"] = float(weekly_used) / float(weekly_total) * 100.0
                return (result or None), None
        return None, "text_generation category not found in response"
    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 print_dryrun() -> None:
    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]]) -> None:
    print(f"\n{_BOLD}Agent Subscription Usage{_RESET}")
    print("========================\n")
    col1, col2, col3 = 12, 10, 8
    header = f"{'Provider':<{col1}}  {'Window':<{col2}}  {'Usage':<{col3}}"
    sep    = f"{'-'*col1}  {'-'*col2}  {'-'*col3}"
    print(f"{_BOLD}{header}{_RESET}")
    print(sep)
    for provider, window, usage in rows:
        print(f"{provider:<{col1}}  {window:<{col2}}  {usage}")
    print()


def main() -> None:
    dryrun = False
    args = sys.argv[1:]

    for arg in args:
        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("  -h, --help      Show this help message and exit")
            sys.exit(0)
        elif arg in ("-n", "--dryrun"):
            dryrun = True
        else:
            print(f"Error: unknown option: {arg}", file=sys.stderr)
            sys.exit(1)

    if dryrun:
        print_dryrun()
        return

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

    # Anthropic
    try:
        token = _read_anthropic_token()
        usage = probe_anthropic(token)
        if usage is None:
            rows.append(("Anthropic", "5-hour", f"{_DIM}UNAVAILABLE{_RESET}"))
            rows.append(("Anthropic", "7-day",  f"{_DIM}UNAVAILABLE{_RESET}"))
        else:
            fh = usage.get("five_hour")
            sd = usage.get("seven_day")
            rows.append(("Anthropic", "5-hour", _colour_pct(fh) if fh is not None else f"{_DIM}N/A{_RESET}"))
            rows.append(("Anthropic", "7-day",  _colour_pct(sd) if sd is not None else f"{_DIM}N/A{_RESET}"))
    except Exception as e:
        rows.append(("Anthropic", "5-hour", f"{_DIM}UNAVAILABLE{_RESET}"))
        rows.append(("Anthropic", "7-day",  f"{_DIM}UNAVAILABLE{_RESET}"))
        print(f"{_DIM}Warning: Anthropic probe failed: {e}{_RESET}", file=sys.stderr)

    # MiniMax
    try:
        key  = _read_minimax_key()
        usage, reason = probe_minimax(key)
        if usage is None:
            rows.append(("MiniMax", "5-hour", f"{_DIM}UNAVAILABLE{_RESET}"))
            rows.append(("MiniMax", "7-day",  f"{_DIM}UNAVAILABLE{_RESET}"))
            if reason:
                print(f"{_DIM}Warning: MiniMax probe: {reason}{_RESET}", file=sys.stderr)
        else:
            fh = usage.get("five_hour")
            sd = usage.get("seven_day")
            rows.append(("MiniMax", "5-hour", _colour_pct(fh) if fh is not None else f"{_DIM}N/A{_RESET}"))
            rows.append(("MiniMax", "7-day",  _colour_pct(sd) if sd is not None else f"{_DIM}N/A{_RESET}"))
    except Exception as e:
        rows.append(("MiniMax", "5-hour", f"{_DIM}UNAVAILABLE{_RESET}"))
        rows.append(("MiniMax", "7-day",  f"{_DIM}UNAVAILABLE{_RESET}"))
        print(f"{_DIM}Warning: MiniMax probe failed: {e}{_RESET}", file=sys.stderr)

    print_table(rows)


if __name__ == "__main__":
    main()