"""Server-Sent Events helper for streaming LLM responses over HTTP."""

from __future__ import annotations

from typing import TYPE_CHECKING

import orjson

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

    from koan.llm.client import LLMClient
    from koan.llm.config import LLMConfig
    from koan.types import Message


async def sse_response(
    client: LLMClient,
    config: LLMConfig,
    messages: list[Message],
) -> AsyncIterator[str]:
    """Yield Server-Sent Events (SSE) formatted strings from an LLM stream.

    Each yielded string is a complete SSE message (``data: ...\\n\\n``).
    The final message is ``data: [DONE]\\n\\n``.

    Suitable for use with FastAPI's ``StreamingResponse``::

        from fastapi.responses import StreamingResponse

        @app.post("/chat")
        async def chat(request: ChatRequest):
            return StreamingResponse(
                sse_response(client, config, messages),
                media_type="text/event-stream",
            )
    """
    async for delta in client.stream(config, messages):
        payload = {}
        if delta.text:
            payload["text"] = delta.text
        if delta.finish_reason:
            payload["finish_reason"] = delta.finish_reason
        if delta.usage:
            payload["usage"] = {
                "input_tokens": delta.usage.input_tokens,
                "output_tokens": delta.usage.output_tokens,
            }
        if delta.tool_call_id:
            payload["tool_call_id"] = delta.tool_call_id
        if delta.tool_name:
            payload["tool_name"] = delta.tool_name
        if delta.tool_args_delta:
            payload["tool_args_delta"] = delta.tool_args_delta

        if payload:
            data = orjson.dumps(payload).decode()
            yield f"data: {data}\n\n"

    yield "data: [DONE]\n\n"
