"""Streaming proxy classes for OpenAI SDK instrumentation."""
from __future__ import annotations

import logging
import time
from typing import Any

from ..context import reset_current_span
from ..types import (
    META_REASONING_TOKENS,
    META_RESPONSE_ID,
    META_STOP_REASON,
    META_TIME_TO_FIRST_TOKEN,
    OPENAI_KEY_CONTENT,
    OPENAI_KEY_TOOL_CALLS,
    SpanStatus,
)

logger = logging.getLogger("staso")


class _OpenAIEventAccumulator:
    """Accumulates OpenAI streaming ChatCompletionChunk events."""

    def __init__(self, span: Any, start_time: float) -> None:
        self._span = span
        self._start_time = start_time
        self._ttft_captured = False
        self.content: str = ""
        self.tool_calls: list[dict] = []
        self.finish_reason: str | None = None
        self.usage: dict[str, Any] | None = None
        self.model: str = ""
        self.response_id: str = ""
        self._tool_call_buffers: dict[int, dict[str, Any]] = {}

    def process(self, chunk: Any) -> None:
        chunk_id = getattr(chunk, "id", None)
        if chunk_id and not self.response_id:
            self.response_id = chunk_id

        chunk_model = getattr(chunk, "model", None)
        if chunk_model and not self.model:
            self.model = chunk_model

        choices = getattr(chunk, "choices", None)
        if choices and len(choices) > 0:
            delta = getattr(choices[0], "delta", None)
            if delta:
                delta_content = getattr(delta, "content", None)
                if delta_content:
                    if not self._ttft_captured:
                        ttft_ms = (time.time() - self._start_time) * 1000
                        self._span.metadata[META_TIME_TO_FIRST_TOKEN] = round(ttft_ms, 2)
                        self._ttft_captured = True
                    self.content += delta_content

                delta_tool_calls = getattr(delta, "tool_calls", None)
                if delta_tool_calls:
                    for tc in delta_tool_calls:
                        idx = getattr(tc, "index", 0)
                        if idx not in self._tool_call_buffers:
                            self._tool_call_buffers[idx] = {
                                "id": getattr(tc, "id", ""),
                                "type": getattr(tc, "type", "function"),
                                "function": {
                                    "name": "",
                                    "arguments": "",
                                },
                            }
                        buf = self._tool_call_buffers[idx]
                        tc_id = getattr(tc, "id", None)
                        if tc_id:
                            buf["id"] = tc_id
                        func = getattr(tc, "function", None)
                        if func:
                            fn_name = getattr(func, "name", None)
                            if fn_name:
                                buf["function"]["name"] = fn_name
                            fn_args = getattr(func, "arguments", None)
                            if fn_args:
                                buf["function"]["arguments"] += fn_args

            finish = getattr(choices[0], "finish_reason", None)
            if finish:
                self.finish_reason = finish

        chunk_usage = getattr(chunk, "usage", None)
        if chunk_usage:
            self.usage = {
                "prompt_tokens": getattr(chunk_usage, "prompt_tokens", 0) or 0,
                "completion_tokens": getattr(chunk_usage, "completion_tokens", 0) or 0,
                "total_tokens": getattr(chunk_usage, "total_tokens", 0) or 0,
            }
            details = getattr(chunk_usage, "completion_tokens_details", None)
            if details:
                reasoning = getattr(details, "reasoning_tokens", None)
                if reasoning:
                    self.usage["reasoning_tokens"] = reasoning

    def finalize_tool_calls(self) -> None:
        if self._tool_call_buffers:
            self.tool_calls = [
                self._tool_call_buffers[k]
                for k in sorted(self._tool_call_buffers)
            ]


def _finalize_openai_span(
    span: Any,
    acc: _OpenAIEventAccumulator,
    ctx_token: Any,
) -> None:
    """Finalize an OpenAI streaming span with accumulated data. Idempotent."""
    try:
        if not span._recording:
            return
        try:
            acc.finalize_tool_calls()

            output: dict[str, Any] = {OPENAI_KEY_CONTENT: acc.content}
            if acc.tool_calls:
                output[OPENAI_KEY_TOOL_CALLS] = acc.tool_calls
            span.output = output

            if acc.finish_reason:
                span.metadata[META_STOP_REASON] = acc.finish_reason
            if acc.response_id:
                span.metadata[META_RESPONSE_ID] = acc.response_id

            if acc.usage:
                span.input_tokens = acc.usage.get("prompt_tokens", 0)
                span.output_tokens = acc.usage.get("completion_tokens", 0)
                span.total_tokens = acc.usage.get("total_tokens", 0)
                reasoning = acc.usage.get("reasoning_tokens")
                if reasoning:
                    span.metadata[META_REASONING_TOKENS] = reasoning

            span.status = SpanStatus.OK
        except Exception:
            logger.debug("staso: failed to finalize openai stream span", exc_info=True)
        finally:
            span.end()
    finally:
        if ctx_token is not None:
            reset_current_span(ctx_token)


class OpenAIStreamProxy:
    """Wraps sync Stream[ChatCompletionChunk] from create(stream=True)."""

    def __init__(self, stream: Any, span: Any, start_time: float, ctx_token: Any) -> None:
        self._stream = stream
        self._span = span
        self._ctx_token = ctx_token
        self._acc = _OpenAIEventAccumulator(span, start_time)
        self._finalized = False

    def __iter__(self):
        try:
            for chunk in self._stream:
                self._acc.process(chunk)
                yield chunk
        except Exception as exc:
            self._span.record_exception(exc)
            self._end()
            raise
        finally:
            self._finalize()

    def __enter__(self):
        if hasattr(self._stream, "__enter__"):
            self._stream.__enter__()
        return self

    def __exit__(self, *args):
        try:
            if hasattr(self._stream, "__exit__"):
                return self._stream.__exit__(*args)
        finally:
            self._finalize()

    def close(self) -> None:
        self._finalize()
        if hasattr(self._stream, "close"):
            self._stream.close()

    def _finalize(self) -> None:
        if self._finalized:
            return
        self._finalized = True
        _finalize_openai_span(self._span, self._acc, self._ctx_token)

    def _end(self) -> None:
        if self._finalized:
            return
        self._finalized = True
        self._span.end()
        if self._ctx_token is not None:
            reset_current_span(self._ctx_token)

    def __getattr__(self, name: str) -> Any:
        return getattr(self._stream, name)


class AsyncOpenAIStreamProxy:
    """Wraps async AsyncStream[ChatCompletionChunk] from create(stream=True)."""

    def __init__(self, stream: Any, span: Any, start_time: float, ctx_token: Any) -> None:
        self._stream = stream
        self._span = span
        self._ctx_token = ctx_token
        self._acc = _OpenAIEventAccumulator(span, start_time)
        self._finalized = False

    async def __aiter__(self):
        try:
            async for chunk in self._stream:
                self._acc.process(chunk)
                yield chunk
        except Exception as exc:
            self._span.record_exception(exc)
            self._end()
            raise
        finally:
            self._finalize()

    async def __aenter__(self):
        if hasattr(self._stream, "__aenter__"):
            await self._stream.__aenter__()
        return self

    async def __aexit__(self, *args):
        try:
            if hasattr(self._stream, "__aexit__"):
                return await self._stream.__aexit__(*args)
        finally:
            self._finalize()

    async def close(self) -> None:
        self._finalize()
        if hasattr(self._stream, "close"):
            await self._stream.close()

    def _finalize(self) -> None:
        if self._finalized:
            return
        self._finalized = True
        _finalize_openai_span(self._span, self._acc, self._ctx_token)

    def _end(self) -> None:
        if self._finalized:
            return
        self._finalized = True
        self._span.end()
        if self._ctx_token is not None:
            reset_current_span(self._ctx_token)

    def __getattr__(self, name: str) -> Any:
        return getattr(self._stream, name)
