"""Streaming wrapper — captures chunks, computes TTFT, reassembles response."""

from __future__ import annotations

import logging
import time
from typing import Any, Iterator

logger = logging.getLogger("llm_observatory")


class StreamWrapper:
    """Wraps an OpenAI streaming response to capture metrics and reassemble the response."""

    def __init__(
        self,
        stream: Any,
        start_time: float,
        request_kwargs: dict[str, Any],
        on_complete: Any,
        on_error: Any = None,
    ) -> None:
        self._stream = stream
        self._start_time = start_time
        self._request_kwargs = request_kwargs
        self._on_complete = on_complete
        self._on_error = on_error
        self._chunks: list[Any] = []
        self._first_token_time: float | None = None
        self._model: str = ""
        self._response_id: str = ""
        self._finalized = False

    def __iter__(self) -> Iterator[Any]:
        return self

    def __next__(self) -> Any:
        try:
            chunk = next(self._stream)
        except StopIteration:
            self._finalize()
            raise
        except Exception as exc:
            self._capture_error(exc)
            raise

        self._track_chunk(chunk)
        return chunk

    def __aiter__(self) -> "StreamWrapper":
        return self

    async def __anext__(self) -> Any:
        try:
            chunk = await self._stream.__anext__()
        except StopAsyncIteration:
            self._finalize()
            raise
        except Exception as exc:
            self._capture_error(exc)
            raise

        self._track_chunk(chunk)
        return chunk

    def _track_chunk(self, chunk: Any) -> None:
        if self._first_token_time is None:
            self._first_token_time = time.time()

        self._chunks.append(chunk)

        if not self._model and hasattr(chunk, "model") and chunk.model:
            self._model = chunk.model
        if not self._response_id and hasattr(chunk, "id") and chunk.id:
            self._response_id = chunk.id

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

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        if hasattr(self._stream, "__exit__"):
            self._stream.__exit__(exc_type, exc_val, exc_tb)
        if exc_val is not None:
            self._capture_error(exc_val)
        else:
            self._finalize()

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

    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        if hasattr(self._stream, "__aexit__"):
            await self._stream.__aexit__(exc_type, exc_val, exc_tb)
        if exc_val is not None:
            self._capture_error(exc_val)
        else:
            self._finalize()

    def __del__(self) -> None:
        try:
            self._finalize()
        except Exception:
            pass

    def _capture_error(self, exc: Exception) -> None:
        """Record an error event for a mid-stream failure. Never raises."""
        if self._finalized:
            return
        self._finalized = True
        try:
            if self._on_error is not None:
                self._on_error(exc=exc, end_time=time.time())
        except Exception:
            logger.debug("llm_observatory: failed to capture stream error event")

    def _finalize(self) -> None:
        if self._finalized or not self._chunks:
            return
        self._finalized = True

        end_time = time.time()

        # Compute TTFT
        ttft_ms: int | None = None
        if self._first_token_time is not None:
            ttft_ms = int((self._first_token_time - self._start_time) * 1000)

        # Reassemble content from chunks
        content_parts: list[str] = []
        reasoning_parts: list[str] = []
        finish_reason = ""
        usage = None

        for chunk in self._chunks:
            if hasattr(chunk, "choices") and chunk.choices:
                delta = getattr(chunk.choices[0], "delta", None)
                if delta:
                    if getattr(delta, "content", None):
                        content_parts.append(delta.content)
                    if getattr(delta, "reasoning_content", None):
                        reasoning_parts.append(delta.reasoning_content)
                fr = getattr(chunk.choices[0], "finish_reason", None)
                if fr:
                    finish_reason = fr
            if hasattr(chunk, "usage") and chunk.usage is not None:
                usage = chunk.usage

        content = "".join(content_parts)
        reasoning_content = "".join(reasoning_parts) if reasoning_parts else None

        # Compute tokens/sec
        tps: float | None = None
        completion_tokens = 0
        if usage is not None:
            completion_tokens = getattr(usage, "completion_tokens", 0) or 0
        if completion_tokens > 0 and self._first_token_time is not None:
            generation_time = end_time - self._first_token_time
            if generation_time > 0:
                tps = completion_tokens / generation_time

        # Build a pseudo-response object for event construction
        response = _ReassembledResponse(
            id=self._response_id,
            model=self._model,
            content=content,
            finish_reason=finish_reason,
            usage=usage,
            reasoning_content=reasoning_content,
        )

        try:
            self._on_complete(
                response=response,
                end_time=end_time,
                time_to_first_token_ms=ttft_ms,
                tokens_per_second=tps,
            )
        except Exception:
            logger.debug("llm_observatory: failed to finalize stream event")


class _ReassembledResponse:
    """Minimal response-like object for event building from streamed chunks."""

    def __init__(
        self,
        id: str,
        model: str,
        content: str,
        finish_reason: str,
        usage: Any,
        reasoning_content: str | None = None,
    ) -> None:
        self.id = id
        self.model = model
        self.choices = [_ReassembledChoice(content, finish_reason, reasoning_content)]
        self.usage = usage


class _ReassembledChoice:
    def __init__(self, content: str, finish_reason: str, reasoning_content: str | None = None) -> None:
        self.index = 0
        self.message = _ReassembledMessage(content, reasoning_content)
        self.finish_reason = finish_reason


class _ReassembledMessage:
    def __init__(self, content: str, reasoning_content: str | None = None) -> None:
        self.role = "assistant"
        self.content = content
        self.reasoning_content = reasoning_content
        self.tool_calls = None


# ── Anthropic streaming ──────────────────────────────────────────


class AnthropicStreamWrapper:
    """Wraps an Anthropic streaming response to capture metrics and reassemble."""

    def __init__(
        self,
        stream: Any,
        start_time: float,
        request_kwargs: dict[str, Any],
        on_complete: Any,
        on_error: Any = None,
    ) -> None:
        self._stream = stream
        self._start_time = start_time
        self._request_kwargs = request_kwargs
        self._on_complete = on_complete
        self._on_error = on_error
        self._first_token_time: float | None = None
        self._finalized = False

        # Accumulation state
        self._response_id: str = ""
        self._model: str = ""
        self._content_parts: list[str] = []
        self._thinking_parts: list[str] = []
        self._stop_reason: str = ""
        self._input_tokens: int = 0
        self._output_tokens: int = 0
        self._has_events = False

    def __iter__(self) -> Iterator[Any]:
        return self

    def __next__(self) -> Any:
        try:
            event = next(self._stream)
        except StopIteration:
            self._finalize()
            raise
        except Exception as exc:
            self._capture_error(exc)
            raise

        self._track_event(event)
        return event

    def __aiter__(self) -> "AnthropicStreamWrapper":
        return self

    async def __anext__(self) -> Any:
        try:
            event = await self._stream.__anext__()
        except StopAsyncIteration:
            self._finalize()
            raise
        except Exception as exc:
            self._capture_error(exc)
            raise

        self._track_event(event)
        return event

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

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        if hasattr(self._stream, "__exit__"):
            self._stream.__exit__(exc_type, exc_val, exc_tb)
        if exc_val is not None:
            self._capture_error(exc_val)
        else:
            self._finalize()

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

    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        if hasattr(self._stream, "__aexit__"):
            await self._stream.__aexit__(exc_type, exc_val, exc_tb)
        if exc_val is not None:
            self._capture_error(exc_val)
        else:
            self._finalize()

    def __del__(self) -> None:
        try:
            self._finalize()
        except Exception:
            pass

    def _track_event(self, event: Any) -> None:
        self._has_events = True
        event_type = getattr(event, "type", "")

        if event_type == "message_start":
            message = getattr(event, "message", None)
            if message:
                self._response_id = getattr(message, "id", "") or ""
                self._model = getattr(message, "model", "") or ""
                usage = getattr(message, "usage", None)
                if usage:
                    self._input_tokens = getattr(usage, "input_tokens", 0) or 0

        elif event_type == "content_block_delta":
            if self._first_token_time is None:
                self._first_token_time = time.time()
            delta = getattr(event, "delta", None)
            if delta:
                delta_type = getattr(delta, "type", "")
                if delta_type == "text_delta":
                    self._content_parts.append(getattr(delta, "text", ""))
                elif delta_type == "thinking_delta":
                    self._thinking_parts.append(getattr(delta, "thinking", ""))

        elif event_type == "message_delta":
            delta = getattr(event, "delta", None)
            if delta:
                sr = getattr(delta, "stop_reason", None)
                if sr:
                    self._stop_reason = sr
            usage = getattr(event, "usage", None)
            if usage:
                self._output_tokens = getattr(usage, "output_tokens", 0) or 0

    def _capture_error(self, exc: Exception) -> None:
        if self._finalized:
            return
        self._finalized = True
        try:
            if self._on_error is not None:
                self._on_error(exc=exc, end_time=time.time())
        except Exception:
            logger.debug("llm_observatory: failed to capture anthropic stream error")

    def _finalize(self) -> None:
        if self._finalized or not self._has_events:
            return
        self._finalized = True

        end_time = time.time()

        # TTFT
        ttft_ms: int | None = None
        if self._first_token_time is not None:
            ttft_ms = int((self._first_token_time - self._start_time) * 1000)

        content = "".join(self._content_parts)
        thinking = "".join(self._thinking_parts) if self._thinking_parts else None

        # tokens/sec
        tps: float | None = None
        if self._output_tokens > 0 and self._first_token_time is not None:
            gen_time = end_time - self._first_token_time
            if gen_time > 0:
                tps = self._output_tokens / gen_time

        total_tokens = self._input_tokens + self._output_tokens
        usage = _AnthropicReassembledUsage(
            self._input_tokens, self._output_tokens, total_tokens
        )

        response = _AnthropicReassembledResponse(
            id=self._response_id,
            model=self._model,
            content=content,
            stop_reason=self._stop_reason,
            usage=usage,
            thinking=thinking,
        )

        try:
            self._on_complete(
                response=response,
                end_time=end_time,
                time_to_first_token_ms=ttft_ms,
                tokens_per_second=tps,
            )
        except Exception:
            logger.debug("llm_observatory: failed to finalize anthropic stream event")


class _AnthropicReassembledResponse:
    """Mimics an anthropic.types.Message for event building."""

    def __init__(
        self,
        id: str,
        model: str,
        content: str,
        stop_reason: str,
        usage: Any,
        thinking: str | None = None,
    ) -> None:
        self.id = id
        self.model = model
        self.stop_reason = stop_reason
        self.usage = usage
        # Build content blocks matching Anthropic's format
        blocks: list[Any] = []
        if thinking:
            blocks.append(_AnthropicTextBlock("thinking", thinking=thinking))
        if content:
            blocks.append(_AnthropicTextBlock("text", text=content))
        self.content = blocks


class _AnthropicTextBlock:
    def __init__(self, type: str, text: str = "", thinking: str = "") -> None:
        self.type = type
        self.text = text
        self.thinking = thinking


class _AnthropicReassembledUsage:
    def __init__(self, input_tokens: int, output_tokens: int, total_tokens: int) -> None:
        self.input_tokens = input_tokens
        self.output_tokens = output_tokens
        self.total_tokens = total_tokens
        self.cache_read_input_tokens = None
        self.cache_creation_input_tokens = None
