import pytest

from cascade import llm_wrappers


class _Usage:
    def __init__(self, input_tokens=0, output_tokens=0, prompt_tokens=0, completion_tokens=0, total_tokens=0):
        self.input_tokens = input_tokens
        self.output_tokens = output_tokens
        self.prompt_tokens = prompt_tokens
        self.completion_tokens = completion_tokens
        self.total_tokens = total_tokens


class _AnthropicTextBlock:
    def __init__(self, text: str):
        self.type = "text"
        self.text = text


class _AnthropicResp:
    def __init__(self, text: str, input_tokens: int, output_tokens: int):
        self.content = [_AnthropicTextBlock(text)]
        self.usage = _Usage(input_tokens=input_tokens, output_tokens=output_tokens)


class _OpenAIMessage:
    def __init__(self, content: str):
        self.content = content
        self.tool_calls = None


class _OpenAIChoice:
    def __init__(self, message: _OpenAIMessage, finish_reason: str = "stop"):
        self.message = message
        self.finish_reason = finish_reason
        self.delta = None


class _OpenAIResp:
    def __init__(self, content: str, prompt_tokens: int, completion_tokens: int):
        self.choices = [_OpenAIChoice(_OpenAIMessage(content))]
        self.usage = _Usage(
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            total_tokens=prompt_tokens + completion_tokens,
        )


class _Delta:
    def __init__(self, content: str):
        self.content = content


class _StreamChoice:
    def __init__(self, content: str = "", finish_reason=None):
        self.delta = _Delta(content)
        self.finish_reason = finish_reason


class _StreamChunk:
    def __init__(self, content: str = "", *, finish_reason=None, usage=None):
        self.choices = [_StreamChoice(content, finish_reason=finish_reason)]
        self.usage = usage


def _patch_fake_tracer(monkeypatch, fake_tracer):
    monkeypatch.setattr(llm_wrappers.trace, "get_current_span", lambda: fake_tracer.current_span)


def test_anthropic_known_model_sets_cost_usd_and_cost(monkeypatch, fake_tracer):
    _patch_fake_tracer(monkeypatch, fake_tracer)

    class _Orig:
        def create(self, *args, **kwargs):
            return _AnthropicResp("ok", input_tokens=1000, output_tokens=1000)

    wrapped = llm_wrappers._AnthropicMessagesWrapper(_Orig(), fake_tracer)
    wrapped.create(
        model="claude-3-5-haiku-20241022",
        max_tokens=10,
        messages=[{"role": "user", "content": "hi"}],
    )

    span = fake_tracer.spans[-1]
    assert span.attributes["cascade.span_type"] == "llm"
    # Pricing for haiku: input 0.0008/1k, output 0.004/1k => 0.0048 total
    assert span.attributes["llm.cost_usd"] == pytest.approx(0.0048)
    assert span.attributes["llm.cost"] == pytest.approx(0.0048)


def test_anthropic_unknown_model_does_not_guess_cost(monkeypatch, fake_tracer):
    _patch_fake_tracer(monkeypatch, fake_tracer)

    class _Orig:
        def create(self, *args, **kwargs):
            return _AnthropicResp("ok", input_tokens=10, output_tokens=10)

    wrapped = llm_wrappers._AnthropicMessagesWrapper(_Orig(), fake_tracer)
    wrapped.create(
        model="claude-unknown-fake",
        max_tokens=10,
        messages=[{"role": "user", "content": "hi"}],
    )

    span = fake_tracer.spans[-1]
    assert "llm.cost_usd" not in span.attributes
    assert "llm.cost" not in span.attributes


def test_openai_non_streaming_does_not_set_cost(monkeypatch, fake_tracer):
    _patch_fake_tracer(monkeypatch, fake_tracer)

    class _Orig:
        def create(self, *args, **kwargs):
            return _OpenAIResp("hello", prompt_tokens=10, completion_tokens=5)

    wrapped = llm_wrappers._OpenAICompletionsWrapper(_Orig(), fake_tracer, base_url="")
    wrapped.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": "hi"}],
    )

    span = fake_tracer.spans[-1]
    assert span.attributes["cascade.span_type"] == "llm"
    assert "llm.cost_usd" not in span.attributes


def test_openai_streaming_accumulates_completion_and_usage(monkeypatch, fake_tracer):
    _patch_fake_tracer(monkeypatch, fake_tracer)

    class _Orig:
        def create(self, *args, **kwargs):
            # last chunk carries usage (OpenRouter-like)
            yield _StreamChunk("hel")
            yield _StreamChunk("lo")
            yield _StreamChunk("", finish_reason="stop", usage=_Usage(prompt_tokens=7, completion_tokens=2, total_tokens=9))

    wrapped = llm_wrappers._OpenAICompletionsWrapper(_Orig(), fake_tracer, base_url="https://openrouter.ai/api/v1")

    stream = wrapped.create(
        model="gpt-4o-mini",
        stream=True,
        messages=[{"role": "user", "content": "say hello"}],
    )

    # Trigger iteration to completion (auto-exit)
    list(stream)

    span = fake_tracer.spans[-1]
    assert span.attributes["llm.streaming"] is True
    assert span.attributes["llm.completion"] == "hello"
    assert span.attributes["llm.total_tokens"] == 9
    assert "llm.cost_usd" not in span.attributes


def test_openai_streaming_error_sets_span_error(monkeypatch, fake_tracer):
    _patch_fake_tracer(monkeypatch, fake_tracer)

    class _Orig:
        def create(self, *args, **kwargs):
            def gen():
                yield _StreamChunk("partial")
                raise RuntimeError("stream broke")

            return gen()

    wrapped = llm_wrappers._OpenAICompletionsWrapper(_Orig(), fake_tracer, base_url="")
    stream = wrapped.create(
        model="gpt-4o-mini",
        stream=True,
        messages=[{"role": "user", "content": "hi"}],
    )

    with pytest.raises(RuntimeError, match="stream broke"):
        with stream:
            for _ in stream:
                pass

    span = fake_tracer.spans[-1]
    assert span.status.status_code.name == "ERROR"
