import json
import os
import random
import string
from importlib import import_module
from importlib.metadata import version
from typing import Tuple, cast

import pytest
from opentelemetry import trace as trace_api
from opentelemetry.sdk.trace import ReadableSpan
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter

from openinference.semconv.trace import OpenInferenceLLMProviderValues, SpanAttributes


@pytest.mark.vcr
def test_tool_calls(
    in_memory_span_exporter: InMemorySpanExporter,
    tracer_provider: trace_api.TracerProvider,
) -> None:
    if _openai_version() < (1, 12, 0):
        pytest.skip("Not supported")
    openai = import_module("openai")
    from openai.types.chat import (  # type: ignore[attr-defined,unused-ignore]
        ChatCompletionToolParam,
    )

    client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY", "sk-"))
    input_tools = [
        ChatCompletionToolParam(
            type="function",
            function={
                "name": "get_weather",
                "description": "finds the weather for a given city",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "city": {
                            "type": "string",
                            "description": "The city to find the weather for, e.g. 'London'",
                        }
                    },
                    "required": ["city"],
                },
            },
        ),
        ChatCompletionToolParam(
            type="function",
            function={
                "name": "get_population",
                "description": "finds the population for a given city",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "city": {
                            "type": "string",
                            "description": "The city to find the population for, e.g. 'London'",
                        }
                    },
                    "required": ["city"],
                },
            },
        ),
    ]
    client.chat.completions.create(
        extra_headers={"Accept-Encoding": "gzip"},
        model="gpt-4o-mini",
        tools=input_tools,
        messages=[
            {
                "role": "assistant",
                "tool_calls": [
                    {
                        "id": "call_62136355",
                        "type": "function",
                        "function": {"name": "get_weather", "arguments": '{"city": "New York"}'},
                    },
                    {
                        "id": "call_62136356",
                        "type": "function",
                        "function": {"name": "get_population", "arguments": '{"city": "New York"}'},
                    },
                ],
            },
            {
                "role": "tool",
                "tool_call_id": "call_62136355",
                "content": '{"city": "New York", "weather": "fine"}',
            },
            {
                "role": "tool",
                "tool_call_id": "call_62136356",
                "content": '{"city": "New York", "weather": "large"}',
            },
            {
                "role": "assistant",
                "content": "In New York the weather is fine and the population is large.",
            },
            {
                "role": "user",
                "content": "What's the weather and population in San Francisco?",
            },
        ],
    )
    spans = get_openai_llm_spans(in_memory_span_exporter.get_finished_spans())
    assert len(spans) == 1
    span = spans[0]
    attributes = dict(span.attributes or {})
    for i in range(len(input_tools)):
        json_schema = attributes.pop(f"llm.tools.{i}.tool.json_schema")
        assert isinstance(json_schema, str)
        assert json.loads(json_schema)
    assert (
        attributes.pop("llm.input_messages.0.message.tool_calls.0.tool_call.id") == "call_62136355"
    )
    assert (
        attributes.pop("llm.input_messages.0.message.tool_calls.0.tool_call.function.name")
        == "get_weather"
    )
    assert (
        attributes.pop("llm.input_messages.0.message.tool_calls.0.tool_call.function.arguments")
        == '{"city": "New York"}'
    )
    assert (
        attributes.pop("llm.input_messages.0.message.tool_calls.1.tool_call.id") == "call_62136356"
    )
    assert (
        attributes.pop("llm.input_messages.0.message.tool_calls.1.tool_call.function.name")
        == "get_population"
    )
    assert (
        attributes.pop("llm.input_messages.0.message.tool_calls.1.tool_call.function.arguments")
        == '{"city": "New York"}'
    )
    assert attributes.pop("llm.input_messages.1.message.role") == "tool"
    assert attributes.pop("llm.input_messages.1.message.tool_call_id") == "call_62136355"
    assert (
        attributes.pop("llm.input_messages.1.message.content")
        == '{"city": "New York", "weather": "fine"}'
    )
    assert attributes.pop("llm.input_messages.2.message.role") == "tool"
    assert attributes.pop("llm.input_messages.2.message.tool_call_id") == "call_62136356"
    assert (
        attributes.pop("llm.input_messages.2.message.content")
        == '{"city": "New York", "weather": "large"}'
    )
    assert attributes.pop("llm.output_messages.0.message.tool_calls.0.tool_call.id")
    assert (
        attributes.pop("llm.output_messages.0.message.tool_calls.0.tool_call.function.name")
        == "get_weather"
    )
    assert (
        attributes.pop("llm.output_messages.0.message.tool_calls.0.tool_call.function.arguments")
        == '{"city": "San Francisco"}'
    )
    assert attributes.pop("llm.output_messages.0.message.tool_calls.1.tool_call.id")
    assert (
        attributes.pop("llm.output_messages.0.message.tool_calls.1.tool_call.function.name")
        == "get_population"
    )
    assert (
        attributes.pop("llm.output_messages.0.message.tool_calls.1.tool_call.function.arguments")
        == '{"city": "San Francisco"}'
    )


@pytest.mark.vcr
def test_cached_tokens(
    in_memory_span_exporter: InMemorySpanExporter,
    tracer_provider: trace_api.TracerProvider,
) -> None:
    if _openai_version() < (1, 12, 0):
        pytest.skip("Not supported")
    openai = import_module("openai")

    client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY", "sk-"))
    random_1024_token_prefix = "".join(random.choices(string.ascii_letters + string.digits, k=2000))
    client.chat.completions.create(
        extra_headers={"Accept-Encoding": "gzip"},
        model="gpt-4o-mini",
        messages=[
            {
                "role": "user",
                "content": "{} Write me a haiku.".format(random_1024_token_prefix),
            },
        ],
    )
    client.chat.completions.create(
        extra_headers={"Accept-Encoding": "gzip"},
        model="gpt-4o-mini",
        messages=[
            {
                "role": "user",
                "content": "{} Write me a sonnet.".format(random_1024_token_prefix),
            },
        ],
    )
    spans = tuple(
        span
        for span in get_openai_llm_spans(in_memory_span_exporter.get_finished_spans())
        if random_1024_token_prefix
        in str((span.attributes or {}).get(SpanAttributes.INPUT_VALUE, ""))
    )
    assert len(spans) == 2
    cache_reads = []
    for span in spans:
        cache_read = dict(span.attributes or {}).get("llm.token_count.prompt_details.cache_read")
        assert isinstance(cache_read, int)
        cache_reads.append(cache_read)
    assert sorted(cache_reads) == [0, 1280]


def _openai_version() -> Tuple[int, int, int]:
    return cast(Tuple[int, int, int], tuple(map(int, version("openai").split(".")[:3])))


def get_openai_llm_spans(spans: Tuple[ReadableSpan, ...]) -> Tuple[ReadableSpan, ...]:
    """Filter spans to only the primary OpenAI LLM response spans to avoid extra internal spans."""
    llm_spans = [
        span
        for span in spans
        if span.attributes
        and span.attributes.get(SpanAttributes.LLM_PROVIDER)
        == OpenInferenceLLMProviderValues.OPENAI.value
    ]
    if not llm_spans:
        raise ValueError("No OpenAI LLM spans found in spans.")
    return tuple(llm_spans)
