"""Python 3 native Gemini SSE chat. Each submitted turn can incur API charges."""
import json
import os
from urllib.request import Request, urlopen

URL = "https://api.seedrouter.ai/v1beta/models/gemini-4-argon:streamGenerateContent?alt=sse"


def events(response):
    data = []
    for raw in response:
        line = raw.decode("utf-8").rstrip("\r\n")
        if line.startswith("data:"):
            data.append(line[5:].lstrip(" "))
        elif not line and data:
            yield json.loads("\n".join(data))
            data = []
    if data:
        yield json.loads("\n".join(data))


def turn(history, text, key):
    user = {"role": "user", "parts": [{"text": text}]}
    body = {
        "contents": [*history, user],
        "generationConfig": {
            "maxOutputTokens": 1024,
            "thinkingConfig": {"thinkingLevel": "low"},
        },
    }
    request = Request(
        URL,
        data=json.dumps(body).encode(),
        headers={
            "Authorization": f"Bearer {key}",
            "Content-Type": "application/json",
            "Accept": "text/event-stream",
        },
        method="POST",
    )
    parts, usage, finish = [], {}, None
    with urlopen(request, timeout=120) as response:
        for event in events(response):
            if "error" in event:
                raise RuntimeError("The stream returned an API error.")
            if "usageMetadata" in event:
                usage.update(event["usageMetadata"])
            for candidate in event.get("candidates", []):
                if candidate.get("index", 0) != 0:
                    continue
                for part in candidate.get("content", {}).get("parts", []):
                    parts.append(part)
                    if not part.get("thought") and "text" in part:
                        print(part["text"], end="", flush=True)
                finish = candidate.get("finishReason", finish)
    if finish != "STOP" or not parts:
        raise RuntimeError(f"Incomplete answer ({finish}); history was not changed.")
    print()
    return [*history, user, {"role": "model", "parts": parts}], usage


if __name__ == "__main__":
    key = os.environ["SEEDROUTER_API_KEY"]
    history = []
    while True:
        try:
            text = input("You (empty to exit): ").strip()
            if not text:
                break
            history, usage = turn(history, text, key)
            print("Usage:", json.dumps(usage))
        except (KeyboardInterrupt, EOFError):
            print("\nStopped. No further request was sent.")
            break
        except Exception as error:
            print(f"\n{type(error).__name__}: request did not complete; history unchanged.")