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

import agentscope
from agentscope.models import load_model_by_config_name


class TestZhipuAIChatWrapper(unittest.TestCase):
    """Test ZhipuAI Chat Wrapper"""

    def setUp(self) -> None:
        self.api_key = "test_api_key.secret_key"
        self.messages = [
            {"role": "user", "content": "Hello, ZhipuAI!"},
            {"role": "assistant", "content": "How can I assist you?"},
        ]

    @patch("agentscope.models.zhipu_model.zhipuai")
    def test_chat(self, mock_zhipuai: MagicMock) -> None:
        """
        Test chat"""
        mock_response = MagicMock()
        mock_response.model_dump.return_value = {
            "choices": [
                {"message": {"content": "Hello, this is a mocked response!"}},
            ],
            "usage": {
                "prompt_tokens": 100,
                "completion_tokens": 5,
                "total_tokens": 105,
            },
        }
        mock_response.choices[
            0
        ].message.content = "Hello, this is a mocked response!"
        mock_zhipuai_client = MagicMock()
        mock_zhipuai.ZhipuAI.return_value = mock_zhipuai_client

        mock_zhipuai_client.chat.completions.create.return_value = (
            mock_response
        )

        agentscope.init(
            model_configs={
                "config_name": "test_config",
                "model_type": "zhipuai_chat",
                "model_name": "glm-4",
                "api_key": self.api_key,
            },
        )

        model = load_model_by_config_name("test_config")

        response = model(messages=self.messages)

        self.assertEqual(response.text, "Hello, this is a mocked response!")

        mock_zhipuai_client.chat.completions.create.assert_called_once()


class TestZhipuAIEmbeddingWrapper(unittest.TestCase):
    """Test ZhipuAI Embedding Wrapper"""

    def setUp(self) -> None:
        self.api_key = "test_api_key"
        self.model_name = "embedding-2"
        self.text_to_embed = "This is a test sentence for embedding."

    @patch("agentscope.models.zhipu_model.zhipuai")
    def test_embedding(self, mock_zhipuai: MagicMock) -> None:
        """Test embedding API"""
        mock_embedding_response = MagicMock()
        mock_embedding_response.model_dump.return_value = {
            "data": [
                {"embedding": [0.1, 0.2, 0.3]},
            ],
            "usage": {
                "prompt_tokens": 10,
                "completion_tokens": 2,
                "total_tokens": 12,
            },
        }

        mock_zhipuai_client = MagicMock()
        mock_zhipuai.ZhipuAI.return_value = mock_zhipuai_client
        mock_zhipuai_client.embeddings.create.return_value = (
            mock_embedding_response
        )

        agentscope.init(
            model_configs={
                "config_name": "test_embedding",
                "model_type": "zhipuai_embedding",
                "model_name": self.model_name,
                "api_key": self.api_key,
            },
        )

        model = load_model_by_config_name("test_embedding")

        response = model(self.text_to_embed)

        expected_embedding = [[0.1, 0.2, 0.3]]
        self.assertEqual(response.embedding, expected_embedding)

        mock_zhipuai_client.embeddings.create.assert_called_once_with(
            input=self.text_to_embed,
            model=self.model_name,
            **{},
        )


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