from backend.vercel_ai_normalize import (
    normalize_vercel_ai_span,
    _infer_cascade_span_type,
)


# ── Existing tests ───────────────────────────────────────────────────────

def test_normalize_toolcall_span():
    attrs = {
        "ai.operationId": "ai.toolCall",
        "ai.toolCall.name": "weather",
        "ai.toolCall.args": {"location": "SF"},
        "ai.toolCall.result": {"temperatureF": 70},
    }

    out = normalize_vercel_ai_span("ai.toolCall", attrs)
    assert out["cascade.span_type"] == "tool"
    assert out["tool.name"] == "weather"
    assert out["tool.input"] == {"location": "SF"}
    assert out["tool.output"] == {"temperatureF": 70}


def test_normalize_llm_span_with_gen_ai_semconv():
    attrs = {
        "ai.operationId": "ai.generateText.doGenerate",
        "gen_ai.request.model": "gpt-5",
        "gen_ai.system": "openai",
        "gen_ai.usage.input_tokens": 11,
        "gen_ai.usage.output_tokens": 22,
        "ai.prompt": "hello",
        "ai.response.text": "world",
        "ai.prompt.messages": [{"role": "user", "content": "hello"}],
    }

    out = normalize_vercel_ai_span("ai.generateText.doGenerate", attrs)
    assert out["cascade.span_type"] == "llm"
    assert out["llm.model"] == "gpt-5"
    assert out["llm.provider"] == "openai"
    assert out["llm.input_tokens"] == 11
    assert out["llm.output_tokens"] == 22
    assert out["llm.total_tokens"] == 33
    assert out["llm.prompt"] == "hello"
    assert out["llm.completion"] == "world"
    assert out["llm.messages"] == [{"role": "user", "content": "hello"}]


def test_normalize_top_level_span_as_function():
    attrs = {"ai.operationId": "ai.generateText"}
    out = normalize_vercel_ai_span("ai.generateText", attrs)
    assert out["cascade.span_type"] == "function"


# ── functionId suffix tests ──────────────────────────────────────────────
# Vercel AI SDK appends `functionId` to span names, e.g.:
#   "ai.generateText.doGenerate My Agent v3.0"
# The old endswith/== checks broke on these names.


class TestFunctionIdSuffix:
    """Span type detection when Vercel AI SDK appends functionId to names."""

    def test_doGenerate_with_functionId_is_llm(self):
        name = "ai.generateText.doGenerate Research Signal Agent v3.0"
        assert _infer_cascade_span_type(name, {}) == "llm"

    def test_doStream_with_functionId_is_llm(self):
        name = "ai.streamText.doStream Extract Research from Helicone"
        assert _infer_cascade_span_type(name, {}) == "llm"

    def test_toolCall_with_functionId_is_tool(self):
        name = "ai.toolCall My Agent"
        assert _infer_cascade_span_type(name, {}) == "tool"

    def test_generateText_with_functionId_is_function(self):
        name = "ai.generateText Research Signal Agent v3.0"
        assert _infer_cascade_span_type(name, {}) == "function"

    def test_streamText_with_functionId_is_function(self):
        name = "ai.streamText My Custom Agent"
        assert _infer_cascade_span_type(name, {}) == "function"

    def test_bare_doGenerate_still_works(self):
        assert _infer_cascade_span_type("ai.generateText.doGenerate", {}) == "llm"

    def test_bare_doStream_still_works(self):
        assert _infer_cascade_span_type("ai.streamText.doStream", {}) == "llm"

    def test_bare_toolCall_still_works(self):
        assert _infer_cascade_span_type("ai.toolCall", {}) == "tool"

    def test_full_normalize_doGenerate_with_functionId(self):
        """End-to-end: doGenerate with functionId should get llm type + extracted data."""
        attrs = {
            "ai.operationId": "ai.generateText.doGenerate",
            "gen_ai.request.model": "claude-sonnet-4-20250514",
            "gen_ai.system": "anthropic",
            "gen_ai.usage.input_tokens": 500,
            "gen_ai.usage.output_tokens": 200,
            "ai.response.text": "Here is my analysis...",
            "ai.prompt.messages": '[{"role": "user", "content": "Analyze this"}]',
        }
        out = normalize_vercel_ai_span(
            "ai.generateText.doGenerate Research Signal Agent v3.0",
            attrs,
        )
        assert out["cascade.span_type"] == "llm"
        assert out["llm.model"] == "claude-sonnet-4-20250514"
        assert out["llm.provider"] == "anthropic"
        assert out["llm.input_tokens"] == 500
        assert out["llm.output_tokens"] == 200
        assert out["llm.total_tokens"] == 700
        assert out["llm.completion"] == "Here is my analysis..."
        assert out["llm.prompt"] == "Analyze this"

    def test_full_normalize_toolCall_with_functionId(self):
        """End-to-end: toolCall with functionId should get tool type + extracted data."""
        attrs = {
            "ai.operationId": "ai.toolCall",
            "ai.toolCall.name": "saveResearchSignals",
            "ai.toolCall.args": '{"signals": [1, 2, 3]}',
            "ai.toolCall.result": '{"saved": true}',
        }
        out = normalize_vercel_ai_span("ai.toolCall My Agent", attrs)
        assert out["cascade.span_type"] == "tool"
        assert out["tool.name"] == "saveResearchSignals"
        assert out["tool.input"] == '{"signals": [1, 2, 3]}'
        assert out["tool.output"] == '{"saved": true}'

    def test_function_wrapper_with_functionId_has_no_llm_data(self):
        """The ai.generateText wrapper span should be 'function', not 'llm'."""
        attrs = {
            "ai.operationId": "ai.generateText",
            "ai.usage.promptTokens": 500,
        }
        out = normalize_vercel_ai_span(
            "ai.generateText Extract Research from Helicone",
            attrs,
        )
        assert out["cascade.span_type"] == "function"

    def test_doGenerate_with_gen_ai_attrs_fallback(self):
        """Even without substring match, gen_ai.* attrs should identify as llm."""
        attrs = {
            "gen_ai.response.model": "gpt-4o",
        }
        out = normalize_vercel_ai_span("custom.span.name", attrs)
        assert out["cascade.span_type"] == "llm"


class TestToolDisplayNameWithFunctionId:
    """Simulate main.py display name logic for tool spans with functionId suffix."""

    @staticmethod
    def _get_display_name(span_name: str, attrs: dict) -> str:
        """Mirrors the logic in main.py lines 317-321."""
        span_display_name = span_name
        if span_name == "ai.toolCall" or span_name.startswith("ai.toolCall "):
            tool_name = attrs.get("tool.name") or attrs.get("ai.toolCall.name")
            if tool_name:
                span_display_name = tool_name
        return span_display_name

    def test_bare_toolCall_gets_tool_name(self):
        assert self._get_display_name(
            "ai.toolCall", {"tool.name": "weather"}
        ) == "weather"

    def test_toolCall_with_functionId_gets_tool_name(self):
        assert self._get_display_name(
            "ai.toolCall My Agent v3", {"tool.name": "saveResearchSignals"}
        ) == "saveResearchSignals"

    def test_toolCall_with_functionId_ai_toolCall_name_attr(self):
        assert self._get_display_name(
            "ai.toolCall Extract Research", {"ai.toolCall.name": "extractBehaviorSignals"}
        ) == "extractBehaviorSignals"

    def test_non_tool_span_keeps_original_name(self):
        assert self._get_display_name(
            "ai.generateText.doGenerate My Agent", {"tool.name": "ignored"}
        ) == "ai.generateText.doGenerate My Agent"

    def test_toolCall_without_tool_name_keeps_span_name(self):
        assert self._get_display_name(
            "ai.toolCall My Agent", {}
        ) == "ai.toolCall My Agent"

