"""Demo streaming agent — tests Litmus proxy with streaming SSE responses."""

from __future__ import annotations

import argparse
import json
import sys


def run_anthropic_stream(prompt: str):
    """Stream a response from Anthropic SDK."""
    import anthropic

    client = anthropic.Anthropic()
    print(f"[Anthropic Streaming] Prompt: {prompt}")
    print("-" * 60)

    collected_text = ""
    with client.messages.stream(
        model="claude-sonnet-4-20250514",
        max_tokens=150,
        messages=[{"role": "user", "content": prompt}],
    ) as stream:
        for text in stream.text_stream:
            collected_text += text
            print(text, end="", flush=True)

    print(f"\n\n[Complete] {len(collected_text)} chars received")
    print("-" * 60)
    return collected_text


def run_openai_stream(prompt: str):
    """Stream a response from OpenAI SDK."""
    import openai

    client = openai.OpenAI()
    print(f"[OpenAI Streaming] Prompt: {prompt}")
    print("-" * 60)

    collected_text = ""
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        max_tokens=150,
        messages=[{"role": "user", "content": prompt}],
        stream=True,
    )

    for chunk in response:
        if chunk.choices and chunk.choices[0].delta.content:
            text = chunk.choices[0].delta.content
            collected_text += text
            print(text, end="", flush=True)

    print(f"\n\n[Complete] {len(collected_text)} chars received")
    print("-" * 60)
    return collected_text


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--provider", choices=["anthropic", "openai"], default="anthropic")
    parser.add_argument("--prompt", default="Count from 1 to 10, one number per line.")
    args = parser.parse_args()

    if args.provider == "anthropic":
        run_anthropic_stream(args.prompt)
    else:
        run_openai_stream(args.prompt)


if __name__ == "__main__":
    main()
