#!/usr/bin/env python3
"""AOCore streaming chat-completion via the OpenAI Python SDK.

Demonstrates:
  - stream=True to get SSE chunks
  - Iteration over the stream and incremental printing
  - Error handling for APIStatusError on 429/5xx with the X-AOSentry-RateLimit-Reset
    header

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

import os
import sys
import time

from openai import OpenAI, APIStatusError


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)

    try:
        stream = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "user", "content": "Count from 1 to 5, one number per line."},
            ],
            stream=True,
        )

        # Each chunk is a ChatCompletionChunk with choices[0].delta.content
        # holding the incremental token(s). Print without buffering so the
        # streaming effect is visible at the terminal.
        for chunk in stream:
            piece = chunk.choices[0].delta.content
            if piece:
                sys.stdout.write(piece)
                sys.stdout.flush()
        sys.stdout.write("\n")
        return 0

    except APIStatusError as e:
        # Handle rate-limit / budget / 5xx — see ../../errors.md for the
        # canonical retry policy.
        if e.status_code == 429:
            body = e.response.json().get("error", {})
            code = body.get("code", "")
            if code == "budget_exceeded":
                print("Budget exhausted — escalate to admin.", file=sys.stderr)
                return 2
            reset = int(e.response.headers.get("X-AOSentry-RateLimit-Reset", "0"))
            wait_s = max(0.0, reset - time.time())
            print(f"Rate limited; retry after {wait_s:.1f}s", file=sys.stderr)
            return 2
        print(f"API error status={e.status_code}: {e.response.text}", file=sys.stderr)
        return 2


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