#!/usr/bin/env python3
"""AOCore chat-completion via the OpenAI Python SDK — drop-in compatible.

The only AOCore-specific configuration is the base_url. Everything else is
the standard OpenAI SDK call shape.

Run:
    pip install openai
    export AOSENTRY_KEY="sk-..."
    python chat.py
"""

import os
import sys

from openai import OpenAI


def main() -> int:
    api_key = os.environ.get("AOSENTRY_KEY")
    if not api_key:
        print("Set AOSENTRY_KEY to your minted API key (sk-...).", file=sys.stderr)
        return 1

    base_url = os.environ.get("AOSENTRY_BASE_URL", "https://gateway.core.aocyber.ai/v1")

    client = OpenAI(base_url=base_url, api_key=api_key)

    # with_raw_response gives us access to response.headers so we can read the
    # X-AOCore-* quota headers — see ../../rate-limits.md.
    raw = client.chat.completions.with_raw_response.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": "You are a concise assistant."},
            {"role": "user", "content": "Say hello in one sentence."},
        ],
        temperature=0.2,
        max_tokens=64,
    )

    completion = raw.parse()
    print(completion.choices[0].message.content)

    # AOCore-specific quota headers — see ../../rate-limits.md for semantics.
    quota_headers = (
        "X-AOSentry-RateLimit-Remaining",
        "X-AOSentry-RPM-Remaining",
        "X-AOSentry-TPM-Remaining",
        "X-AOSentry-Budget-Remaining",
        "X-AOSentry-RateLimit-Reset",
        "X-AOSentry-Budget-Reset",
    )
    print("\n-- Quota --")
    for header in quota_headers:
        if header in raw.headers:
            print(f"{header}: {raw.headers[header]}")

    return 0


if __name__ == "__main__":
    raise SystemExit(main())