# -*- coding: utf-8 -*-
"""dashscope test"""
import unittest
from unittest.mock import patch, MagicMock

from agentscope.models import (
    ModelResponse,
    DashScopeChatWrapper,
    DashScopeImageSynthesisWrapper,
    DashScopeTextEmbeddingWrapper,
    DashScopeMultiModalWrapper,
)


class TestDashScopeChatWrapper(unittest.TestCase):
    """Test DashScope Chat Wrapper"""

    def setUp(self) -> None:
        self.config_name = "test_config"
        self.model_name = "test_model"
        self.api_key = "test_api_key"
        self.wrapper = DashScopeChatWrapper(
            config_name=self.config_name,
            model_name=self.model_name,
            api_key=self.api_key,
        )

    @patch("agentscope.models.dashscope_model.dashscope.Generation.call")
    def test_call_success(self, mock_generation_call: MagicMock) -> None:
        """Test call success"""
        # Set up the mock response for a successful API call
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.request_id = "test_request_id"
        mock_response.usage = {
            "input_tokens": 3,
            "output_tokens": 5,
            "total_tokens": 8,
        }
        mock_response.output = {
            "choices": [{"message": {"content": "Hello, world!"}}],
        }

        mock_generation_call.return_value = mock_response

        # Define test input
        messages = [
            {"role": "user", "content": "Hi!"},
            {"role": "assistant", "content": "Hello!"},
        ]

        # Call the wrapper method
        response = self.wrapper(messages)

        # Verify the response
        self.assertIsInstance(response, ModelResponse)
        self.assertEqual(response.text, "Hello, world!")
        self.assertEqual(response.raw, mock_response)

        # Verify call to dashscope.Generation.call
        mock_generation_call.assert_called_once_with(
            model=self.model_name,
            messages=messages,
            result_format="message",
        )

    @patch("agentscope.models.dashscope_model.dashscope.Generation.call")
    def test_call_failure(self, mock_generation_call: MagicMock) -> None:
        """test call failure"""
        # Set up the mock response for a failed API call
        mock_response = MagicMock()
        mock_response.status_code = "400"
        mock_response.request_id = "Test_request_id"
        mock_response.code = "Error Code"
        mock_response.message = "Error Message"
        mock_generation_call.return_value = mock_response

        # Define test input
        messages = [
            {"role": "user", "content": "Hi!"},
            {"role": "assistant", "content": "Hello!"},
        ]

        # Call the wrapper method and expect an exception
        with self.assertRaises(RuntimeError) as context:
            self.wrapper(messages)

        # Assert the expected exception message contains the error details
        self.assertIn("Error Code", str(context.exception))
        self.assertIn("Error Message", str(context.exception))
        self.assertIn("400", str(context.exception))
        self.assertIn("Test_request_id", str(context.exception))

        # Verify call to dashscope.Generation.call
        mock_generation_call.assert_called_once_with(
            model=self.model_name,
            messages=messages,
            result_format="message",
        )


class TestDashScopeImageSynthesisWrapper(unittest.TestCase):
    """Test DashScope Image Synthesis Wrapper"""

    def setUp(self) -> None:
        self.config_name = "config_name"
        self.model_name = "test_model"
        self.api_key = "test_api_key"
        self.wrapper = DashScopeImageSynthesisWrapper(
            config_name=self.config_name,
            model_name=self.model_name,
            api_key=self.api_key,
        )

    @patch(
        "agentscope.file_manager.file_manager.save_image",
        side_effect=lambda x: f'/local/path/{x.split("/")[-1]}',
    )
    @patch("agentscope.models.dashscope_model.dashscope.ImageSynthesis.call")
    def test_image_synthesis_wrapper_call_success(
        self,
        mock_call: MagicMock,
        mock_save_image: MagicMock,
    ) -> None:
        """Test call success"""
        # Set up the mock response for a successful API call
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.output = {
            "results": [{"url": "http://example.com/image.jpg"}],
        }
        mock_call.return_value = mock_response
        # Call the wrapper with prompt
        prompt = "Generate an image of a sunset"
        response = self.wrapper(prompt, save_local=False)
        # Check if response is an instance of ModelResponse and contains
        # expected data
        self.assertIsInstance(response, ModelResponse)
        self.assertEqual(
            response.image_urls,
            ["http://example.com/image.jpg"],
        )

        # Call the wrapper method with save_local set to True
        response_with_save = self.wrapper(prompt, save_local=True)

        # Verify save_image call and local image url in response
        mock_save_image.assert_called_with("http://example.com/image.jpg")
        self.assertEqual(
            response_with_save.image_urls,
            ["/local/path/image.jpg"],
        )

    @patch("agentscope.models.dashscope_model.dashscope.ImageSynthesis.call")
    def test_image_synthesis_wrapper_call_failure(
        self,
        mock_call: MagicMock,
    ) -> None:
        """Test call failure"""
        # Set up the mock response for a failed API call
        mock_response = MagicMock()
        mock_response.status_code = 400
        mock_response.request_id = "Test_request_id"
        mock_response.code = "Error Code"
        mock_response.message = "Error Message"
        mock_call.return_value = mock_response

        # Call the wrapper with prompt and expect a RuntimeError
        prompt = "Generate an image of a sunset"
        with self.assertRaises(RuntimeError) as context:
            self.wrapper(prompt, save_local=False, n=1)

        # Assert the expected exception message
        self.assertIn("Error Code", str(context.exception))
        self.assertIn("Error Message", str(context.exception))
        self.assertIn("Test_request_id", str(context.exception))
        self.assertIn("400", str(context.exception))

        # Verify call to dashscope.ImageSynthesis.call
        mock_call.assert_called_once_with(
            model=self.model_name,
            prompt=prompt,
            n=1,  # Assuming this is a default value used to call the API
        )


class TestDashScopeTextEmbeddingWrapper(unittest.TestCase):
    """Test DashScope Text Embedding Wrapper"""

    def setUp(self) -> None:
        # Initialize DashScopeTextEmbeddingWrapper instance
        self.wrapper = DashScopeTextEmbeddingWrapper(
            config_name="test_config",
            model_name="test_model",
            api_key="test_key",
        )

    @patch("agentscope.models.dashscope_model.dashscope.TextEmbedding.call")
    def test_call_success(self, mock_call: MagicMock) -> None:
        """Test call success"""
        # Mock a successful API response
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.usage = {
            "total_tokens": 8,
        }
        mock_response.output = {
            "embeddings": [{"embedding": [0.1, 0.2, 0.3]}],
        }
        mock_call.return_value = mock_response

        # Call the wrapper with a mock API
        texts = ["Hello, world!"]
        response = self.wrapper(texts)

        # Assert the API was called correctly
        mock_call.assert_called_once_with(
            input=texts,
            model=self.wrapper.model_name,
            **self.wrapper.generate_args,
        )

        # Assert the response is as expected
        self.assertEqual(response.embedding, [[0.1, 0.2, 0.3]])

    @patch("agentscope.models.dashscope_model.dashscope.TextEmbedding.call")
    def test_call_failure(self, mock_call: MagicMock) -> None:
        """Test call failure"""
        # Mock a failed API response
        mock_response = MagicMock()
        mock_response.status_code = 400
        mock_response.request_id = "Test_request_id"
        mock_response.code = "Error Code"
        mock_response.message = "Error Message"
        mock_call.return_value = mock_response

        # Call the wrapper with a mock API and expect an exception
        texts = ["Hello, world!"]
        with self.assertRaises(RuntimeError) as context:
            self.wrapper(texts)

        # Assert the expected exception message
        self.assertIn("Error Code", str(context.exception))
        self.assertIn("Error Message", str(context.exception))
        self.assertIn("Test_request_id", str(context.exception))
        self.assertIn("400", str(context.exception))

        # Verify call to dashscope.TextEmbedding.call
        mock_call.assert_called_once_with(
            input=texts,
            model=self.wrapper.model_name,
            **self.wrapper.generate_args,
        )


class TestDashScopeMultiModalWrapper(unittest.TestCase):
    """Test DashScope MultiModal Wrapper"""

    def setUp(self) -> None:
        # Initialize DashScopeMultiModalWrapper instance
        self.wrapper = DashScopeMultiModalWrapper(
            config_name="test_config",
            model_name="test_model",
            api_key="test_key",
        )

    @patch(
        "agentscope.models.dashscope_model."
        "dashscope.MultiModalConversation.call",
    )
    def test_call_success(self, mock_call: MagicMock) -> None:
        """Test call success"""
        # Mocking the response from the API
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.output = {
            "choices": [
                {"message": {"content": [{"text": "This is the result."}]}},
            ],
        }
        mock_response.usage = {
            "input_tokens": 23,
            "output_tokens": 5,
            "image_tokens": 17,
        }
        mock_call.return_value = mock_response

        messages = [
            {
                "role": "user",
                "content": [
                    {"text": "What does this picture depict?"},
                    {"image": "http://example.com/image.jpg"},
                ],
            },
        ]
        # Calling the wrapper and validating the response
        response = self.wrapper(messages=messages)
        self.assertIsInstance(response, ModelResponse)
        self.assertEqual(response.text, "This is the result.")
        self.assertEqual(response.raw, mock_response)

        # Verify call to dashscope.MultiModalConversation.call
        mock_call.assert_called_once_with(
            model=self.wrapper.model_name,
            messages=messages,
        )

    @patch(
        "agentscope.models.dashscope_model."
        "dashscope.MultiModalConversation.call",
    )
    def test_call_failure(self, mock_call: MagicMock) -> None:
        """Test call failure"""
        # Simulating a failed API call
        mock_response = MagicMock()
        mock_response.status_code = 400
        mock_response.request_id = "Test_request_id"
        mock_response.code = "Error Code"
        mock_response.message = "Error Message"
        mock_call.return_value = mock_response

        messages = [
            {
                "role": "user",
                "content": [
                    {"text": "What does this picture depict?"},
                    {"image": "http://example.com/image.jpg"},
                ],
            },
        ]
        # Expecting a RuntimeError to be raised
        with self.assertRaises(RuntimeError) as context:
            self.wrapper(messages=messages)

        # Assert the expected exception message
        self.assertIn("Error Code", str(context.exception))
        self.assertIn("Error Message", str(context.exception))
        self.assertIn("Test_request_id", str(context.exception))
        self.assertIn("400", str(context.exception))

        # Verify call to dashscope.MultiModalConversation.call
        mock_call.assert_called_once_with(
            model=self.wrapper.model_name,
            messages=messages,
        )


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