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

import logging
import time
from typing import Any

from ..context import reset_current_span
from ..types import (
    ANTHROPIC_KEY_CONTENT,
    META_CACHE_CREATION_TOKENS,
    META_CACHE_READ_TOKENS,
    META_RESPONSE_ID,
    META_STOP_REASON,
    META_TIME_TO_FIRST_TOKEN,
    SpanStatus,
)

logger = logging.getLogger("staso")


def _finalize_stream_span(
    span: Any,
    message_data: dict[str, Any],
    ctx_token: Any,
) -> None:
    """Finalize a streaming span with accumulated data. Idempotent.

    Always resets ctx_token, even if span was already ended.
    """
    try:
        if not span._recording:
            return
        try:
            span.output = {ANTHROPIC_KEY_CONTENT: message_data.get("content", [])}

            if message_data.get("stop_reason"):
                span.metadata[META_STOP_REASON] = message_data["stop_reason"]
            if message_data.get("id"):
                span.metadata[META_RESPONSE_ID] = message_data["id"]

            span.input_tokens = message_data.get("input_tokens", 0)
            span.output_tokens = message_data.get("output_tokens", 0)
            span.total_tokens = span.input_tokens + span.output_tokens

            cache_read = message_data.get("cache_read_input_tokens")
            if cache_read:
                span.metadata[META_CACHE_READ_TOKENS] = cache_read
            cache_creation = message_data.get("cache_creation_input_tokens")
            if cache_creation:
                span.metadata[META_CACHE_CREATION_TOKENS] = cache_creation

            content = message_data.get("content", [])
            if any(isinstance(b, dict) and b.get("type") == "thinking" for b in content):
                span.metadata["has_thinking"] = True

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


def _capture_message_response(span: Any, message: Any, ctx_token: Any) -> None:
    """Finalize span from a fully-accumulated Message object (from MessageStream).

    Always resets ctx_token, even if span was already ended.
    """
    try:
        if not span._recording:
            return
        try:
            from .anthropic import _capture_response
            _capture_response(span, message)

            content = getattr(message, "content", None)
            if content and any(getattr(b, "type", None) == "thinking" for b in content):
                span.metadata["has_thinking"] = True

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


class _EventAccumulator:
    """Accumulates raw stream events into a message data dict."""

    def __init__(self, span: Any, start_time: float) -> None:
        self._span = span
        self._start_time = start_time
        self._ttft_captured = False
        self.data: dict[str, Any] = {
            "content": [],
            "input_tokens": 0,
            "output_tokens": 0,
        }
        self._blocks: dict[int, dict] = {}

    def process(self, event: Any) -> None:
        event_type = getattr(event, "type", None)
        if not event_type:
            return

        if not self._ttft_captured and event_type == "content_block_delta":
            ttft_ms = (time.time() - self._start_time) * 1000
            self._span.metadata[META_TIME_TO_FIRST_TOKEN] = round(ttft_ms, 2)
            self._ttft_captured = True

        if event_type == "message_start":
            msg = getattr(event, "message", None)
            if msg:
                self.data["id"] = getattr(msg, "id", None)
                self.data["model"] = getattr(msg, "model", None)
                usage = getattr(msg, "usage", None)
                if usage:
                    self.data["input_tokens"] = getattr(usage, "input_tokens", 0) or 0
                    cr = getattr(usage, "cache_read_input_tokens", None)
                    if cr:
                        self.data["cache_read_input_tokens"] = cr
                    cc = getattr(usage, "cache_creation_input_tokens", None)
                    if cc:
                        self.data["cache_creation_input_tokens"] = cc

        elif event_type == "content_block_start":
            idx = getattr(event, "index", 0)
            block = getattr(event, "content_block", None)
            if block:
                btype = getattr(block, "type", "text")
                self._blocks[idx] = {"type": btype, "text": getattr(block, "text", "")}

        elif event_type == "content_block_delta":
            idx = getattr(event, "index", 0)
            delta = getattr(event, "delta", None)
            if delta and idx in self._blocks:
                dtype = getattr(delta, "type", "")
                if dtype == "text_delta":
                    self._blocks[idx]["text"] += getattr(delta, "text", "")
                elif dtype == "thinking_delta":
                    self._blocks[idx].setdefault("thinking", "")
                    self._blocks[idx]["thinking"] += getattr(delta, "thinking", "")

        elif event_type == "content_block_stop":
            idx = getattr(event, "index", 0)
            if idx in self._blocks:
                self.data["content"].append(self._blocks.pop(idx))

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


# ---------------------------------------------------------------------------
# Sync proxies
# ---------------------------------------------------------------------------

class StreamProxy:
    """Wraps Stream[RawMessageStreamEvent] 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 = _EventAccumulator(span, start_time)
        self._finalized = False

    def __iter__(self):
        try:
            for event in self._stream:
                self._acc.process(event)
                yield event
        except Exception as exc:
            self._span.record_exception(exc)
            self._end()
            raise
        finally:
            self._finalize()  # idempotent — safe after _end(), handles GeneratorExit

    def __enter__(self):
        self._stream.__enter__()
        return self

    def __exit__(self, *args):
        try:
            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_stream_span(self._span, self._acc.data, 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 MessageStreamProxy:
    """Wraps MessageStream from messages.stream() context."""

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

    def __iter__(self):
        for event in self._stream:
            self._track_ttft(event)
            yield event

    @property
    def text_stream(self):
        first = True
        for chunk in self._stream.text_stream:
            if first:
                ttft_ms = (time.time() - self._start_time) * 1000
                self._span.metadata[META_TIME_TO_FIRST_TOKEN] = round(ttft_ms, 2)
                first = False
            yield chunk

    def get_final_message(self):
        return self._stream.get_final_message()

    def get_final_text(self) -> str:
        return self._stream.get_final_text()

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

    def _track_ttft(self, event: Any) -> None:
        if not self._ttft_captured and getattr(event, "type", None) == "content_block_delta":
            ttft_ms = (time.time() - self._start_time) * 1000
            self._span.metadata[META_TIME_TO_FIRST_TOKEN] = round(ttft_ms, 2)
            self._ttft_captured = True

    def _finalize(self) -> None:
        if self._finalized:
            return
        self._finalized = True
        try:
            msg = self._stream.get_final_message()
            _capture_message_response(self._span, msg, self._ctx_token)
        except Exception:
            _finalize_stream_span(self._span, {}, self._ctx_token)

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


class MessageStreamManagerProxy:
    """Wraps MessageStreamManager from messages.stream()."""

    def __init__(self, manager: Any, span: Any, start_time: float, ctx_token: Any) -> None:
        self._manager = manager
        self._span = span
        self._start_time = start_time
        self._ctx_token = ctx_token
        self._proxy: MessageStreamProxy | None = None

    def __enter__(self):
        stream = self._manager.__enter__()
        self._proxy = MessageStreamProxy(stream, self._span, self._start_time, self._ctx_token)
        return self._proxy

    def __exit__(self, *args):
        if self._proxy:
            self._proxy._finalize()
        return self._manager.__exit__(*args)

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


# ---------------------------------------------------------------------------
# Async proxies
# ---------------------------------------------------------------------------

class AsyncStreamProxy:
    """Wraps AsyncStream[RawMessageStreamEvent] 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 = _EventAccumulator(span, start_time)
        self._finalized = False

    async def __aiter__(self):
        try:
            async for event in self._stream:
                self._acc.process(event)
                yield event
        except Exception as exc:
            self._span.record_exception(exc)
            self._end()
            raise
        finally:
            self._finalize()  # idempotent — safe after _end(), handles GeneratorExit

    async def __aenter__(self):
        await self._stream.__aenter__()
        return self

    async def __aexit__(self, *args):
        try:
            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_stream_span(self._span, self._acc.data, 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 AsyncMessageStreamProxy:
    """Wraps AsyncMessageStream."""

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

    async def __aiter__(self):
        async for event in self._stream:
            self._track_ttft(event)
            yield event

    async def text_stream_iter(self):
        first = True
        async for chunk in self._stream.text_stream:
            if first:
                ttft_ms = (time.time() - self._start_time) * 1000
                self._span.metadata[META_TIME_TO_FIRST_TOKEN] = round(ttft_ms, 2)
                first = False
            yield chunk

    async def get_final_message(self):
        return await self._stream.get_final_message()

    async def get_final_text(self) -> str:
        return await self._stream.get_final_text()

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

    def _track_ttft(self, event: Any) -> None:
        if not self._ttft_captured and getattr(event, "type", None) == "content_block_delta":
            ttft_ms = (time.time() - self._start_time) * 1000
            self._span.metadata[META_TIME_TO_FIRST_TOKEN] = round(ttft_ms, 2)
            self._ttft_captured = True

    def _finalize(self) -> None:
        if self._finalized:
            return
        self._finalized = True
        try:
            # For async, we can't await here; use sync path from accumulated data
            msg = self._stream.current_message_snapshot
            if msg:
                _capture_message_response(self._span, msg, self._ctx_token)
            else:
                _finalize_stream_span(self._span, {}, self._ctx_token)
        except Exception:
            _finalize_stream_span(self._span, {}, self._ctx_token)

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


class AsyncMessageStreamManagerProxy:
    """Wraps AsyncMessageStreamManager."""

    def __init__(self, manager: Any, span: Any, start_time: float, ctx_token: Any) -> None:
        self._manager = manager
        self._span = span
        self._start_time = start_time
        self._ctx_token = ctx_token
        self._proxy: AsyncMessageStreamProxy | None = None

    async def __aenter__(self):
        stream = await self._manager.__aenter__()
        self._proxy = AsyncMessageStreamProxy(stream, self._span, self._start_time, self._ctx_token)
        return self._proxy

    async def __aexit__(self, *args):
        if self._proxy:
            self._proxy._finalize()
        return await self._manager.__aexit__(*args)

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