"""Tests for staso.integrations._streaming proxy classes."""
from __future__ import annotations

import asyncio
import json
import time
from types import SimpleNamespace
from unittest.mock import MagicMock

import pytest

from staso.context import set_current_span
from staso.integrations._streaming import (
    AsyncStreamProxy,
    MessageStreamManagerProxy,
    MessageStreamProxy,
    StreamProxy,
    _EventAccumulator,
    _finalize_stream_span,
)
from staso.provider import Span
from staso.types import SpanKind, SpanStatus


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _make_span(transport=None):
    """Create a test Span with deterministic IDs."""
    return Span(
        name="anthropic.messages.create",
        trace_id="aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee",
        span_id="11111111-2222-3333-4444-555555555555",
        parent_span_id=None,
        transport=transport,
        agent_id="test-agent",
        session_id="test-session",
        kind=SpanKind.LLM,
    )


def _make_stream_events():
    """Generate a realistic sequence of Anthropic stream events."""
    return [
        SimpleNamespace(type="message_start", message=SimpleNamespace(
            id="msg_123", model="claude-test",
            usage=SimpleNamespace(input_tokens=50, cache_read_input_tokens=None, cache_creation_input_tokens=None),
        )),
        SimpleNamespace(type="content_block_start", index=0, content_block=SimpleNamespace(type="text", text="")),
        SimpleNamespace(type="content_block_delta", index=0, delta=SimpleNamespace(type="text_delta", text="Hello")),
        SimpleNamespace(type="content_block_delta", index=0, delta=SimpleNamespace(type="text_delta", text=" world")),
        SimpleNamespace(type="content_block_stop", index=0),
        SimpleNamespace(type="message_delta", delta=SimpleNamespace(stop_reason="end_turn"), usage=SimpleNamespace(output_tokens=25)),
        SimpleNamespace(type="message_stop"),
    ]


def _make_thinking_events():
    """Generate stream events that include a thinking block."""
    return [
        SimpleNamespace(type="message_start", message=SimpleNamespace(
            id="msg_456", model="claude-test",
            usage=SimpleNamespace(input_tokens=30, cache_read_input_tokens=None, cache_creation_input_tokens=None),
        )),
        SimpleNamespace(type="content_block_start", index=0, content_block=SimpleNamespace(type="thinking", text="")),
        SimpleNamespace(type="content_block_delta", index=0, delta=SimpleNamespace(type="thinking_delta", thinking="Let me think...")),
        SimpleNamespace(type="content_block_stop", index=0),
        SimpleNamespace(type="content_block_start", index=1, content_block=SimpleNamespace(type="text", text="")),
        SimpleNamespace(type="content_block_delta", index=1, delta=SimpleNamespace(type="text_delta", text="Answer")),
        SimpleNamespace(type="content_block_stop", index=1),
        SimpleNamespace(type="message_delta", delta=SimpleNamespace(stop_reason="end_turn"), usage=SimpleNamespace(output_tokens=10)),
        SimpleNamespace(type="message_stop"),
    ]


class FakeStream:
    """Mimics anthropic.Stream for testing."""
    def __init__(self, events):
        self._events = events

    def __iter__(self):
        yield from self._events

    def __enter__(self):
        return self

    def __exit__(self, *args):
        pass

    def close(self):
        pass


class FakeAsyncStream:
    """Mimics anthropic.AsyncStream for testing."""
    def __init__(self, events):
        self._events = events

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

    async def __aenter__(self):
        return self

    async def __aexit__(self, *args):
        pass

    async def close(self):
        pass


class FakeMessageStream:
    """Mimics anthropic.MessageStream for testing."""
    def __init__(self, events, final_message):
        self._events = events
        self._final = final_message

    def __iter__(self):
        yield from self._events

    def get_final_message(self):
        return self._final

    def get_final_text(self):
        return "Hello world"

    @property
    def text_stream(self):
        yield "Hello"
        yield " world"


class FakeMessageStreamManager:
    """Mimics anthropic.MessageStreamManager for testing."""
    def __init__(self, stream):
        self._stream = stream

    def __enter__(self):
        return self._stream

    def __exit__(self, *args):
        pass


# ---------------------------------------------------------------------------
# StreamProxy (sync)
# ---------------------------------------------------------------------------

class TestStreamProxy:
    def test_events_are_passed_through(self):
        """Events are yielded to the caller, not consumed internally."""
        events = _make_stream_events()
        span = _make_span()
        token = set_current_span(span)
        proxy = StreamProxy(FakeStream(events), span, time.time(), token)

        collected = list(proxy)
        assert len(collected) == len(events)
        for original, received in zip(events, collected):
            assert original is received

    def test_accumulates_message_data(self):
        """message_start captures id/model/input_tokens; deltas accumulate text;
        message_delta captures stop_reason/output_tokens."""
        events = _make_stream_events()
        transport = MagicMock()
        span = _make_span(transport=transport)
        token = set_current_span(span)
        proxy = StreamProxy(FakeStream(events), span, time.time(), token)

        list(proxy)  # consume

        assert span._recording is False
        enqueued = transport.enqueue.call_args[0][0]
        assert enqueued["input_tokens"] == 50
        assert enqueued["output_tokens"] == 25
        assert enqueued["total_tokens"] == 75

        output = json.loads(enqueued["output"])
        assert output["content"] == [{"type": "text", "text": "Hello world"}]

        attrs = json.loads(enqueued["attributes"])
        assert attrs["stop_reason"] == "end_turn"
        assert attrs["response_id"] == "msg_123"

    def test_span_finalized_after_iteration(self):
        """Span is ended (recording=False) after the stream is fully consumed."""
        events = _make_stream_events()
        span = _make_span()
        token = set_current_span(span)
        proxy = StreamProxy(FakeStream(events), span, time.time(), token)

        assert span._recording is True
        list(proxy)
        assert span._recording is False

    def test_ttft_captured_on_first_content_block_delta(self):
        """time_to_first_token is set when the first content_block_delta arrives."""
        events = _make_stream_events()
        transport = MagicMock()
        span = _make_span(transport=transport)
        token = set_current_span(span)
        proxy = StreamProxy(FakeStream(events), span, time.time(), token)

        list(proxy)

        attrs = json.loads(transport.enqueue.call_args[0][0]["attributes"])
        assert "time_to_first_token" in attrs
        assert isinstance(attrs["time_to_first_token"], float)
        assert attrs["time_to_first_token"] >= 0

    def test_error_during_iteration_records_exception(self):
        """If the underlying stream raises, the exception is recorded on the span."""
        class ErrorStream:
            def __iter__(self):
                yield SimpleNamespace(type="message_start", message=SimpleNamespace(
                    id="msg_err", model="claude-test",
                    usage=SimpleNamespace(input_tokens=10, cache_read_input_tokens=None, cache_creation_input_tokens=None),
                ))
                raise RuntimeError("stream broken")

            def __enter__(self):
                return self

            def __exit__(self, *args):
                pass

        span = _make_span()
        token = set_current_span(span)
        proxy = StreamProxy(ErrorStream(), span, time.time(), token)

        with pytest.raises(RuntimeError, match="stream broken"):
            list(proxy)

        assert span._recording is False
        assert span.status == SpanStatus.ERROR
        assert span.error_message == "stream broken"

    def test_context_manager_enter_exit(self):
        """__enter__/__exit__ delegates to underlying stream and finalizes span."""
        events = _make_stream_events()
        span = _make_span()
        token = set_current_span(span)
        proxy = StreamProxy(FakeStream(events), span, time.time(), token)

        with proxy as p:
            assert p is proxy
            collected = list(p)

        assert len(collected) == len(events)
        assert span._recording is False

    def test_thinking_blocks_set_has_thinking(self):
        """When a thinking block is present, has_thinking metadata is set."""
        events = _make_thinking_events()
        transport = MagicMock()
        span = _make_span(transport=transport)
        token = set_current_span(span)
        proxy = StreamProxy(FakeStream(events), span, time.time(), token)

        list(proxy)

        attrs = json.loads(transport.enqueue.call_args[0][0]["attributes"])
        assert attrs.get("has_thinking") is True

    def test_close_finalizes_span(self):
        """Calling close() finalizes the span even without consuming events."""
        events = _make_stream_events()
        span = _make_span()
        token = set_current_span(span)
        proxy = StreamProxy(FakeStream(events), span, time.time(), token)

        proxy.close()
        assert span._recording is False

    def test_getattr_delegates_to_underlying_stream(self):
        """Unknown attributes are delegated to the underlying stream."""
        events = _make_stream_events()
        span = _make_span()
        token = set_current_span(span)
        fake_stream = FakeStream(events)
        fake_stream.custom_attr = "test_value"
        proxy = StreamProxy(fake_stream, span, time.time(), token)

        assert proxy.custom_attr == "test_value"

    def test_finalize_is_idempotent(self):
        """Calling _finalize multiple times does not double-end the span."""
        events = _make_stream_events()
        transport = MagicMock()
        span = _make_span(transport=transport)
        token = set_current_span(span)
        proxy = StreamProxy(FakeStream(events), span, time.time(), token)

        list(proxy)  # triggers _finalize
        proxy._finalize()  # second call should be a no-op

        assert transport.enqueue.call_count == 1

    def test_cache_tokens_captured(self):
        """Cache read and creation tokens are captured from message_start usage."""
        events = [
            SimpleNamespace(type="message_start", message=SimpleNamespace(
                id="msg_cache", model="claude-test",
                usage=SimpleNamespace(input_tokens=100, cache_read_input_tokens=40, cache_creation_input_tokens=20),
            )),
            SimpleNamespace(type="content_block_start", index=0, content_block=SimpleNamespace(type="text", text="")),
            SimpleNamespace(type="content_block_delta", index=0, delta=SimpleNamespace(type="text_delta", text="cached")),
            SimpleNamespace(type="content_block_stop", index=0),
            SimpleNamespace(type="message_delta", delta=SimpleNamespace(stop_reason="end_turn"), usage=SimpleNamespace(output_tokens=5)),
            SimpleNamespace(type="message_stop"),
        ]
        transport = MagicMock()
        span = _make_span(transport=transport)
        token = set_current_span(span)
        proxy = StreamProxy(FakeStream(events), span, time.time(), token)

        list(proxy)

        attrs = json.loads(transport.enqueue.call_args[0][0]["attributes"])
        assert attrs["cache_read_input_tokens"] == 40
        assert attrs["cache_creation_input_tokens"] == 20


# ---------------------------------------------------------------------------
# AsyncStreamProxy
# ---------------------------------------------------------------------------

class TestAsyncStreamProxy:
    def test_events_are_passed_through(self):
        """Async events are yielded to the caller unchanged."""
        events = _make_stream_events()
        span = _make_span()
        proxy = AsyncStreamProxy(FakeAsyncStream(events), span, time.time(), None)

        async def _collect():
            return [e async for e in proxy]

        collected = asyncio.run(_collect())
        assert len(collected) == len(events)
        for original, received in zip(events, collected):
            assert original is received

    def test_accumulates_message_data(self):
        """Accumulated data matches expected tokens, content, and metadata."""
        events = _make_stream_events()
        transport = MagicMock()
        span = _make_span(transport=transport)
        proxy = AsyncStreamProxy(FakeAsyncStream(events), span, time.time(), None)

        async def _consume():
            async for _ in proxy:
                pass

        asyncio.run(_consume())

        assert span._recording is False
        enqueued = transport.enqueue.call_args[0][0]
        assert enqueued["input_tokens"] == 50
        assert enqueued["output_tokens"] == 25
        assert enqueued["total_tokens"] == 75

        output = json.loads(enqueued["output"])
        assert output["content"] == [{"type": "text", "text": "Hello world"}]

        attrs = json.loads(enqueued["attributes"])
        assert attrs["stop_reason"] == "end_turn"

    def test_span_finalized_after_async_iteration(self):
        """Span recording stops after async iteration completes."""
        events = _make_stream_events()
        span = _make_span()
        proxy = AsyncStreamProxy(FakeAsyncStream(events), span, time.time(), None)

        async def _consume():
            async for _ in proxy:
                pass

        assert span._recording is True
        asyncio.run(_consume())
        assert span._recording is False

    def test_ttft_captured(self):
        """time_to_first_token is captured on async path."""
        events = _make_stream_events()
        transport = MagicMock()
        span = _make_span(transport=transport)
        proxy = AsyncStreamProxy(FakeAsyncStream(events), span, time.time(), None)

        async def _consume():
            async for _ in proxy:
                pass

        asyncio.run(_consume())

        attrs = json.loads(transport.enqueue.call_args[0][0]["attributes"])
        assert "time_to_first_token" in attrs
        assert attrs["time_to_first_token"] >= 0

    def test_error_during_async_iteration_records_exception(self):
        """Errors during async iteration are recorded on the span."""
        class ErrorAsyncStream:
            async def __aiter__(self):
                yield SimpleNamespace(type="message_start", message=SimpleNamespace(
                    id="msg_err", model="claude-test",
                    usage=SimpleNamespace(input_tokens=10, cache_read_input_tokens=None, cache_creation_input_tokens=None),
                ))
                raise RuntimeError("async stream broken")

            async def __aenter__(self):
                return self

            async def __aexit__(self, *args):
                pass

        span = _make_span()
        proxy = AsyncStreamProxy(ErrorAsyncStream(), span, time.time(), None)

        async def _consume():
            async for _ in proxy:
                pass

        with pytest.raises(RuntimeError, match="async stream broken"):
            asyncio.run(_consume())

        assert span._recording is False
        assert span.status == SpanStatus.ERROR
        assert span.error_message == "async stream broken"

    def test_async_context_manager(self):
        """__aenter__/__aexit__ delegates and finalizes."""
        events = _make_stream_events()
        span = _make_span()
        proxy = AsyncStreamProxy(FakeAsyncStream(events), span, time.time(), None)

        async def _run():
            async with proxy as p:
                assert p is proxy
                return [e async for e in p]

        collected = asyncio.run(_run())
        assert len(collected) == len(events)
        assert span._recording is False

    def test_thinking_blocks_set_has_thinking(self):
        """Thinking blocks set has_thinking on async path."""
        events = _make_thinking_events()
        transport = MagicMock()
        span = _make_span(transport=transport)
        proxy = AsyncStreamProxy(FakeAsyncStream(events), span, time.time(), None)

        async def _consume():
            async for _ in proxy:
                pass

        asyncio.run(_consume())

        attrs = json.loads(transport.enqueue.call_args[0][0]["attributes"])
        assert attrs.get("has_thinking") is True

    def test_finalize_is_idempotent(self):
        """Double finalization does not double-end the span."""
        events = _make_stream_events()
        transport = MagicMock()
        span = _make_span(transport=transport)
        proxy = AsyncStreamProxy(FakeAsyncStream(events), span, time.time(), None)

        async def _consume():
            async for _ in proxy:
                pass

        asyncio.run(_consume())
        proxy._finalize()

        assert transport.enqueue.call_count == 1


# ---------------------------------------------------------------------------
# Wrapper integration tests (via _wrap_sync / _wrap_async with stream=True)
# ---------------------------------------------------------------------------

class TestSyncStreamingWrapper:
    def test_returns_stream_proxy(self, capturing):
        """Calling create(stream=True) returns a StreamProxy wrapping the stream."""
        class FakeMessages:
            def create(self, *args, **kwargs):
                return FakeStream(_make_stream_events())

        from staso.integrations.anthropic import _wrap_sync

        _wrap_sync(FakeMessages)
        result = FakeMessages().create(model="claude-test", max_tokens=1024, stream=True)

        assert isinstance(result, StreamProxy)

    def test_proxy_yields_all_events_unchanged(self, capturing):
        """The proxy yields every event from the underlying stream."""
        events = _make_stream_events()

        class FakeMessages:
            def create(self, *args, **kwargs):
                return FakeStream(events)

        from staso.integrations.anthropic import _wrap_sync

        _wrap_sync(FakeMessages)
        proxy = FakeMessages().create(model="claude-test", stream=True)

        collected = list(proxy)
        assert len(collected) == len(events)

    def test_span_captured_with_correct_data(self, capturing):
        """After consuming the stream, the span is captured with correct token
        counts, content, and metadata."""
        class FakeMessages:
            def create(self, *args, **kwargs):
                return FakeStream(_make_stream_events())

        from staso.integrations.anthropic import _wrap_sync

        _wrap_sync(FakeMessages)
        proxy = FakeMessages().create(model="claude-test", max_tokens=512, stream=True)
        list(proxy)  # consume

        assert len(capturing) == 1
        span = capturing[0]

        assert span["name"] == "anthropic.messages.create"
        assert span["kind"] == "llm"
        assert span["model"] == "claude-test"
        assert span["input_tokens"] == 50
        assert span["output_tokens"] == 25
        assert span["total_tokens"] == 75
        assert span["status"] == "ok"

        output = json.loads(span["output"])
        assert output["content"] == [{"type": "text", "text": "Hello world"}]

        attrs = json.loads(span["attributes"])
        assert attrs["stop_reason"] == "end_turn"
        assert attrs["response_id"] == "msg_123"
        assert "time_to_first_token" in attrs
        assert attrs["provider"] == "anthropic"
        assert attrs["max_tokens"] == 512

    def test_input_captured(self, capturing):
        """Span input includes model and system prompt."""
        class FakeMessages:
            def create(self, *args, **kwargs):
                return FakeStream(_make_stream_events())

        from staso.integrations.anthropic import _wrap_sync

        _wrap_sync(FakeMessages)
        proxy = FakeMessages().create(
            model="claude-test",
            system="Be helpful.",
            stream=True,
        )
        list(proxy)

        inp = json.loads(capturing[0]["input"])
        assert inp["model"] == "claude-test"
        assert inp["system"] == "Be helpful."

    def test_error_before_stream_captures_error_span(self, capturing):
        """If the underlying create() raises, an error span is still captured."""
        class FakeMessages:
            def create(self, *args, **kwargs):
                raise ConnectionError("API down")

        from staso.integrations.anthropic import _wrap_sync

        _wrap_sync(FakeMessages)

        with pytest.raises(ConnectionError, match="API down"):
            FakeMessages().create(model="claude-test", stream=True)

        assert len(capturing) == 1
        assert capturing[0]["status"] == "error"
        assert capturing[0]["error_message"] == "API down"


class TestAsyncStreamingWrapper:
    def test_returns_async_stream_proxy(self, capturing):
        """Calling async create(stream=True) returns an AsyncStreamProxy."""
        class FakeAsyncMessages:
            async def create(self, *args, **kwargs):
                return FakeAsyncStream(_make_stream_events())

        from staso.integrations.anthropic import _wrap_async

        _wrap_async(FakeAsyncMessages)
        result = asyncio.run(FakeAsyncMessages().create(model="claude-test", stream=True))

        assert isinstance(result, AsyncStreamProxy)

    def test_span_captured_with_correct_data(self, capturing):
        """After consuming the async stream, the span has correct data."""
        class FakeAsyncMessages:
            async def create(self, *args, **kwargs):
                return FakeAsyncStream(_make_stream_events())

        from staso.integrations.anthropic import _wrap_async

        _wrap_async(FakeAsyncMessages)

        async def _run():
            proxy = await FakeAsyncMessages().create(model="claude-async", stream=True)
            async for _ in proxy:
                pass

        asyncio.run(_run())

        assert len(capturing) == 1
        span = capturing[0]

        assert span["model"] == "claude-async"
        assert span["input_tokens"] == 50
        assert span["output_tokens"] == 25
        assert span["total_tokens"] == 75
        assert span["status"] == "ok"

        output = json.loads(span["output"])
        assert output["content"] == [{"type": "text", "text": "Hello world"}]

    def test_error_before_async_stream_captures_error_span(self, capturing):
        """If async create() raises, an error span is captured."""
        class FakeAsyncMessages:
            async def create(self, *args, **kwargs):
                raise ConnectionError("async API down")

        from staso.integrations.anthropic import _wrap_async

        _wrap_async(FakeAsyncMessages)

        with pytest.raises(ConnectionError, match="async API down"):
            asyncio.run(FakeAsyncMessages().create(model="claude-async", stream=True))

        assert len(capturing) == 1
        assert capturing[0]["status"] == "error"


# ---------------------------------------------------------------------------
# MessageStreamManagerProxy
# ---------------------------------------------------------------------------

class TestMessageStreamManagerProxy:
    def test_enter_returns_message_stream_proxy(self):
        """__enter__ returns a MessageStreamProxy wrapping the inner stream."""
        final_message = SimpleNamespace(
            id="msg_final",
            model="claude-test",
            content=[SimpleNamespace(type="text", text="Hello world")],
            usage=SimpleNamespace(
                input_tokens=50, output_tokens=25,
                cache_read_input_tokens=None, cache_creation_input_tokens=None,
            ),
            stop_reason="end_turn",
        )
        inner_stream = FakeMessageStream(_make_stream_events(), final_message)
        manager = FakeMessageStreamManager(inner_stream)

        span = _make_span()
        token = set_current_span(span)
        proxy = MessageStreamManagerProxy(manager, span, time.time(), token)

        with proxy as stream:
            assert isinstance(stream, MessageStreamProxy)

    def test_exit_finalizes_span_with_final_message(self):
        """On __exit__, the span is finalized using get_final_message()."""
        final_message = SimpleNamespace(
            id="msg_final",
            model="claude-test",
            content=[SimpleNamespace(type="text", text="Hello world", model_dump=lambda: {"type": "text", "text": "Hello world"})],
            usage=SimpleNamespace(
                input_tokens=50, output_tokens=25,
                cache_read_input_tokens=None, cache_creation_input_tokens=None,
            ),
            stop_reason="end_turn",
        )
        inner_stream = FakeMessageStream(_make_stream_events(), final_message)
        manager = FakeMessageStreamManager(inner_stream)

        transport = MagicMock()
        span = _make_span(transport=transport)
        token = set_current_span(span)
        proxy = MessageStreamManagerProxy(manager, span, time.time(), token)

        with proxy as stream:
            for _ in stream:
                pass

        assert span._recording is False
        enqueued = transport.enqueue.call_args[0][0]
        assert enqueued["input_tokens"] == 50
        assert enqueued["output_tokens"] == 25
        assert enqueued["status"] == "ok"

        attrs = json.loads(enqueued["attributes"])
        assert attrs["stop_reason"] == "end_turn"
        assert attrs["response_id"] == "msg_final"

    def test_exit_finalizes_even_without_consuming(self):
        """Exiting the context manager without iterating still finalizes."""
        final_message = SimpleNamespace(
            id="msg_unused",
            model="claude-test",
            content=[SimpleNamespace(type="text", text="unused", model_dump=lambda: {"type": "text", "text": "unused"})],
            usage=SimpleNamespace(
                input_tokens=10, output_tokens=5,
                cache_read_input_tokens=None, cache_creation_input_tokens=None,
            ),
            stop_reason="end_turn",
        )
        inner_stream = FakeMessageStream([], final_message)
        manager = FakeMessageStreamManager(inner_stream)

        transport = MagicMock()
        span = _make_span(transport=transport)
        token = set_current_span(span)
        proxy = MessageStreamManagerProxy(manager, span, time.time(), token)

        with proxy:
            pass

        assert span._recording is False
        assert transport.enqueue.call_count == 1

    def test_getattr_delegates_to_manager(self):
        """Unknown attributes are forwarded to the underlying manager."""
        final_message = SimpleNamespace(
            id="msg_attr", model="claude-test",
            content=[], usage=SimpleNamespace(input_tokens=0, output_tokens=0),
            stop_reason="end_turn",
        )
        inner_stream = FakeMessageStream([], final_message)
        manager = FakeMessageStreamManager(inner_stream)
        manager.custom_prop = "manager_value"

        span = _make_span()
        token = set_current_span(span)
        proxy = MessageStreamManagerProxy(manager, span, time.time(), token)

        assert proxy.custom_prop == "manager_value"


# ---------------------------------------------------------------------------
# _EventAccumulator (unit tests)
# ---------------------------------------------------------------------------

class TestEventAccumulator:
    def test_message_start_captures_id_model_input_tokens(self):
        """message_start event populates id, model, and input_tokens."""
        span = _make_span()
        acc = _EventAccumulator(span, time.time())

        acc.process(SimpleNamespace(type="message_start", message=SimpleNamespace(
            id="msg_acc", model="claude-test",
            usage=SimpleNamespace(input_tokens=42, cache_read_input_tokens=None, cache_creation_input_tokens=None),
        )))

        assert acc.data["id"] == "msg_acc"
        assert acc.data["model"] == "claude-test"
        assert acc.data["input_tokens"] == 42

    def test_content_blocks_accumulated(self):
        """content_block_start + delta + stop produce a complete content block."""
        span = _make_span()
        acc = _EventAccumulator(span, time.time())

        acc.process(SimpleNamespace(type="content_block_start", index=0,
                                    content_block=SimpleNamespace(type="text", text="")))
        acc.process(SimpleNamespace(type="content_block_delta", index=0,
                                    delta=SimpleNamespace(type="text_delta", text="Hello")))
        acc.process(SimpleNamespace(type="content_block_delta", index=0,
                                    delta=SimpleNamespace(type="text_delta", text=" world")))
        acc.process(SimpleNamespace(type="content_block_stop", index=0))

        assert acc.data["content"] == [{"type": "text", "text": "Hello world"}]

    def test_message_delta_captures_stop_reason_and_output_tokens(self):
        """message_delta captures stop_reason and output_tokens."""
        span = _make_span()
        acc = _EventAccumulator(span, time.time())

        acc.process(SimpleNamespace(type="message_delta",
                                    delta=SimpleNamespace(stop_reason="end_turn"),
                                    usage=SimpleNamespace(output_tokens=99)))

        assert acc.data["stop_reason"] == "end_turn"
        assert acc.data["output_tokens"] == 99

    def test_ttft_set_on_first_content_block_delta(self):
        """TTFT metadata is set on the span for the first content_block_delta."""
        span = _make_span()
        start = time.time()
        acc = _EventAccumulator(span, start)

        acc.process(SimpleNamespace(type="content_block_start", index=0,
                                    content_block=SimpleNamespace(type="text", text="")))
        acc.process(SimpleNamespace(type="content_block_delta", index=0,
                                    delta=SimpleNamespace(type="text_delta", text="Hi")))

        assert "time_to_first_token" in span.metadata
        assert span.metadata["time_to_first_token"] >= 0

    def test_ttft_only_captured_once(self):
        """TTFT is only set on the first content_block_delta, not subsequent ones."""
        span = _make_span()
        acc = _EventAccumulator(span, time.time())

        acc.process(SimpleNamespace(type="content_block_start", index=0,
                                    content_block=SimpleNamespace(type="text", text="")))
        acc.process(SimpleNamespace(type="content_block_delta", index=0,
                                    delta=SimpleNamespace(type="text_delta", text="a")))
        first_ttft = span.metadata["time_to_first_token"]

        acc.process(SimpleNamespace(type="content_block_delta", index=0,
                                    delta=SimpleNamespace(type="text_delta", text="b")))

        assert span.metadata["time_to_first_token"] == first_ttft

    def test_events_without_type_are_ignored(self):
        """Events without a type attribute are silently skipped."""
        span = _make_span()
        acc = _EventAccumulator(span, time.time())

        acc.process(SimpleNamespace())  # no type attribute
        acc.process(SimpleNamespace(type=None))  # type is None

        assert acc.data["content"] == []
        assert acc.data["input_tokens"] == 0

    def test_multiple_content_blocks(self):
        """Multiple content blocks at different indices accumulate independently."""
        span = _make_span()
        acc = _EventAccumulator(span, time.time())

        acc.process(SimpleNamespace(type="content_block_start", index=0,
                                    content_block=SimpleNamespace(type="text", text="")))
        acc.process(SimpleNamespace(type="content_block_delta", index=0,
                                    delta=SimpleNamespace(type="text_delta", text="First")))
        acc.process(SimpleNamespace(type="content_block_stop", index=0))

        acc.process(SimpleNamespace(type="content_block_start", index=1,
                                    content_block=SimpleNamespace(type="text", text="")))
        acc.process(SimpleNamespace(type="content_block_delta", index=1,
                                    delta=SimpleNamespace(type="text_delta", text="Second")))
        acc.process(SimpleNamespace(type="content_block_stop", index=1))

        assert len(acc.data["content"]) == 2
        assert acc.data["content"][0] == {"type": "text", "text": "First"}
        assert acc.data["content"][1] == {"type": "text", "text": "Second"}

    def test_thinking_delta_accumulated(self):
        """thinking_delta text is accumulated into the thinking field."""
        span = _make_span()
        acc = _EventAccumulator(span, time.time())

        acc.process(SimpleNamespace(type="content_block_start", index=0,
                                    content_block=SimpleNamespace(type="thinking", text="")))
        acc.process(SimpleNamespace(type="content_block_delta", index=0,
                                    delta=SimpleNamespace(type="thinking_delta", thinking="Let me ")))
        acc.process(SimpleNamespace(type="content_block_delta", index=0,
                                    delta=SimpleNamespace(type="thinking_delta", thinking="think...")))
        acc.process(SimpleNamespace(type="content_block_stop", index=0))

        block = acc.data["content"][0]
        assert block["type"] == "thinking"
        assert block["thinking"] == "Let me think..."


# ---------------------------------------------------------------------------
# _finalize_stream_span (unit tests)
# ---------------------------------------------------------------------------

class TestFinalizeStreamSpan:
    def test_sets_span_fields_and_ends_span(self):
        """Finalize populates output, metadata, tokens, and ends the span."""
        transport = MagicMock()
        span = _make_span(transport=transport)
        token = set_current_span(span)

        message_data = {
            "id": "msg_fin",
            "content": [{"type": "text", "text": "done"}],
            "stop_reason": "end_turn",
            "input_tokens": 10,
            "output_tokens": 20,
        }

        _finalize_stream_span(span, message_data, token)

        assert span._recording is False
        assert span.input_tokens == 10
        assert span.output_tokens == 20
        assert span.total_tokens == 30
        assert span.status == SpanStatus.OK

    def test_idempotent_when_already_ended(self):
        """Calling finalize on an already-ended span does nothing."""
        transport = MagicMock()
        span = _make_span(transport=transport)
        span.end()  # end it first

        _finalize_stream_span(span, {"content": [{"type": "text", "text": "late"}]}, None)

        # Only one enqueue call from the first end()
        assert transport.enqueue.call_count == 1

    def test_thinking_content_sets_has_thinking(self):
        """Content with thinking blocks sets has_thinking metadata."""
        transport = MagicMock()
        span = _make_span(transport=transport)
        token = set_current_span(span)

        message_data = {
            "content": [
                {"type": "thinking", "thinking": "hmm"},
                {"type": "text", "text": "answer"},
            ],
            "input_tokens": 5,
            "output_tokens": 10,
        }

        _finalize_stream_span(span, message_data, token)

        enqueued = transport.enqueue.call_args[0][0]
        attrs = json.loads(enqueued["attributes"])
        assert attrs["has_thinking"] is True

    def test_missing_fields_use_defaults(self):
        """Missing fields in message_data fall back to defaults."""
        transport = MagicMock()
        span = _make_span(transport=transport)
        token = set_current_span(span)

        _finalize_stream_span(span, {}, token)

        assert span._recording is False
        assert span.input_tokens == 0
        assert span.output_tokens == 0
        assert span.total_tokens == 0


# ---------------------------------------------------------------------------
# GeneratorExit / break during iteration
# ---------------------------------------------------------------------------

class TestGeneratorExit:
    def test_sync_break_finalizes_span(self):
        """Breaking out of bare iteration (no context manager) still ends the span."""
        events = _make_stream_events()
        transport = MagicMock()
        span = _make_span(transport=transport)
        token = set_current_span(span)
        proxy = StreamProxy(FakeStream(events), span, time.time(), token)

        for event in proxy:
            break  # triggers GeneratorExit

        assert span._recording is False
        assert proxy._finalized is True

    def test_sync_break_resets_context(self):
        """Breaking out of iteration resets the ContextVar."""
        from staso.context import get_current_span

        events = _make_stream_events()
        span = _make_span()
        token = set_current_span(span)
        proxy = StreamProxy(FakeStream(events), span, time.time(), token)

        for event in proxy:
            break

        assert get_current_span() is not span

    def test_sync_return_from_loop_finalizes(self):
        """Returning from a function mid-iteration finalizes the span."""
        events = _make_stream_events()
        transport = MagicMock()
        span = _make_span(transport=transport)
        token = set_current_span(span)
        proxy = StreamProxy(FakeStream(events), span, time.time(), token)

        def consume_one():
            for event in proxy:
                return event

        consume_one()
        assert span._recording is False
        assert proxy._finalized is True

    def test_async_break_finalizes_span(self):
        """Breaking out of async iteration still ends the span."""
        events = _make_stream_events()
        transport = MagicMock()
        span = _make_span(transport=transport)
        proxy = AsyncStreamProxy(FakeAsyncStream(events), span, time.time(), None)

        async def _consume():
            async for event in proxy:
                break

        asyncio.run(_consume())
        assert span._recording is False
        assert proxy._finalized is True


# ---------------------------------------------------------------------------
# _wrap_stream_sync / _wrap_stream_async integration
# ---------------------------------------------------------------------------

class TestWrapStreamSyncIntegration:
    def test_returns_manager_proxy(self, capturing):
        """Messages.stream() returns a MessageStreamManagerProxy."""
        final_msg = SimpleNamespace(
            id="msg_ws", model="claude-test",
            content=[SimpleNamespace(type="text", text="Hello",
                                     model_dump=lambda: {"type": "text", "text": "Hello"})],
            usage=SimpleNamespace(input_tokens=20, output_tokens=10,
                                  cache_read_input_tokens=None, cache_creation_input_tokens=None),
            stop_reason="end_turn",
        )
        inner = FakeMessageStream(_make_stream_events(), final_msg)
        mgr = FakeMessageStreamManager(inner)

        class FakeMessages:
            def stream(self, *args, **kwargs):
                return mgr

        from staso.integrations.anthropic import _wrap_stream_sync

        _wrap_stream_sync(FakeMessages)
        result = FakeMessages().stream(model="claude-test", max_tokens=256)

        assert isinstance(result, MessageStreamManagerProxy)

    def test_span_captured_after_context_exit(self, capturing):
        """Exiting the stream context manager captures a complete span."""
        final_msg = SimpleNamespace(
            id="msg_ws2", model="claude-test",
            content=[SimpleNamespace(type="text", text="world",
                                     model_dump=lambda: {"type": "text", "text": "world"})],
            usage=SimpleNamespace(input_tokens=30, output_tokens=15,
                                  cache_read_input_tokens=None, cache_creation_input_tokens=None),
            stop_reason="end_turn",
        )
        inner = FakeMessageStream(_make_stream_events(), final_msg)
        mgr = FakeMessageStreamManager(inner)

        class FakeMessages:
            def stream(self, *args, **kwargs):
                return mgr

        from staso.integrations.anthropic import _wrap_stream_sync

        _wrap_stream_sync(FakeMessages)

        with FakeMessages().stream(model="claude-test", max_tokens=128) as stream:
            for _ in stream:
                pass

        assert len(capturing) == 1
        span = capturing[0]
        assert span["name"] == "anthropic.messages.stream"
        assert span["kind"] == "llm"
        assert span["model"] == "claude-test"
        assert span["input_tokens"] == 30
        assert span["output_tokens"] == 15
        assert span["status"] == "ok"

        attrs = json.loads(span["attributes"])
        assert attrs["provider"] == "anthropic"
        assert attrs["max_tokens"] == 128
