import email.utils
import json
import os
import random
import time
from datetime import datetime, timezone
from pathlib import Path

import requests


def retry_delay(value, attempt):
    if value:
        try:
            seconds = float(value)
        except ValueError:
            try:
                date = email.utils.parsedate_to_datetime(value)
                if date.tzinfo is None:
                    date = date.replace(tzinfo=timezone.utc)
                seconds = (date - datetime.now(timezone.utc)).total_seconds()
            except (TypeError, ValueError):
                seconds = None
        if seconds is not None:
            if seconds > 60:
                raise RuntimeError("Retry-After exceeds 60s; retry manually later")
            return max(0, seconds)
    return min(30, 2 ** attempt) + random.uniform(0, 1)


def open_stream(body, headers):
    for attempt in range(3):
        response = requests.post(
            "https://api.quickrouter.ai/v1/chat/completions",
            json=body,
            headers=headers,
            stream=True,
            timeout=(10, 60),
        )
        if response.status_code == 429 and attempt < 2:
            retry_after = response.headers.get("Retry-After")
            response.close()
            time.sleep(retry_delay(retry_after, attempt))
            continue
        try:
            response.raise_for_status()
            if "text/event-stream" not in response.headers.get("Content-Type", "").lower():
                raise ValueError("Expected text/event-stream")
            return response
        except Exception:
            response.close()
            raise
    raise RuntimeError("Retry limit reached")


def sse_data(lines):
    parts = []
    for line in lines:
        if isinstance(line, bytes):
            line = line.decode("utf-8")
        if not line:
            if parts:
                yield "\n".join(parts)
                parts = []
        elif line.startswith("data:"):
            value = line[5:]
            parts.append(value[1:] if value.startswith(" ") else value)
    if parts:
        yield "\n".join(parts)


def consume(response, output):
    usage = None
    complete = False
    for data in sse_data(response.iter_lines()):
        if data == "[DONE]":
            complete = True
            break
        event = json.loads(data)
        if event.get("error"):
            raise RuntimeError(str(event["error"]))
        if event.get("usage"):
            usage = event["usage"]
        for choice in event.get("choices", []):
            text = choice.get("delta", {}).get("content")
            if text:
                output.write(text)
                output.flush()
                print(text, end="", flush=True)
    if not complete:
        raise RuntimeError("Stream ended without [DONE]; partial output kept, request not replayed")
    return usage


def main():
    headers = {"Authorization": "Bearer " + os.environ["QUICKROUTER_API_KEY"]}
    body = {
        "model": "gpt-4o",
        "messages": [{"role": "user", "content": "Explain API streaming in one short sentence."}],
        "max_tokens": 64,
        "stream": True,
        "stream_options": {"include_usage": True},
    }
    try:
        with open_stream(body, headers) as response, Path("partial-output.txt").open("w", encoding="utf-8") as output:
            usage = consume(response, output)
        print("\nusage:", usage)
    except (requests.RequestException, RuntimeError, ValueError) as error:
        print("\nStopped:", str(error))
        raise


if __name__ == "__main__":
    main()