"""
SSE streaming support for the kvict Proxy.

Handles streaming responses from LLM providers:
1. Forwards each SSE chunk to the client immediately (low latency).
2. Buffers the full response in parallel for caching + token counting.
3. Returns aggregated usage data after the stream completes.
"""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass, field
from typing import Any, AsyncIterator, Dict, List, Optional

import httpx

from proxy.adapters.base import ProviderAdapter
from proxy.schemas import (
    ChatCompletionChunk,
    ChatCompletionResponse,
    Choice,
    ChoiceMessage,
    Usage,
)

logger = logging.getLogger(__name__)


@dataclass
class StreamResult:
    """Aggregated result after a stream completes."""

    full_content: str = ""
    completion_tokens: int = 0
    prompt_tokens: int = 0
    model: str = ""
    response_id: str = ""
    finish_reason: Optional[str] = None
    chunks_received: int = 0
    # Usage from the final chunk (OpenAI sends it with stream_options.include_usage)
    usage: Optional[Usage] = None


async def stream_provider_response(
    response: httpx.Response,
    adapter: ProviderAdapter,
) -> AsyncIterator[tuple[str, StreamResult]]:
    """Stream an SSE response from the provider, yielding each line.

    Yields:
        Tuples of (sse_line, accumulated_result).
        The final yield contains the complete StreamResult.

    Usage:
        async for sse_line, result in stream_provider_response(response, adapter):
            # Forward sse_line to client
            yield sse_line
        # result now contains full content and token counts
    """
    result = StreamResult()
    content_parts: List[str] = []

    async for raw_line in response.aiter_lines():
        line = raw_line.strip()
        if not line:
            continue

        # Forward the raw SSE line as-is
        sse_line = line

        # Parse if it's a data line
        if line.startswith("data: "):
            data_str = line[6:]  # Strip "data: " prefix

            if data_str.strip() == "[DONE]":
                yield "data: [DONE]\n\n", result
                break

            chunk = adapter.parse_stream_chunk(data_str)
            if chunk is not None:
                result.chunks_received += 1
                result.model = chunk.model or result.model
                result.response_id = chunk.id or result.response_id

                for choice in chunk.choices:
                    if choice.delta.content:
                        content_parts.append(choice.delta.content)
                    if choice.finish_reason:
                        result.finish_reason = choice.finish_reason

            # Check for usage in the chunk (OpenAI: usage; Gemini: usageMetadata)
            try:
                raw_data = json.loads(data_str)
                if "usage" in raw_data and raw_data["usage"]:
                    usage_data = raw_data["usage"]
                    result.usage = Usage(
                        prompt_tokens=usage_data.get("prompt_tokens", 0),
                        completion_tokens=usage_data.get("completion_tokens", 0),
                        total_tokens=usage_data.get("total_tokens", 0),
                    )
                elif "usageMetadata" in raw_data and raw_data["usageMetadata"]:
                    um = raw_data["usageMetadata"]
                    pt = um.get("promptTokenCount", 0)
                    ct = um.get("candidatesTokenCount", 0)
                    result.usage = Usage(
                        prompt_tokens=pt,
                        completion_tokens=ct,
                        total_tokens=um.get("totalTokenCount", pt + ct),
                    )
            except (json.JSONDecodeError, KeyError):
                pass

        yield f"{sse_line}\n\n", result

    # Finalize
    result.full_content = "".join(content_parts)
    if result.usage:
        result.prompt_tokens = result.usage.prompt_tokens
        result.completion_tokens = result.usage.completion_tokens
    else:
        # Estimate completion tokens from content (fallback)
        result.completion_tokens = max(1, len(result.full_content) // 4)


def stream_result_to_response(result: StreamResult) -> ChatCompletionResponse:
    """Convert a completed StreamResult to a ChatCompletionResponse for caching."""
    return ChatCompletionResponse(
        id=result.response_id,
        model=result.model,
        choices=[
            Choice(
                index=0,
                message=ChoiceMessage(
                    role="assistant",
                    content=result.full_content,
                ),
                finish_reason=result.finish_reason or "stop",
            )
        ],
        usage=result.usage
        or Usage(
            prompt_tokens=result.prompt_tokens,
            completion_tokens=result.completion_tokens,
            total_tokens=result.prompt_tokens + result.completion_tokens,
        ),
    )
