# tests/test_minicpm_messages.py
from __future__ import annotations

import json
import os
import unittest

os.environ.setdefault("BORDERLESS_INFERENCE_MODE", "hub")
os.environ.setdefault("BORDERLESS_PRELOAD_MODEL", "0")

from langchain_core.messages import AIMessage, ToolMessage

from ui.agent.minicpm.messages import (
    append_tool_instructions,
    normalize_messages,
)


class NormalizeMessagesTests(unittest.TestCase):
    def test_dict_messages(self) -> None:
        messages = [
            {"role": "system", "content": "You are helpful."},
            {"role": "user", "content": "Hello"},
        ]
        normalized = normalize_messages(messages)
        self.assertEqual(normalized[0]["role"], "system")
        self.assertEqual(normalized[1]["content"], "Hello")

    def test_tool_message_becomes_user_context(self) -> None:
        messages = [
            ToolMessage(content='{"results": []}', tool_call_id="call-1"),
        ]
        normalized = normalize_messages(messages)
        self.assertEqual(normalized[0]["role"], "user")
        self.assertIn("Tool result", normalized[0]["content"])

    def test_ai_message(self) -> None:
        messages = [AIMessage(content="Done.")]
        normalized = normalize_messages(messages)
        self.assertEqual(normalized[0]["role"], "assistant")
        self.assertEqual(normalized[0]["content"], "Done.")


class ToolInstructionTests(unittest.TestCase):
    def test_appends_tool_block_to_system(self) -> None:
        tools = [
            {
                "type": "function",
                "function": {
                    "name": "search_immigration_info",
                    "description": "Search immigration sources",
                    "parameters": {
                        "type": "object",
                        "properties": {"query": {"type": "string"}},
                    },
                },
            }
        ]
        messages = [
            {"role": "system", "content": "Base prompt"},
            {"role": "user", "content": "Research Canada"},
        ]
        updated = append_tool_instructions(messages, tools)
        self.assertIn("search_immigration_info", updated[0]["content"])
        self.assertIn(json.dumps(tools[0]["function"]["parameters"]), updated[0]["content"])


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