import asyncio
from typing import Any, Dict, List, Mapping, Optional, Type, Union, cast

import pytest
from groq import AsyncGroq, Groq
from groq._base_client import _StreamT
from groq._types import Body, RequestFiles, RequestOptions, ResponseT
from groq.resources.chat.completions import AsyncCompletions, Completions
from groq.types.chat.chat_completion import (  # type: ignore[attr-defined]
    ChatCompletion,
    ChatCompletionMessage,
    Choice,
)
from groq.types.completion_usage import CompletionUsage
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
from opentelemetry.util._importlib_metadata import entry_points
from opentelemetry.util.types import AttributeValue

from openinference.instrumentation import OITracer, using_attributes
from openinference.instrumentation.groq import GroqInstrumentor
from openinference.semconv.trace import MessageAttributes, SpanAttributes

MOCK_COMPLETION = ChatCompletion(
    id="chat_comp_0",
    choices=[
        Choice(
            finish_reason="stop",
            index=0,
            logprobs=None,
            message=ChatCompletionMessage(
                content="idk, sorry!", role="assistant", function_call=None, tool_calls=None
            ),
        )
    ],
    created=1722531851,
    model="fake_model",
    object="chat.completion",
    system_fingerprint="fp0",
    usage=CompletionUsage(
        completion_tokens=379,
        prompt_tokens=25,
        total_tokens=404,
        completion_time=0.616262398,
        prompt_time=0.002549632,
        queue_time=None,
        total_time=0.6188120300000001,
    ),
)


def _mock_post(
    self: Any,
    path: str = "fake/url",
    *,
    cast_to: Type[ResponseT],
    body: Optional[Body] = None,
    options: RequestOptions = {},
    files: Optional[RequestFiles] = None,
    stream: bool = False,
    stream_cls: Optional[Type[_StreamT]] = None,
) -> Union[ResponseT, _StreamT]:
    """
    opts = FinalRequestOptions.construct(
        method="post", url=path, json_data=body, files=to_httpx_files(files), **options
    )
    return cast(ResponseT, self.request(cast_to, opts, stream=stream, stream_cls=stream_cls))
    """
    return cast(ResponseT, MOCK_COMPLETION)


async def _async_mock_post(
    self: Any,
    path: str = "fake/url",
    *,
    cast_to: Type[ResponseT],
    body: Optional[Body] = None,
    options: RequestOptions = {},
    files: Optional[RequestFiles] = None,
    stream: bool = False,
    stream_cls: Optional[Type[_StreamT]] = None,
) -> Union[ResponseT, _StreamT]:
    """
    opts = FinalRequestOptions.construct(
        method="post", url=path, json_data=body, files=to_httpx_files(files), **options
    )
    return cast(ResponseT, self.request(cast_to, opts, stream=stream, stream_cls=stream_cls))
    """
    return cast(ResponseT, MOCK_COMPLETION)


def _make_mock_post(finish_reason: str) -> Any:
    completion = ChatCompletion(
        id="chat_comp_0",
        choices=[
            Choice(
                finish_reason=finish_reason,  # type: ignore[arg-type]
                index=0,
                logprobs=None,
                message=ChatCompletionMessage(
                    content="hi", role="assistant", function_call=None, tool_calls=None
                ),
            )
        ],
        created=1722531851,
        model="fake_model",
        object="chat.completion",
        system_fingerprint="fp0",
        usage=CompletionUsage(
            completion_tokens=5,
            prompt_tokens=5,
            total_tokens=10,
        ),
    )

    def _post(
        self: Any,
        path: str = "fake/url",
        *,
        cast_to: Type[ResponseT],
        body: Optional[Body] = None,
        options: RequestOptions = {},
        files: Optional[RequestFiles] = None,
        stream: bool = False,
        stream_cls: Optional[Type[_StreamT]] = None,
    ) -> Union[ResponseT, _StreamT]:
        return cast(ResponseT, completion)

    return _post


@pytest.fixture()
def session_id() -> str:
    return "my-test-session-id"


@pytest.fixture()
def user_id() -> str:
    return "my-test-user-id"


@pytest.fixture()
def metadata() -> Dict[str, Any]:
    return {
        "test-int": 1,
        "test-str": "string",
        "test-list": [1, 2, 3],
        "test-dict": {
            "key-1": "val-1",
            "key-2": "val-2",
        },
    }


@pytest.fixture()
def tags() -> List[str]:
    return ["tag-1", "tag-2"]


@pytest.fixture
def prompt_template() -> str:
    return (
        "This is a test prompt template with int {var_int}, "
        "string {var_string}, and list {var_list}"
    )


@pytest.fixture
def prompt_template_version() -> str:
    return "v1.0"


@pytest.fixture
def prompt_template_variables() -> Dict[str, Any]:
    return {
        "var_int": 1,
        "var_str": "2",
        "var_list": [1, 2, 3],
    }


def _check_context_attributes(
    attributes: Any,
) -> None:
    assert attributes.get(SESSION_ID, None)
    assert attributes.get(USER_ID, None)
    assert attributes.get(METADATA, None)
    assert attributes.get(TAG_TAGS, None)
    assert attributes.get(LLM_PROMPT_TEMPLATE, None)
    assert attributes.get(LLM_PROMPT_TEMPLATE_VERSION, None)
    assert attributes.get(LLM_PROMPT_TEMPLATE_VARIABLES, None)


class TestInstrumentor:
    def test_entrypoint_for_opentelemetry_instrument(self) -> None:
        (instrumentor_entrypoint,) = entry_points(
            group="opentelemetry_instrumentor",
            name="groq",
        )
        instrumentor = instrumentor_entrypoint.load()()
        assert isinstance(instrumentor, GroqInstrumentor)

    # Ensure we're using the common OITracer from common openinference-instrumentation pkg
    def test_oitracer(self, setup_groq_instrumentation: Any) -> None:
        assert isinstance(GroqInstrumentor()._tracer, OITracer)


def test_groq_instrumentation(
    tracer_provider: TracerProvider,
    in_memory_span_exporter: InMemorySpanExporter,
    setup_groq_instrumentation: Any,
    session_id: str,
    user_id: str,
    metadata: Dict[str, Any],
    tags: List[str],
    prompt_template: str,
    prompt_template_version: str,
    prompt_template_variables: Dict[str, Any],
) -> None:
    client = Groq(api_key="fake")
    client.chat.completions._post = _mock_post  # type: ignore[assignment]

    with using_attributes(
        session_id=session_id,
        user_id=user_id,
        metadata=metadata,
        tags=tags,
        prompt_template=prompt_template,
        prompt_template_version=prompt_template_version,
        prompt_template_variables=prompt_template_variables,
    ):
        client.chat.completions.create(
            messages=[
                {
                    "role": "user",
                    "content": "Explain the importance of low latency LLMs",
                }
            ],
            model="fake_model",
        )
    spans = in_memory_span_exporter.get_finished_spans()

    assert spans[0].name == "Completions"
    assert spans[0].attributes and spans[0].attributes.get("openinference.span.kind") == "LLM"

    for span in spans:
        att = span.attributes
        _check_context_attributes(
            att,
        )

    attributes = dict(cast(Mapping[str, AttributeValue], span.attributes))
    assert (
        attributes[f"{SpanAttributes.LLM_INPUT_MESSAGES}.0.{MessageAttributes.MESSAGE_ROLE}"]
        == "user"
    )
    assert (
        attributes[f"{SpanAttributes.LLM_INPUT_MESSAGES}.0.{MessageAttributes.MESSAGE_CONTENT}"]
        == "Explain the importance of low latency LLMs"
    )
    assert (
        attributes[f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.0.{MessageAttributes.MESSAGE_ROLE}"]
        == "assistant"
    )
    assert (
        attributes[f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.0.{MessageAttributes.MESSAGE_CONTENT}"]
        == "idk, sorry!"
    )
    assert attributes[SpanAttributes.LLM_MODEL_NAME] == "fake_model"
    assert attributes[SpanAttributes.LLM_FINISH_REASON] == "stop"


def test_groq_async_instrumentation(
    tracer_provider: TracerProvider,
    in_memory_span_exporter: InMemorySpanExporter,
    setup_groq_instrumentation: Any,
    session_id: str,
    user_id: str,
    metadata: Dict[str, Any],
    tags: List[str],
    prompt_template: str,
    prompt_template_version: str,
    prompt_template_variables: Dict[str, Any],
) -> None:
    client = AsyncGroq(api_key="fake")
    client.chat.completions._post = _async_mock_post  # type: ignore[assignment]

    async def exec_comp() -> None:
        await client.chat.completions.create(
            messages=[
                {
                    "role": "user",
                    "content": "Explain the importance of low latency LLMs",
                }
            ],
            model="fake_model",
        )

    with using_attributes(
        session_id=session_id,
        user_id=user_id,
        metadata=metadata,
        tags=tags,
        prompt_template=prompt_template,
        prompt_template_version=prompt_template_version,
        prompt_template_variables=prompt_template_variables,
    ):
        asyncio.run(exec_comp())

    spans = in_memory_span_exporter.get_finished_spans()

    assert spans[0].name == "AsyncCompletions"
    assert spans[0].attributes and spans[0].attributes.get("openinference.span.kind") == "LLM"

    for span in spans:
        att = span.attributes
        _check_context_attributes(
            att,
        )

    attributes = dict(cast(Mapping[str, AttributeValue], span.attributes))
    print(attributes)
    assert (
        attributes[f"{SpanAttributes.LLM_INPUT_MESSAGES}.0.{MessageAttributes.MESSAGE_ROLE}"]
        == "user"
    )
    assert (
        attributes[f"{SpanAttributes.LLM_INPUT_MESSAGES}.0.{MessageAttributes.MESSAGE_CONTENT}"]
        == "Explain the importance of low latency LLMs"
    )
    assert (
        attributes[f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.0.{MessageAttributes.MESSAGE_ROLE}"]
        == "assistant"
    )
    assert (
        attributes[f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.0.{MessageAttributes.MESSAGE_CONTENT}"]
        == "idk, sorry!"
    )
    assert attributes[SpanAttributes.LLM_MODEL_NAME] == "fake_model"
    assert attributes[SpanAttributes.LLM_FINISH_REASON] == "stop"


@pytest.mark.parametrize("finish_reason", ["stop", "length", "tool_calls", "function_call"])
def test_finish_reason_values(
    finish_reason: str,
    tracer_provider: TracerProvider,
    in_memory_span_exporter: InMemorySpanExporter,
    setup_groq_instrumentation: Any,
) -> None:
    client = Groq(api_key="fake")
    client.chat.completions._post = _make_mock_post(finish_reason)

    client.chat.completions.create(
        messages=[{"role": "user", "content": "hello"}],
        model="fake_model",
    )
    spans = in_memory_span_exporter.get_finished_spans()
    assert len(spans) == 1
    attributes = dict(cast(Mapping[str, AttributeValue], spans[0].attributes))
    assert attributes.get(SpanAttributes.LLM_FINISH_REASON) == finish_reason


def test_groq_uninstrumentation(
    tracer_provider: TracerProvider,
) -> None:
    GroqInstrumentor().instrument(tracer_provider=tracer_provider)

    assert hasattr(Completions.create, "__wrapped__")
    assert hasattr(AsyncCompletions.create, "__wrapped__")

    GroqInstrumentor().uninstrument()

    assert not hasattr(Completions.create, "__wrapped__")
    assert not hasattr(AsyncCompletions.create, "__wrapped__")


# Ensure we're using the common OITracer from common opeinference-instrumentation pkg
def test_oitracer(
    setup_groq_instrumentation: Any,
) -> None:
    assert isinstance(GroqInstrumentor()._tracer, OITracer)


SESSION_ID = SpanAttributes.SESSION_ID
USER_ID = SpanAttributes.USER_ID
METADATA = SpanAttributes.METADATA
TAG_TAGS = SpanAttributes.TAG_TAGS
LLM_PROMPT_TEMPLATE = SpanAttributes.LLM_PROMPT_TEMPLATE
LLM_PROMPT_TEMPLATE_VERSION = SpanAttributes.LLM_PROMPT_TEMPLATE_VERSION
LLM_PROMPT_TEMPLATE_VARIABLES = SpanAttributes.LLM_PROMPT_TEMPLATE_VARIABLES
