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

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

import pytest

from staso.context import set_current_span
from staso.integrations._openai_streaming import (
    AsyncOpenAIStreamProxy,
    OpenAIStreamProxy,
    _OpenAIEventAccumulator,
)
from staso.provider import Span
from staso.types import SpanKind


def _make_span(transport=None):
    """Create a test Span with deterministic IDs."""
    return Span(
        name="openai.chat.completions.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_chunks(content_parts=None, finish_reason="stop", include_usage=False):
    """Generate a sequence of ChatCompletionChunk-like objects."""
    if content_parts is None:
        content_parts = ["Hello", " world", "!"]

    chunks = []
    for i, part in enumerate(content_parts):
        chunks.append(
            SimpleNamespace(
                id="chatcmpl-test",
                choices=[
                    SimpleNamespace(
                        delta=SimpleNamespace(
                            content=part,
                            role="assistant" if i == 0 else None,
                            tool_calls=None,
                        ),
                        finish_reason=None,
                    )
                ],
                model="gpt-4",
                usage=None,
            )
        )

    chunks.append(
        SimpleNamespace(
            id="chatcmpl-test",
            choices=[
                SimpleNamespace(
                    delta=SimpleNamespace(content=None, tool_calls=None),
                    finish_reason=finish_reason,
                )
            ],
            model="gpt-4",
            usage=SimpleNamespace(
                prompt_tokens=10,
                completion_tokens=20,
                total_tokens=30,
                completion_tokens_details=None,
            )
            if include_usage
            else None,
        )
    )
    return chunks


class FakeStream:
    """Mimics openai.Stream for testing."""

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

    def __iter__(self):
        yield from self._chunks

    def __enter__(self):
        return self

    def __exit__(self, *args):
        pass


class FakeAsyncStream:
    """Mimics openai.AsyncStream for testing."""

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

    async def __aiter__(self):
        for chunk in self._chunks:
            yield chunk

    async def __aenter__(self):
        return self

    async def __aexit__(self, *args):
        pass


class TestEventAccumulator:
    def test_accumulates_content(self):
        span = _make_span()
        acc = _OpenAIEventAccumulator(span, 0.0)
        chunks = _make_chunks(["Hello", " world"])
        for chunk in chunks:
            acc.process(chunk)
        assert acc.content == "Hello world"

    def test_captures_finish_reason(self):
        span = _make_span()
        acc = _OpenAIEventAccumulator(span, 0.0)
        chunks = _make_chunks(finish_reason="length")
        for chunk in chunks:
            acc.process(chunk)
        assert acc.finish_reason == "length"

    def test_captures_usage(self):
        span = _make_span()
        acc = _OpenAIEventAccumulator(span, 0.0)
        chunks = _make_chunks(include_usage=True)
        for chunk in chunks:
            acc.process(chunk)
        assert acc.usage is not None
        assert acc.usage["prompt_tokens"] == 10
        assert acc.usage["completion_tokens"] == 20

    def test_captures_response_id(self):
        span = _make_span()
        acc = _OpenAIEventAccumulator(span, 0.0)
        chunks = _make_chunks()
        for chunk in chunks:
            acc.process(chunk)
        assert acc.response_id == "chatcmpl-test"


class TestSyncStreaming:
    def test_yields_all_chunks(self, capturing):
        transport = MagicMock()
        transport.enqueue.side_effect = lambda d: capturing.append(d)
        span = _make_span(transport=transport)
        token = set_current_span(span)
        chunks = _make_chunks(["Hello", " world", "!"])
        stream = FakeStream(chunks)

        proxy = OpenAIStreamProxy(stream, span, 0.0, token)
        collected = list(proxy)

        assert len(collected) == len(chunks)

    def test_emits_span_with_content(self, capturing):
        transport = MagicMock()
        transport.enqueue.side_effect = lambda d: capturing.append(d)
        span = _make_span(transport=transport)
        token = set_current_span(span)
        chunks = _make_chunks(["Hello", " world"])
        stream = FakeStream(chunks)

        proxy = OpenAIStreamProxy(stream, span, 0.0, token)
        list(proxy)

        assert len(capturing) == 1
        output = json.loads(capturing[0]["output"])
        assert output["content"] == "Hello world"

    def test_emits_span_with_usage(self, capturing):
        transport = MagicMock()
        transport.enqueue.side_effect = lambda d: capturing.append(d)
        span = _make_span(transport=transport)
        token = set_current_span(span)
        chunks = _make_chunks(include_usage=True)
        stream = FakeStream(chunks)

        proxy = OpenAIStreamProxy(stream, span, 0.0, token)
        list(proxy)

        assert capturing[0]["input_tokens"] == 10
        assert capturing[0]["output_tokens"] == 20
        assert capturing[0]["total_tokens"] == 30

    def test_context_manager(self, capturing):
        transport = MagicMock()
        transport.enqueue.side_effect = lambda d: capturing.append(d)
        span = _make_span(transport=transport)
        token = set_current_span(span)
        chunks = _make_chunks(["Hi"])
        stream = FakeStream(chunks)

        proxy = OpenAIStreamProxy(stream, span, 0.0, token)
        with proxy as p:
            for _ in p:
                pass

        assert len(capturing) == 1


class TestAsyncStreaming:
    def test_yields_all_chunks(self, capturing):
        transport = MagicMock()
        transport.enqueue.side_effect = lambda d: capturing.append(d)
        span = _make_span(transport=transport)
        chunks = _make_chunks(["Hello", " world"])
        stream = FakeAsyncStream(chunks)

        async def collect():
            token = set_current_span(span)
            proxy = AsyncOpenAIStreamProxy(stream, span, 0.0, token)
            result = []
            async for chunk in proxy:
                result.append(chunk)
            return result

        collected = asyncio.run(collect())
        assert len(collected) == len(chunks)

    def test_emits_span_with_content(self, capturing):
        transport = MagicMock()
        transport.enqueue.side_effect = lambda d: capturing.append(d)
        span = _make_span(transport=transport)
        chunks = _make_chunks(["Hello", " async"])
        stream = FakeAsyncStream(chunks)

        async def collect():
            token = set_current_span(span)
            proxy = AsyncOpenAIStreamProxy(stream, span, 0.0, token)
            async for _ in proxy:
                pass

        asyncio.run(collect())

        assert len(capturing) == 1
        output = json.loads(capturing[0]["output"])
        assert output["content"] == "Hello async"
