from __future__ import annotations

"""Streaming helpers."""

import json

from .errors import AgentError
from .types import StopReason, StreamChunk, ToolCallDelta


def parse_sse_chunk(payload: bytes) -> StreamChunk | None:
    """Parse a single SSE byte chunk into a unified stream chunk."""

    try:
        text = payload.decode("utf-8")
    except UnicodeDecodeError as error:
        raise AgentError.invalid_stream() from error

    for line in text.splitlines():
        if not line.startswith("data:"):
            continue
        data = line[5:].strip()
        if not data:
            continue
        if data == "[DONE]":
            return None
        try:
            value = json.loads(data)
        except json.JSONDecodeError as error:
            raise AgentError.invalid_stream() from error
        return _map_openai_like_chunk(value)
    return None


def _map_openai_like_chunk(value: dict) -> StreamChunk:
    choice = ((value.get("choices") or [{}]) or [{}])[0] or {}
    delta_value = choice.get("delta") or {}
    delta = delta_value.get("content") or ""

    tool_call_delta = None
    items = delta_value.get("tool_calls")
    if isinstance(items, list) and items:
        tool_call_delta = []
        for item in items:
            function = item.get("function") or {}
            tool_call_delta.append(
                ToolCallDelta(
                    index=int(item.get("index") or 0),
                    id=item.get("id"),
                    name=function.get("name"),
                    arguments_delta=function.get("arguments"),
                )
            )

    finish_reason = None
    if choice.get("finish_reason") is not None:
        finish_reason = map_stop_reason(str(choice["finish_reason"]))

    return StreamChunk(
        delta=delta, tool_call_delta=tool_call_delta, finish_reason=finish_reason
    )


def map_stop_reason(value: str) -> StopReason:
    """Map a provider stop reason into the unified enum."""

    if value in {"stop", "end_turn"}:
        return StopReason.STOP
    if value in {"tool_use", "tool_calls"}:
        return StopReason.TOOL_USE
    if value in {"length", "max_tokens"}:
        return StopReason.MAX_TOKENS
    return StopReason.OTHER
