"""Tests for the streaming wrapper."""

import time

import pytest

from llm_observatory.streaming import StreamWrapper, AnthropicStreamWrapper


class FakeDelta:
    def __init__(self, content=None, reasoning_content=None):
        self.content = content
        self.reasoning_content = reasoning_content
        self.role = None
        self.tool_calls = None


class FakeStreamChoice:
    def __init__(self, content=None, finish_reason=None, reasoning_content=None):
        self.index = 0
        self.delta = FakeDelta(content, reasoning_content)
        self.finish_reason = finish_reason


class FakeStreamUsage:
    def __init__(self):
        self.prompt_tokens = 10
        self.completion_tokens = 5
        self.total_tokens = 15


class FakeChunk:
    def __init__(self, content=None, finish_reason=None, model="gpt-4o", chunk_id="chatcmpl-x", usage=None, reasoning_content=None):
        self.id = chunk_id
        self.model = model
        self.choices = [FakeStreamChoice(content, finish_reason, reasoning_content)]
        self.usage = usage


def test_stream_wrapper_captures_chunks():
    chunks = [
        FakeChunk(content="Hello"),
        FakeChunk(content=" world"),
        FakeChunk(content="!", finish_reason="stop", usage=FakeStreamUsage()),
    ]

    captured = {}

    def on_complete(response, end_time, time_to_first_token_ms, tokens_per_second):
        captured["response"] = response
        captured["ttft"] = time_to_first_token_ms
        captured["tps"] = tokens_per_second

    wrapper = StreamWrapper(
        stream=iter(chunks),
        start_time=time.time(),
        request_kwargs={"model": "gpt-4o", "stream": True},
        on_complete=on_complete,
    )

    collected = list(wrapper)

    assert len(collected) == 3
    assert captured["response"].model == "gpt-4o"
    assert captured["response"].choices[0].message.content == "Hello world!"
    assert captured["ttft"] is not None
    assert captured["response"].usage.total_tokens == 15


def test_stream_wrapper_ttft():
    start = time.time()

    chunks = [
        FakeChunk(content="Hi"),
        FakeChunk(content="!", finish_reason="stop"),
    ]

    captured = {}

    def on_complete(response, end_time, time_to_first_token_ms, tokens_per_second):
        captured["ttft"] = time_to_first_token_ms

    wrapper = StreamWrapper(
        stream=iter(chunks),
        start_time=start,
        request_kwargs={"model": "gpt-4o", "stream": True},
        on_complete=on_complete,
    )

    # Iterate with a small delay to make TTFT measurable
    for _ in wrapper:
        pass

    assert captured["ttft"] is not None
    assert captured["ttft"] >= 0


def test_stream_wrapper_empty():
    captured = {}

    def on_complete(**kw):
        captured.update(kw)

    wrapper = StreamWrapper(
        stream=iter([]),
        start_time=time.time(),
        request_kwargs={"model": "gpt-4o", "stream": True},
        on_complete=on_complete,
    )

    collected = list(wrapper)
    assert len(collected) == 0
    # Empty stream should not call on_complete (no chunks)
    assert "response" not in captured


class FakeAsyncStream:
    """Async iterator that yields FakeChunk objects."""

    def __init__(self, chunks):
        self._chunks = iter(chunks)

    def __aiter__(self):
        return self

    async def __anext__(self):
        try:
            return next(self._chunks)
        except StopIteration:
            raise StopAsyncIteration


@pytest.mark.asyncio
async def test_async_stream_wrapper_captures_chunks():
    chunks = [
        FakeChunk(content="Hello"),
        FakeChunk(content=" async"),
        FakeChunk(content="!", finish_reason="stop", usage=FakeStreamUsage()),
    ]

    captured = {}

    def on_complete(response, end_time, time_to_first_token_ms, tokens_per_second):
        captured["response"] = response
        captured["ttft"] = time_to_first_token_ms
        captured["tps"] = tokens_per_second

    wrapper = StreamWrapper(
        stream=FakeAsyncStream(chunks),
        start_time=time.time(),
        request_kwargs={"model": "gpt-4o", "stream": True},
        on_complete=on_complete,
    )

    collected = []
    async for chunk in wrapper:
        collected.append(chunk)

    assert len(collected) == 3
    assert captured["response"].model == "gpt-4o"
    assert captured["response"].choices[0].message.content == "Hello async!"
    assert captured["ttft"] is not None
    assert captured["response"].usage.total_tokens == 15


def test_stream_wrapper_error_mid_stream():
    """Exception raised mid-stream should propagate and trigger on_error callback."""

    def bad_stream():
        yield FakeChunk(content="Hello")
        raise ConnectionError("connection reset by peer")

    on_complete_called = []
    error_captured = {}

    def on_complete(response, end_time, time_to_first_token_ms, tokens_per_second):
        on_complete_called.append(True)

    def on_error(exc, end_time):
        error_captured["exc"] = exc
        error_captured["end_time"] = end_time

    wrapper = StreamWrapper(
        stream=bad_stream(),
        start_time=time.time(),
        request_kwargs={"model": "gpt-4o", "stream": True},
        on_complete=on_complete,
        on_error=on_error,
    )

    raised = None
    try:
        list(wrapper)
    except ConnectionError as exc:
        raised = exc

    assert raised is not None, "ConnectionError should have propagated to caller"
    assert "connection reset" in str(raised)
    # on_error should have been called
    assert "exc" in error_captured
    assert isinstance(error_captured["exc"], ConnectionError)
    # on_complete should NOT have been called
    assert not on_complete_called


def test_stream_wrapper_error_before_first_token():
    """Exception raised before any chunk should propagate and trigger on_error."""

    def empty_then_error():
        raise RuntimeError("upstream unavailable")
        yield  # make it a generator

    on_complete_called = []
    error_captured = {}

    def on_complete(response, end_time, time_to_first_token_ms, tokens_per_second):
        on_complete_called.append(True)

    def on_error(exc, end_time):
        error_captured["exc"] = exc

    wrapper = StreamWrapper(
        stream=empty_then_error(),
        start_time=time.time(),
        request_kwargs={"model": "gpt-4o", "stream": True},
        on_complete=on_complete,
        on_error=on_error,
    )

    raised = None
    try:
        list(wrapper)
    except RuntimeError as exc:
        raised = exc

    assert raised is not None, "RuntimeError should have propagated to caller"
    assert "upstream unavailable" in str(raised)
    # on_complete should not fire — no chunks were collected
    assert not on_complete_called
    # on_error should fire even before first token
    assert "exc" in error_captured
    assert isinstance(error_captured["exc"], RuntimeError)


@pytest.mark.asyncio
async def test_async_stream_wrapper_error_mid_stream():
    """Async mid-stream exception should propagate and trigger on_error."""

    class ErroringAsyncStream:
        def __init__(self):
            self._chunks = iter([FakeChunk(content="Hi")])
            self._count = 0

        def __aiter__(self):
            return self

        async def __anext__(self):
            self._count += 1
            if self._count == 1:
                return next(self._chunks)
            raise ConnectionError("async connection reset")

    on_complete_called = []
    error_captured = {}

    def on_complete(response, end_time, time_to_first_token_ms, tokens_per_second):
        on_complete_called.append(True)

    def on_error(exc, end_time):
        error_captured["exc"] = exc

    wrapper = StreamWrapper(
        stream=ErroringAsyncStream(),
        start_time=time.time(),
        request_kwargs={"model": "gpt-4o", "stream": True},
        on_complete=on_complete,
        on_error=on_error,
    )

    raised = None
    try:
        async for _ in wrapper:
            pass
    except ConnectionError as exc:
        raised = exc

    assert raised is not None, "ConnectionError should have propagated to caller"
    assert "async connection reset" in str(raised)
    assert "exc" in error_captured
    assert isinstance(error_captured["exc"], ConnectionError)


@pytest.mark.asyncio
async def test_async_stream_wrapper_context_manager():
    chunks = [
        FakeChunk(content="test", finish_reason="stop"),
    ]

    captured = {}

    def on_complete(response, end_time, time_to_first_token_ms, tokens_per_second):
        captured["response"] = response

    wrapper = StreamWrapper(
        stream=FakeAsyncStream(chunks),
        start_time=time.time(),
        request_kwargs={"model": "gpt-4o", "stream": True},
        on_complete=on_complete,
    )

    collected = []
    async with wrapper as w:
        async for chunk in w:
            collected.append(chunk)

    assert len(collected) == 1
    assert captured["response"].choices[0].message.content == "test"


def test_stream_wrapper_captures_reasoning_content():
    """Reasoning content (delta.reasoning_content) should be reassembled separately."""
    chunks = [
        FakeChunk(reasoning_content="Let me think"),
        FakeChunk(reasoning_content=" step by step..."),
        FakeChunk(content="The answer is 42.", finish_reason="stop", usage=FakeStreamUsage()),
    ]

    captured = {}

    def on_complete(response, end_time, time_to_first_token_ms, tokens_per_second):
        captured["response"] = response

    wrapper = StreamWrapper(
        stream=iter(chunks),
        start_time=time.time(),
        request_kwargs={"model": "o3-mini", "stream": True},
        on_complete=on_complete,
    )

    list(wrapper)

    msg = captured["response"].choices[0].message
    assert msg.content == "The answer is 42."
    assert msg.reasoning_content == "Let me think step by step..."


def test_stream_wrapper_reasoning_only_no_content():
    """Some models stream only reasoning_content with no regular content (e.g. Parasail)."""
    chunks = [
        FakeChunk(reasoning_content="Thinking about SQS vs Kafka..."),
        FakeChunk(reasoning_content=" SQS wins for low-ops."),
        FakeChunk(finish_reason="length", usage=FakeStreamUsage()),
    ]

    captured = {}

    def on_complete(response, end_time, time_to_first_token_ms, tokens_per_second):
        captured["response"] = response

    wrapper = StreamWrapper(
        stream=iter(chunks),
        start_time=time.time(),
        request_kwargs={"model": "parasail/parasail-gpt-oss-20b-fast-mem0", "stream": True},
        on_complete=on_complete,
    )

    list(wrapper)

    msg = captured["response"].choices[0].message
    assert msg.content == ""
    assert msg.reasoning_content == "Thinking about SQS vs Kafka... SQS wins for low-ops."


def test_stream_wrapper_no_reasoning_content():
    """When no reasoning_content is present, it should be None."""
    chunks = [
        FakeChunk(content="Hello"),
        FakeChunk(content="!", finish_reason="stop"),
    ]

    captured = {}

    def on_complete(response, end_time, time_to_first_token_ms, tokens_per_second):
        captured["response"] = response

    wrapper = StreamWrapper(
        stream=iter(chunks),
        start_time=time.time(),
        request_kwargs={"model": "gpt-4o", "stream": True},
        on_complete=on_complete,
    )

    list(wrapper)

    msg = captured["response"].choices[0].message
    assert msg.content == "Hello!"
    assert msg.reasoning_content is None


def test_stream_wrapper_early_break_finalizes():
    """Breaking out of iteration early should still finalize via __del__."""
    chunks = [
        FakeChunk(content="Hello"),
        FakeChunk(content=" world"),
        FakeChunk(content="!", finish_reason="stop", usage=FakeStreamUsage()),
    ]

    captured = {}

    def on_complete(response, end_time, time_to_first_token_ms, tokens_per_second):
        captured["response"] = response

    wrapper = StreamWrapper(
        stream=iter(chunks),
        start_time=time.time(),
        request_kwargs={"model": "gpt-4o", "stream": True},
        on_complete=on_complete,
    )

    # Only consume the first chunk, then break
    for chunk in wrapper:
        break

    # Explicitly delete to trigger __del__
    del wrapper

    assert "response" in captured, "on_complete should fire via __del__ after early break"
    assert captured["response"].choices[0].message.content == "Hello"


def test_stream_wrapper_context_manager_early_break_finalizes():
    """Breaking inside a context manager should still finalize."""
    chunks = [
        FakeChunk(content="Hello"),
        FakeChunk(content=" world"),
        FakeChunk(content="!", finish_reason="stop"),
    ]

    captured = {}

    def on_complete(response, end_time, time_to_first_token_ms, tokens_per_second):
        captured["response"] = response

    wrapper = StreamWrapper(
        stream=iter(chunks),
        start_time=time.time(),
        request_kwargs={"model": "gpt-4o", "stream": True},
        on_complete=on_complete,
    )

    with wrapper as w:
        for chunk in w:
            break

    assert "response" in captured, "on_complete should fire via __exit__ after early break"


def test_stream_wrapper_on_error_failure_does_not_raise():
    """If on_error callback itself raises, the original exception should still propagate."""

    def bad_stream():
        yield FakeChunk(content="Hi")
        raise ConnectionError("connection reset")

    def on_complete(**kw):
        pass

    def on_error(**kw):
        raise RuntimeError("on_error callback exploded")

    wrapper = StreamWrapper(
        stream=bad_stream(),
        start_time=time.time(),
        request_kwargs={"model": "gpt-4o", "stream": True},
        on_complete=on_complete,
        on_error=on_error,
    )

    # The original ConnectionError should propagate, not the RuntimeError from on_error
    raised = None
    try:
        list(wrapper)
    except ConnectionError as exc:
        raised = exc

    assert raised is not None, "Original exception should propagate even if on_error fails"


def test_stream_wrapper_on_complete_failure_does_not_raise():
    """If on_complete callback itself raises, it should not propagate to the caller."""
    chunks = [
        FakeChunk(content="Hello", finish_reason="stop"),
    ]

    def on_complete(**kw):
        raise RuntimeError("on_complete callback exploded")

    wrapper = StreamWrapper(
        stream=iter(chunks),
        start_time=time.time(),
        request_kwargs={"model": "gpt-4o", "stream": True},
        on_complete=on_complete,
    )

    # Should NOT raise — the callback failure is swallowed
    collected = list(wrapper)
    assert len(collected) == 1


def test_stream_wrapper_no_on_error_still_propagates():
    """Mid-stream error without on_error callback should still propagate the exception."""

    def bad_stream():
        yield FakeChunk(content="Hi")
        raise ConnectionError("connection reset")

    wrapper = StreamWrapper(
        stream=bad_stream(),
        start_time=time.time(),
        request_kwargs={"model": "gpt-4o", "stream": True},
        on_complete=lambda **kw: None,
        # no on_error
    )

    raised = None
    try:
        list(wrapper)
    except ConnectionError as exc:
        raised = exc

    assert raised is not None, "Exception should propagate even without on_error"


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


class FakeAnthropicStreamEvent:
    """Fake Anthropic streaming event."""
    def __init__(self, type, **kwargs):
        self.type = type
        for k, v in kwargs.items():
            setattr(self, k, v)


class FakeAnthropicDelta:
    def __init__(self, type, text=None, thinking=None, stop_reason=None):
        self.type = type
        self.text = text
        self.thinking = thinking
        self.stop_reason = stop_reason


class FakeAnthropicUsage:
    def __init__(self, input_tokens=0, output_tokens=0):
        self.input_tokens = input_tokens
        self.output_tokens = output_tokens


class FakeAnthropicMessage:
    def __init__(self, id="msg_x", model="claude-sonnet-4-20250514", input_tokens=20):
        self.id = id
        self.model = model
        self.usage = FakeAnthropicUsage(input_tokens=input_tokens)


def test_anthropic_stream_wrapper_captures_text():
    """AnthropicStreamWrapper should reassemble text from content_block_delta events."""
    events = [
        FakeAnthropicStreamEvent("message_start", message=FakeAnthropicMessage()),
        FakeAnthropicStreamEvent("content_block_delta", delta=FakeAnthropicDelta("text_delta", text="Hello")),
        FakeAnthropicStreamEvent("content_block_delta", delta=FakeAnthropicDelta("text_delta", text=" world")),
        FakeAnthropicStreamEvent("message_delta", delta=FakeAnthropicDelta("message_delta", stop_reason="end_turn"), usage=FakeAnthropicUsage(output_tokens=10)),
        FakeAnthropicStreamEvent("message_stop"),
    ]

    captured = {}

    def on_complete(response, end_time, time_to_first_token_ms, tokens_per_second):
        captured["response"] = response
        captured["ttft"] = time_to_first_token_ms

    wrapper = AnthropicStreamWrapper(
        stream=iter(events),
        start_time=time.time(),
        request_kwargs={"model": "claude-sonnet-4-20250514", "stream": True},
        on_complete=on_complete,
    )

    collected = list(wrapper)
    assert len(collected) == 5

    resp = captured["response"]
    # Should have text content
    text_blocks = [b for b in resp.content if b.type == "text"]
    assert len(text_blocks) == 1
    assert text_blocks[0].text == "Hello world"

    assert resp.model == "claude-sonnet-4-20250514"
    assert resp.stop_reason == "end_turn"
    assert resp.usage.input_tokens == 20
    assert resp.usage.output_tokens == 10
    assert captured["ttft"] is not None


def test_anthropic_stream_wrapper_captures_thinking():
    """AnthropicStreamWrapper should reassemble thinking content."""
    events = [
        FakeAnthropicStreamEvent("message_start", message=FakeAnthropicMessage()),
        FakeAnthropicStreamEvent("content_block_delta", delta=FakeAnthropicDelta("thinking_delta", thinking="Let me think")),
        FakeAnthropicStreamEvent("content_block_delta", delta=FakeAnthropicDelta("thinking_delta", thinking="...")),
        FakeAnthropicStreamEvent("content_block_delta", delta=FakeAnthropicDelta("text_delta", text="Answer")),
        FakeAnthropicStreamEvent("message_delta", delta=FakeAnthropicDelta("message_delta", stop_reason="end_turn"), usage=FakeAnthropicUsage(output_tokens=5)),
        FakeAnthropicStreamEvent("message_stop"),
    ]

    captured = {}

    def on_complete(response, end_time, time_to_first_token_ms, tokens_per_second):
        captured["response"] = response

    wrapper = AnthropicStreamWrapper(
        stream=iter(events),
        start_time=time.time(),
        request_kwargs={"model": "claude-sonnet-4-20250514", "stream": True},
        on_complete=on_complete,
    )

    list(wrapper)

    resp = captured["response"]
    thinking_blocks = [b for b in resp.content if b.type == "thinking"]
    text_blocks = [b for b in resp.content if b.type == "text"]
    assert len(thinking_blocks) == 1
    assert thinking_blocks[0].thinking == "Let me think..."
    assert text_blocks[0].text == "Answer"


def test_anthropic_stream_wrapper_empty():
    """Empty stream should not call on_complete."""
    captured = {}

    def on_complete(**kw):
        captured.update(kw)

    wrapper = AnthropicStreamWrapper(
        stream=iter([]),
        start_time=time.time(),
        request_kwargs={"model": "claude-sonnet-4-20250514", "stream": True},
        on_complete=on_complete,
    )

    collected = list(wrapper)
    assert len(collected) == 0
    assert "response" not in captured


def test_anthropic_stream_wrapper_error_mid_stream():
    """Mid-stream error should trigger on_error."""
    def bad_stream():
        yield FakeAnthropicStreamEvent("message_start", message=FakeAnthropicMessage())
        yield FakeAnthropicStreamEvent("content_block_delta", delta=FakeAnthropicDelta("text_delta", text="Hi"))
        raise ConnectionError("stream interrupted")

    error_captured = {}

    def on_error(exc, end_time):
        error_captured["exc"] = exc

    wrapper = AnthropicStreamWrapper(
        stream=bad_stream(),
        start_time=time.time(),
        request_kwargs={"model": "claude-sonnet-4-20250514", "stream": True},
        on_complete=lambda **kw: None,
        on_error=on_error,
    )

    raised = None
    try:
        list(wrapper)
    except ConnectionError as exc:
        raised = exc

    assert raised is not None
    assert "exc" in error_captured


@pytest.mark.asyncio
async def test_anthropic_stream_wrapper_async():
    """AnthropicStreamWrapper should work with async iteration."""
    events = [
        FakeAnthropicStreamEvent("message_start", message=FakeAnthropicMessage()),
        FakeAnthropicStreamEvent("content_block_delta", delta=FakeAnthropicDelta("text_delta", text="async hello")),
        FakeAnthropicStreamEvent("message_delta", delta=FakeAnthropicDelta("message_delta", stop_reason="end_turn"), usage=FakeAnthropicUsage(output_tokens=3)),
        FakeAnthropicStreamEvent("message_stop"),
    ]

    captured = {}

    def on_complete(response, end_time, time_to_first_token_ms, tokens_per_second):
        captured["response"] = response

    class AsyncEvents:
        def __init__(self, events):
            self._events = iter(events)
        def __aiter__(self):
            return self
        async def __anext__(self):
            try:
                return next(self._events)
            except StopIteration:
                raise StopAsyncIteration

    wrapper = AnthropicStreamWrapper(
        stream=AsyncEvents(events),
        start_time=time.time(),
        request_kwargs={"model": "claude-sonnet-4-20250514", "stream": True},
        on_complete=on_complete,
    )

    collected = []
    async for event in wrapper:
        collected.append(event)

    assert len(collected) == 4
    text_blocks = [b for b in captured["response"].content if b.type == "text"]
    assert text_blocks[0].text == "async hello"


def test_anthropic_stream_wrapper_context_manager():
    """AnthropicStreamWrapper should work as a context manager."""
    events = [
        FakeAnthropicStreamEvent("message_start", message=FakeAnthropicMessage()),
        FakeAnthropicStreamEvent("content_block_delta", delta=FakeAnthropicDelta("text_delta", text="ctx")),
        FakeAnthropicStreamEvent("message_delta", delta=FakeAnthropicDelta("message_delta", stop_reason="end_turn"), usage=FakeAnthropicUsage(output_tokens=1)),
        FakeAnthropicStreamEvent("message_stop"),
    ]

    captured = {}

    def on_complete(response, end_time, time_to_first_token_ms, tokens_per_second):
        captured["response"] = response

    wrapper = AnthropicStreamWrapper(
        stream=iter(events),
        start_time=time.time(),
        request_kwargs={"model": "claude-sonnet-4-20250514", "stream": True},
        on_complete=on_complete,
    )

    with wrapper as w:
        for _ in w:
            pass

    assert "response" in captured
