"""Streaming chat usage examples

Demonstrates SSE streaming responses from the Octen SDK.
"""

from octen import Octen
from octen.models import ChatMessage, WebSearchOptions


def example_basic_stream():
    """Stream content tokens to stdout in real time."""
    print("=== Basic Streaming ===\n")
    with Octen(api_key="octen-your-key") as client:
        print("Answer: ", end="", flush=True)
        for event in client.chat.create(
            model="openai/gpt-5.4",
            messages=[ChatMessage(role="user", content="Explain quantum entanglement in three sentences")],
            stream=True,
            web_search="off",
        ):
            if event.type == "content" and event.choices:
                print(event.choices[0].delta.content or "", end="", flush=True)
            elif event.type == "finish" and event.choices:
                print(f"\n[finish] reason={event.choices[0].finish_reason}")
            elif event.type == "usage" and event.usage:
                u = event.usage
                print(f"[usage] prompt={u.prompt_tokens} "
                      f"completion={u.completion_tokens} "
                      f"total={u.total_tokens}")


def example_stream_with_web_search():
    """Stream with web search -- handle search_done and content events."""
    print("\n=== Streaming with Web Search ===\n")
    with Octen(api_key="octen-your-key") as client:
        for event in client.chat.create(
            model="openai/gpt-5.4",
            messages=[ChatMessage(role="user", content="What are the latest AI breakthroughs?")],
            stream=True,
            web_search="on",
            web_search_options=WebSearchOptions(safesearch="off"),
        ):
            if event.type == "search_done":
                total = sum(len(g.results) for g in (event.search_results or []))
                print(f"[search_done] {total} results found\n")
                print("Answer: ", end="", flush=True)

            elif event.type == "content" and event.choices:
                print(event.choices[0].delta.content or "", end="", flush=True)

            elif event.type == "finish":
                print()

            elif event.type == "usage" and event.usage:
                print(f"\n[usage] search_queries={event.usage.num_search_queries} "
                      f"total_tokens={event.usage.total_tokens}")


def example_collect_full_response():
    """Collect all stream chunks and print assembled text."""
    print("\n=== Collect Full Streaming Response ===\n")
    with Octen(api_key="octen-your-key") as client:
        chunks = []
        finish_reason = None
        usage = None

        for event in client.chat.create(
            model="anthropic/claude-sonnet-4.6",
            messages=[ChatMessage(role="user", content="Write a four-line poem about Python")],
            stream=True,
            web_search="off",
        ):
            if event.type == "content" and event.choices:
                chunk = event.choices[0].delta.content or ""
                chunks.append(chunk)
            elif event.type == "finish" and event.choices:
                finish_reason = event.choices[0].finish_reason
            elif event.type == "usage":
                usage = event.usage

        full_text = "".join(chunks)
        print(f"Full text:\n{full_text}")
        print(f"\nFinish reason: {finish_reason}")
        if usage:
            print(f"Total tokens:  {usage.total_tokens}")


if __name__ == "__main__":
    example_basic_stream()
    example_stream_with_web_search()
    example_collect_full_response()
