import importlib.util
import os
import warnings
from dataclasses import dataclass
from typing import Any, Callable, Dict, List
from unittest.mock import MagicMock

import pytest

from guardrails.llm_providers import (
    ArbitraryCallable,
    AsyncArbitraryCallable,
    LLMResponse,
    PromptCallableException,
    chat_prompt,
    get_async_llm_ask,
    get_llm_ask,
)
from guardrails.utils.safe_get import safe_get_with_brackets

from .mocks import MockAsyncOpenAILlm, MockOpenAILlm


def test_openai_callable_does_not_retry_on_success(mocker):
    llm = MockOpenAILlm()
    succeed_spy = mocker.spy(llm, "succeed")

    arbitrary_callable = ArbitraryCallable(
        llm.succeed, messages=[{"role": "user", "content": "Hello"}]
    )
    response = arbitrary_callable()

    assert succeed_spy.call_count == 1
    assert isinstance(response, LLMResponse) is True
    assert response.output == "Hello world!"
    assert response.prompt_token_count is None
    assert response.response_token_count is None


@pytest.mark.asyncio
async def test_async_openai_callable_does_not_retry_on_success(mocker):
    llm = MockAsyncOpenAILlm()
    succeed_spy = mocker.spy(llm, "succeed")

    arbitrary_callable = AsyncArbitraryCallable(
        llm.succeed, messages=[{"role": "user", "content": "Hello"}]
    )
    response = await arbitrary_callable()

    assert succeed_spy.call_count == 1
    assert isinstance(response, LLMResponse) is True
    assert response.output == "Hello world!"
    assert response.prompt_token_count is None
    assert response.response_token_count is None


@pytest.fixture(scope="module")
def openai_chat_mock():
    from openai.types import CompletionUsage
    from openai.types.chat import ChatCompletion, ChatCompletionMessage
    from openai.types.chat.chat_completion import Choice

    return ChatCompletion(
        id="",
        choices=[
            Choice(
                finish_reason="stop",
                index=0,
                message=ChatCompletionMessage(
                    content="Mocked LLM output",
                    role="assistant",
                ),
            ),
        ],
        created=0,
        model="",
        object="chat.completion",
        usage=CompletionUsage(
            completion_tokens=20,
            prompt_tokens=10,
            total_tokens=30,
        ),
    )


@pytest.fixture(scope="module")
def openai_chat_stream_mock():
    def gen():
        # Returns a generator object
        for i in range(4, 8):
            yield {
                "choices": [
                    {
                        "index": 0,
                        "delta": {"content": f"{i},"},
                        "finish_reason": None,
                    }
                ]
            }

    return gen()


@pytest.fixture(scope="module")
def openai_mock():
    @dataclass
    class MockCompletionUsage:
        completion_tokens: int
        prompt_tokens: int
        total_tokens: int

    @dataclass
    class MockCompletionChoice:
        finish_reason: str
        index: int
        logprobs: Any
        text: str

    @dataclass
    class MockCompletion:
        id: str
        choices: List[MockCompletionChoice]
        created: int
        model: str
        object: str
        usage: MockCompletionUsage

    return MockCompletion(
        id="",
        choices=[
            MockCompletionChoice(
                finish_reason="stop",
                index=0,
                logprobs=None,
                text="Mocked LLM output",
            ),
        ],
        created=0,
        model="",
        object="text_completion",
        usage=MockCompletionUsage(
            completion_tokens=20,
            prompt_tokens=10,
            total_tokens=30,
        ),
    )


@pytest.fixture(scope="module")
def openai_stream_mock():
    def gen():
        # Returns a generator object
        for i in range(4, 8):
            yield {
                "choices": [{"text": f"{i},", "finish_reason": None}],
                "model": "openai-model-name",
            }

    return gen()


@pytest.mark.skipif(
    not importlib.util.find_spec("manifest"),
    reason="manifest-ml is not installed",
)
def test_manifest_callable():
    client = MagicMock()
    client.run.return_value = "Hello world!"

    from guardrails.llm_providers import ManifestCallable

    manifest_callable = ManifestCallable()
    response = manifest_callable(text="Hello", client=client)

    assert isinstance(response, LLMResponse) is True
    assert response.output == "Hello world!"
    assert response.prompt_token_count is None
    assert response.response_token_count is None


@pytest.mark.skipif(
    not importlib.util.find_spec("manifest"),
    reason="manifest-ml is not installed",
)
@pytest.mark.asyncio
async def test_async_manifest_callable():
    client = MagicMock()

    async def return_async():
        return ["Hello world!"]

    client.arun_batch.return_value = return_async()

    from guardrails.llm_providers import AsyncManifestCallable

    manifest_callable = AsyncManifestCallable()
    response = await manifest_callable(text="Hello", client=client)

    assert isinstance(response, LLMResponse) is True
    assert response.output == "Hello world!"
    assert response.prompt_token_count is None
    assert response.response_token_count is None


@pytest.mark.skipif(
    not importlib.util.find_spec("transformers")
    and not importlib.util.find_spec("torch"),
    reason="transformers or torch is not installed",
)
@pytest.mark.parametrize(
    "model_inputs,tokenizer_call_count", [(None, 1), ({"input_ids": ["Hello"]}, 0)]
)
def test_hugging_face_model_callable(mocker, model_inputs, tokenizer_call_count):
    class MockTokenizer:
        def __call__(self, prompt: str, *args: Any, **kwds: Any) -> Dict[str, Any]:
            self.prompt = prompt
            return self

        def to(self, *args, **kwargs):
            return {"input_ids": [self.prompt]}

        def decode(self, output: str, *args, **kwargs) -> str:
            return output

    tokenizer = MockTokenizer()

    tokenizer_call_spy = mocker.spy(tokenizer, "to")
    tokenizer_decode_spy = mocker.spy(tokenizer, "decode")

    model_generate = MagicMock()
    model_generate.return_value = ["Hello there!"]

    from guardrails.llm_providers import HuggingFaceModelCallable

    hf_model_callable = HuggingFaceModelCallable()
    response = hf_model_callable(
        model_generate=model_generate,
        messages=[{"role": "user", "content": "Hello"}],
        tokenizer=tokenizer,
    )

    assert tokenizer_call_spy.call_count == 1
    assert tokenizer_decode_spy.call_count == 1
    assert isinstance(response, LLMResponse) is True
    assert response.output == "Hello there!"
    assert response.prompt_token_count is None
    assert response.response_token_count is None


@pytest.mark.skipif(
    not importlib.util.find_spec("transformers")
    and not importlib.util.find_spec("torch"),
    reason="transformers or torch is not installed",
)
def test_hugging_face_pipeline_callable():
    pipeline = MagicMock()
    pipeline.return_value = [{"generated_text": "Hello there!"}]

    from guardrails.llm_providers import HuggingFacePipelineCallable

    hf_model_callable = HuggingFacePipelineCallable()
    response = hf_model_callable(
        pipeline=pipeline, messages=[{"role": "user", "content": "Hello"}]
    )

    assert isinstance(response, LLMResponse) is True
    assert response.output == "Hello there!"
    assert response.prompt_token_count is None
    assert response.response_token_count is None


@pytest.mark.skipif(
    not importlib.util.find_spec("litellm"),
    reason="`litellm` is not installed",
)
def test_litellm_callable(mocker):
    # Mock the litellm.completion function and
    # the classes it returns
    @dataclass
    class Message:
        content: str

    @dataclass
    class Choice:
        message: Message

    @dataclass
    class Usage:
        prompt_tokens: int
        completion_tokens: int

    @dataclass
    class MockResponse:
        choices: List[Choice]
        usage: Usage

    class MockCompletion:
        @staticmethod
        def create() -> MockResponse:
            return MockResponse(
                choices=[Choice(message=Message(content="Hello there!"))],
                usage=Usage(prompt_tokens=10, completion_tokens=20),
            )

    mocker.patch("litellm.completion", return_value=MockCompletion.create())

    from guardrails.llm_providers import LiteLLMCallable

    litellm_callable = LiteLLMCallable()
    response = litellm_callable("Hello")

    assert isinstance(response, LLMResponse) is True
    assert response.output == "Hello there!"
    assert response.prompt_token_count == 10
    assert response.response_token_count == 20


class ReturnTempCallable(Callable):
    def __call__(self, *args, messages=None, **kwargs) -> Any:
        return ""


@pytest.mark.parametrize(
    "llm_api, args, kwargs, expected_temperature",
    [
        (ReturnTempCallable(), [], {"temperature": 0.5}, 0.5),
        (ReturnTempCallable(), [], {}, 0),
        (ReturnTempCallable(), [], {"model": "gpt-5-nano"}, None),
    ],
)
def test_get_llm_ask_temperature(llm_api, args, kwargs, expected_temperature):
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        result = get_llm_ask(llm_api, *args, **kwargs)
        if expected_temperature is None:
            assert "temperature" not in result.init_kwargs
            assert len(w) == 0
        else:
            assert "temperature" in result.init_kwargs
            assert result.init_kwargs["temperature"] == expected_temperature
            if "temperature" not in kwargs:
                assert len(w) == 1
                assert issubclass(w[0].category, DeprecationWarning)
                assert "default value of 0 for temperature is deprecated" in str(
                    w[0].message
                )


@pytest.mark.skipif(
    not importlib.util.find_spec("manifest"),
    reason="manifest is not installed",
)
def test_get_llm_ask_manifest(mocker):
    def mock_os_environ_get(key, *args):
        if key == "OPENAI_API_KEY":
            return "sk-xxxxxxxxxxxxxx"
        return safe_get_with_brackets(os.environ, key, *args)

    mocker.patch("os.environ.get", side_effect=mock_os_environ_get)

    from manifest import Manifest

    from guardrails.llm_providers import ManifestCallable

    manifest_client = Manifest("openai")

    prompt_callable = get_llm_ask(manifest_client)

    assert isinstance(prompt_callable, ManifestCallable)


@pytest.mark.skipif(
    not importlib.util.find_spec("transformers"),
    reason="transformers is not installed",
)
def test_get_llm_ask_hugging_face_model(mocker):
    from transformers import PreTrainedModel, GenerationMixin

    from guardrails.llm_providers import HuggingFaceModelCallable

    class MockModel(PreTrainedModel, GenerationMixin):
        _modules: Any

        def __init__(self, *args, **kwargs):
            self._modules = {}

    mock_model = MockModel()

    prompt_callable = get_llm_ask(mock_model.generate)

    assert isinstance(prompt_callable, HuggingFaceModelCallable)


@pytest.mark.skipif(
    not importlib.util.find_spec("transformers"),
    reason="transformers is not installed",
)
def test_get_llm_ask_hugging_face_pipeline():
    from transformers import Pipeline

    from guardrails.llm_providers import HuggingFacePipelineCallable

    class MockPipeline(Pipeline):
        task = "text-generation"

        def __init__(self, *args, **kwargs):
            pass

        def _forward():
            pass

        def _sanitize_parameters():
            pass

        def postprocess():
            pass

        def preprocess():
            pass

    mock_pipeline = MockPipeline()

    prompt_callable = get_llm_ask(mock_pipeline)

    assert isinstance(prompt_callable, HuggingFacePipelineCallable)


@pytest.mark.skipif(
    not importlib.util.find_spec("litellm"),
    reason="`litellm` is not installed",
)
def test_get_llm_ask_litellm():
    from litellm import completion

    from guardrails.llm_providers import LiteLLMCallable

    prompt_callable = get_llm_ask(completion)

    assert isinstance(prompt_callable, LiteLLMCallable)


def test_get_llm_ask_custom_llm():
    from guardrails.llm_providers import ArbitraryCallable

    def my_llm(prompt: str, *, messages=None, **kwargs) -> str:
        return f"Hello {prompt}!"

    prompt_callable = get_llm_ask(my_llm)

    assert isinstance(prompt_callable, ArbitraryCallable)


def test_get_llm_ask_custom_llm_warning():
    from guardrails.llm_providers import ArbitraryCallable

    def my_llm(prompt: str, **kwargs) -> str:
        return f"Hello {prompt}!"

    with pytest.warns(
        UserWarning,
        match=(
            "We recommend including 'messages'"
            " as keyword-only arguments for custom LLM callables."
            " Doing so ensures these arguments are not unintentionally"
            " passed through to other calls via \\*\\*kwargs."
        ),
    ):
        prompt_callable = get_llm_ask(my_llm)

        assert isinstance(prompt_callable, ArbitraryCallable)


def test_get_llm_ask_custom_llm_must_accept_kwargs():
    def my_llm(messages: str) -> str:
        return f"Hello {messages}!"

    with pytest.raises(
        ValueError, match="Custom LLM callables must accept \\*\\*kwargs!"
    ):
        get_llm_ask(my_llm)


def test_get_async_llm_ask_custom_llm():
    from guardrails.llm_providers import AsyncArbitraryCallable

    async def my_llm(messages: str, **kwargs) -> str:
        return f"Hello {messages}!"

    prompt_callable = get_async_llm_ask(my_llm)

    assert isinstance(prompt_callable, AsyncArbitraryCallable)


def test_get_async_llm_ask_custom_llm_warning():
    from guardrails.llm_providers import AsyncArbitraryCallable

    async def my_llm(**kwargs) -> str:
        return "Hello world!"

    with pytest.warns(
        UserWarning,
        match=(
            "We recommend including 'messages'"
            " as keyword-only arguments for custom LLM callables."
            " Doing so ensures these arguments are not unintentionally"
            " passed through to other calls via \\*\\*kwargs."
        ),
    ):
        prompt_callable = get_async_llm_ask(my_llm)

        assert isinstance(prompt_callable, AsyncArbitraryCallable)


def test_get_async_llm_ask_custom_llm_must_accept_kwargs():
    def my_llm(prompt: str) -> str:
        return f"Hello {prompt}!"

    with pytest.raises(
        ValueError, match="Custom LLM callables must accept \\*\\*kwargs!"
    ):
        get_async_llm_ask(my_llm)


def test_chat_prompt():
    # raises when messages are not provided
    with pytest.raises(PromptCallableException):
        chat_prompt(None)