"""Check whether the Groq daily token quota has enough headroom for a full run.

Makes one small completion that requests a moderate number of tokens. If it
succeeds, there is daily headroom; if it returns a 429 'tokens per day' error,
it prints how many tokens are still used vs the daily limit.

Usage:
    python check_quota.py
When it reports READY, run:
    $env:GAIA_BACKEND="groq"; python submit_official.py
"""

from __future__ import annotations

import os
import re
import sys

try:
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
except Exception:  # noqa: BLE001
    pass

from dotenv import load_dotenv

load_dotenv()

import requests

MODEL = os.getenv("GAIA_GROQ_MODEL", "openai/gpt-oss-120b")
# A full agentic run needs tens of thousands of tokens spread over many calls;
# we probe with a modest request as a readiness signal.
PROBE_MAX_TOKENS = 800


def main() -> None:
    key = os.getenv("GROQ_API_KEY")
    if not key:
        print("GROQ_API_KEY not set."); return

    resp = requests.post(
        "https://api.groq.com/openai/v1/chat/completions",
        headers={"Authorization": f"Bearer {key}"},
        json={
            "model": MODEL,
            "messages": [{"role": "user", "content": "Reply with the word ok."}],
            "max_tokens": PROBE_MAX_TOKENS,
        },
        timeout=30,
    )

    print(f"Model: {MODEL}")
    print(f"HTTP:  {resp.status_code}")

    if resp.status_code == 200:
        rem_min = resp.headers.get("x-ratelimit-remaining-tokens", "?")
        print(f"Per-minute tokens remaining: {rem_min} / "
              f"{resp.headers.get('x-ratelimit-limit-tokens', '?')}")
        print(
            f"\n==> A {PROBE_MAX_TOKENS}-token call SUCCEEDS, so the daily quota "
            "has at least some headroom."
        )
        print(
            "    NOTE: Groq does not expose daily-tokens-remaining on success, so\n"
            "    this only proves small calls work, not a full ~30-60k-token run.\n"
            "    If submit_official.py dies early with a 429 'tokens per day', the\n"
            "    daily budget has not recovered enough yet -- wait longer and retry."
        )
        print('\n    Try the run:  $env:GAIA_BACKEND="groq"; python submit_official.py')
        return

    body = resp.text
    print("Response:", body[:400])
    m = re.search(r"Used (\d+).*?Limit (\d+)|Limit (\d+).*?Used (\d+)", body)
    if "tokens per day" in body.lower() or "TPD" in body:
        print("\n==> NOT READY: daily token quota still exhausted. Wait longer.")
    elif "tokens per minute" in body.lower():
        print("\n==> Per-minute cap hit; daily may be fine. Wait ~60s and retry this check.")
    else:
        print("\n==> Not ready / unexpected error above.")


if __name__ == "__main__":
    main()