# Copyright 2023 OpenSPG Authors
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
# in compliance with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software distributed under the License
# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
# or implied.

import sys
import unittest
import unittest.mock
from dataclasses import dataclass

from nn4k.invoker import NNInvoker


@dataclass
class MockCompletion:
    choices: list


@dataclass
class MockMessage:
    content: str


@dataclass
class MockChoice:
    message: MockMessage


@dataclass
class MockEmbeddings:
    data: list


@dataclass
class MockEmbedding:
    embedding: list


class TestOpenAIInvoker(unittest.TestCase):
    """
    OpenAIInvoker unittest
    """

    def setUp(self):
        self._saved_openai = sys.modules.get("openai")
        self._mocked_openai = unittest.mock.MagicMock()
        sys.modules["openai"] = self._mocked_openai

    def tearDown(self):
        del sys.modules["openai"]
        if self._saved_openai is not None:
            sys.modules["openai"] = self._saved_openai

    def testOpenAICompletion(self):
        self._mocked_openai.__version__ = "1.7.0"
        self._mocked_openai.OpenAI = unittest.mock.MagicMock

        nn_config = {
            "nn_name": "gpt-3.5-turbo",
            "openai_api_key": "EMPTY",
            "openai_api_base": "http://localhost:38080/v1",
            "openai_max_tokens": 2000,
        }
        invoker = NNInvoker.from_config(nn_config)
        self.assertEqual(invoker.init_args, nn_config)
        self.assertEqual(invoker.client.api_key, nn_config["openai_api_key"])
        self.assertEqual(invoker.client.base_url, nn_config["openai_api_base"])

        mock_completion = MockCompletion(
            choices=[MockChoice(message=MockMessage(content="a dog named Bolt ..."))]
        )
        invoker.client.chat.completions.create.return_value = mock_completion

        result = invoker.remote_inference("Long long ago, ")
        invoker.client.chat.completions.create.assert_called_with(
            model=nn_config["nn_name"],
            messages=[{"role": "user", "content": "Long long ago, "}],
            max_tokens=nn_config["openai_max_tokens"],
        )
        self.assertEqual(result, [mock_completion.choices[0].message.content])

    def testOpenAIEmbedding(self):
        self._mocked_openai.__version__ = "1.7.0"
        self._mocked_openai.OpenAI = unittest.mock.MagicMock

        nn_config = {
            "nn_name": "text-embedding-ada-002",
            "openai_api_key": "EMPTY",
            "openai_api_base": "http://localhost:38080/v1",
        }
        invoker = NNInvoker.from_config(nn_config)
        self.assertEqual(invoker.init_args, nn_config)
        self.assertEqual(invoker.client.api_key, nn_config["openai_api_key"])
        self.assertEqual(invoker.client.base_url, nn_config["openai_api_base"])

        mock_embeddings = MockEmbeddings(data=[MockEmbedding(embedding=[0.1, 0.2])])
        invoker.client.embeddings.create.return_value = mock_embeddings

        result = invoker.remote_inference("How old are you?", type="Embedding")
        invoker.client.embeddings.create.assert_called_with(
            model=nn_config["nn_name"],
            input=["How old are you?"],
        )
        self.assertEqual(result, [mock_embeddings.data[0].embedding])


if __name__ == "__main__":
    unittest.main()