#!/usr/bin/env python3
"""Unit tests for the generic LLM Chat MCP Server (mcp-servers/llm-chat/server.py).

Tests cover:
- JSON-RPC request handling (initialize, ping, tools/list, tools/call)
- call_llm: success, API errors, 504 retry + fallback model logic
- Notification handling (no response)
"""

import os
import sys
import unittest
from unittest.mock import patch, MagicMock, call

sys.path.insert(0, os.path.join(os.path.dirname(__file__)))


class TestHandleRequest(unittest.TestCase):
    """Test JSON-RPC request routing."""

    def test_initialize_response(self):
        """initialize should return protocol version and server info."""
        from tests._llm_chat_helpers import handle_request
        resp = handle_request({"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}})
        self.assertEqual(resp["id"], 1)
        self.assertEqual(resp["result"]["protocolVersion"], "2024-11-05")
        self.assertIn("tools", resp["result"]["capabilities"])
        self.assertEqual(resp["result"]["serverInfo"]["name"], "llm-chat")

    def test_ping_response(self):
        """ping should return empty result."""
        from tests._llm_chat_helpers import handle_request
        resp = handle_request({"jsonrpc": "2.0", "id": 2, "method": "ping", "params": {}})
        self.assertEqual(resp["id"], 2)
        self.assertEqual(resp["result"], {})

    def test_notification_returns_none(self):
        """Requests without id are notifications and should return None."""
        from tests._llm_chat_helpers import handle_request
        resp = handle_request({"jsonrpc": "2.0", "method": "notifications/initialized"})
        self.assertIsNone(resp)

    def test_unknown_method_returns_error(self):
        """Unknown methods should return JSON-RPC error -32601."""
        from tests._llm_chat_helpers import handle_request
        resp = handle_request({"jsonrpc": "2.0", "id": 3, "method": "unknown/method", "params": {}})
        self.assertIn("error", resp)
        self.assertEqual(resp["error"]["code"], -32601)

    def test_unknown_tool_returns_error(self):
        """Unknown tool name inside tools/call should return error."""
        from tests._llm_chat_helpers import handle_request
        resp = handle_request({
            "jsonrpc": "2.0", "id": 4, "method": "tools/call",
            "params": {"name": "nonexistent_tool", "arguments": {}}
        })
        self.assertIn("error", resp)
        self.assertEqual(resp["error"]["code"], -32601)


class TestToolsList(unittest.TestCase):
    """Test tools/list response structure."""

    def test_tools_list_returns_chat_tool(self):
        """tools/list should return a single 'chat' tool."""
        from tests._llm_chat_helpers import handle_request
        resp = handle_request({"jsonrpc": "2.0", "id": 5, "method": "tools/list", "params": {}})
        tools = resp["result"]["tools"]
        self.assertEqual(len(tools), 1)
        self.assertEqual(tools[0]["name"], "chat")

    def test_tools_list_schema_has_required_prompt(self):
        """The chat tool schema must require 'prompt'."""
        from tests._llm_chat_helpers import handle_request
        resp = handle_request({"jsonrpc": "2.0", "id": 6, "method": "tools/list", "params": {}})
        schema = resp["result"]["tools"][0]["inputSchema"]
        self.assertIn("prompt", schema["required"])

    def test_tools_list_schema_has_optional_model_and_system(self):
        """The chat tool schema should expose optional 'model' and 'system' parameters."""
        from tests._llm_chat_helpers import handle_request
        resp = handle_request({"jsonrpc": "2.0", "id": 7, "method": "tools/list", "params": {}})
        props = resp["result"]["tools"][0]["inputSchema"]["properties"]
        self.assertIn("model", props)
        self.assertIn("system", props)


class TestToolCallNoApiKey(unittest.TestCase):
    """Test tool call behavior when LLM_API_KEY is missing."""

    @patch("tests._llm_chat_helpers.LLM_API_KEY", "")
    def test_missing_api_key_returns_error(self):
        """Tool call without API key should return isError result."""
        from tests._llm_chat_helpers import handle_request
        resp = handle_request({
            "jsonrpc": "2.0", "id": 8, "method": "tools/call",
            "params": {"name": "chat", "arguments": {"prompt": "hello"}}
        })
        self.assertTrue(resp["result"]["isError"])
        self.assertIn("LLM_API_KEY", resp["result"]["content"][0]["text"])


class TestCallLlmSuccess(unittest.TestCase):
    """Test call_llm for successful API responses."""

    @patch("tests._llm_chat_helpers.LLM_API_KEY", "test-key")
    @patch("httpx.Client")
    def test_successful_call_returns_content(self, mock_client_cls):
        """A 200 response should return the message content."""
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = {
            "choices": [{"message": {"content": "Hello from LLM!"}}]
        }
        mock_client = MagicMock()
        mock_client.__enter__ = MagicMock(return_value=mock_client)
        mock_client.__exit__ = MagicMock(return_value=False)
        mock_client.post.return_value = mock_response
        mock_client_cls.return_value = mock_client

        from tests._llm_chat_helpers import call_llm
        content, error = call_llm([{"role": "user", "content": "hi"}])
        self.assertEqual(content, "Hello from LLM!")
        self.assertIsNone(error)

    @patch("tests._llm_chat_helpers.LLM_API_KEY", "test-key")
    @patch("httpx.Client")
    def test_custom_model_is_passed(self, mock_client_cls):
        """The requested model name should appear in the API payload."""
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = {
            "choices": [{"message": {"content": "OK"}}]
        }
        mock_client = MagicMock()
        mock_client.__enter__ = MagicMock(return_value=mock_client)
        mock_client.__exit__ = MagicMock(return_value=False)
        mock_client.post.return_value = mock_response
        mock_client_cls.return_value = mock_client

        from tests._llm_chat_helpers import call_llm
        call_llm([{"role": "user", "content": "test"}], model="deepseek-chat")
        payload = mock_client.post.call_args[1]["json"]
        self.assertEqual(payload["model"], "deepseek-chat")

    @patch("tests._llm_chat_helpers.LLM_API_KEY", "test-key")
    @patch("httpx.Client")
    def test_api_error_status_returns_error_message(self, mock_client_cls):
        """Non-200, non-504 status should return an error string."""
        mock_response = MagicMock()
        mock_response.status_code = 401
        mock_response.text = "Unauthorized"
        mock_client = MagicMock()
        mock_client.__enter__ = MagicMock(return_value=mock_client)
        mock_client.__exit__ = MagicMock(return_value=False)
        mock_client.post.return_value = mock_response
        mock_client_cls.return_value = mock_client

        from tests._llm_chat_helpers import call_llm
        content, error = call_llm([{"role": "user", "content": "test"}])
        self.assertIsNone(content)
        self.assertIn("401", error)

    @patch("tests._llm_chat_helpers.LLM_API_KEY", "")
    def test_missing_api_key_returns_error(self):
        """call_llm without API key should return error immediately."""
        from tests._llm_chat_helpers import call_llm
        content, error = call_llm([{"role": "user", "content": "test"}])
        self.assertIsNone(content)
        self.assertIn("LLM_API_KEY", error)


class TestCallLlm504Retry(unittest.TestCase):
    """Test the 504 retry and fallback model logic in call_llm."""

    @patch("tests._llm_chat_helpers.LLM_API_KEY", "test-key")
    @patch("tests._llm_chat_helpers.DEFAULT_MODEL", "gpt-4o")
    @patch("tests._llm_chat_helpers.FALLBACK_MODEL", "gpt-4o-mini")
    @patch("httpx.Client")
    def test_504_twice_then_fallback_succeeds(self, mock_client_cls):
        """Two 504s should trigger fallback model on attempt 3."""
        resp_504 = MagicMock()
        resp_504.status_code = 504

        resp_ok = MagicMock()
        resp_ok.status_code = 200
        resp_ok.json.return_value = {
            "choices": [{"message": {"content": "Fallback reply"}}]
        }

        mock_client = MagicMock()
        mock_client.__enter__ = MagicMock(return_value=mock_client)
        mock_client.__exit__ = MagicMock(return_value=False)
        mock_client.post.side_effect = [resp_504, resp_504, resp_ok]
        mock_client_cls.return_value = mock_client

        from tests._llm_chat_helpers import call_llm
        content, error = call_llm([{"role": "user", "content": "test"}])

        self.assertIsNone(error)
        self.assertIn("Fallback reply", content)
        # Fallback note should be prepended
        self.assertIn("[Note: Used fallback model gpt-4o-mini", content)

    @patch("tests._llm_chat_helpers.LLM_API_KEY", "test-key")
    @patch("tests._llm_chat_helpers.DEFAULT_MODEL", "gpt-4o")
    @patch("tests._llm_chat_helpers.FALLBACK_MODEL", "gpt-4o-mini")
    @patch("httpx.Client")
    def test_504_once_then_retry_succeeds_no_fallback_note(self, mock_client_cls):
        """A single 504 followed by success should use original model (no fallback note)."""
        resp_504 = MagicMock()
        resp_504.status_code = 504

        resp_ok = MagicMock()
        resp_ok.status_code = 200
        resp_ok.json.return_value = {
            "choices": [{"message": {"content": "Retry success"}}]
        }

        mock_client = MagicMock()
        mock_client.__enter__ = MagicMock(return_value=mock_client)
        mock_client.__exit__ = MagicMock(return_value=False)
        mock_client.post.side_effect = [resp_504, resp_ok]
        mock_client_cls.return_value = mock_client

        from tests._llm_chat_helpers import call_llm
        content, error = call_llm([{"role": "user", "content": "test"}])

        self.assertIsNone(error)
        self.assertEqual(content, "Retry success")
        self.assertNotIn("[Note:", content)

    @patch("tests._llm_chat_helpers.LLM_API_KEY", "test-key")
    @patch("httpx.Client")
    def test_three_504s_returns_error(self, mock_client_cls):
        """Three consecutive 504s should return the gateway timeout error."""
        resp_504 = MagicMock()
        resp_504.status_code = 504

        mock_client = MagicMock()
        mock_client.__enter__ = MagicMock(return_value=mock_client)
        mock_client.__exit__ = MagicMock(return_value=False)
        mock_client.post.return_value = resp_504
        mock_client_cls.return_value = mock_client

        from tests._llm_chat_helpers import call_llm
        content, error = call_llm([{"role": "user", "content": "test"}])

        self.assertIsNone(content)
        self.assertIn("504", error)

    @patch("tests._llm_chat_helpers.LLM_API_KEY", "test-key")
    @patch("tests._llm_chat_helpers.DEFAULT_MODEL", "primary-model")
    @patch("tests._llm_chat_helpers.FALLBACK_MODEL", "fallback-model")
    @patch("httpx.Client")
    def test_fallback_uses_different_model_name(self, mock_client_cls):
        """On attempt 3, the payload model should be FALLBACK_MODEL, not DEFAULT_MODEL."""
        resp_504 = MagicMock()
        resp_504.status_code = 504

        resp_ok = MagicMock()
        resp_ok.status_code = 200
        resp_ok.json.return_value = {
            "choices": [{"message": {"content": "OK"}}]
        }

        mock_client = MagicMock()
        mock_client.__enter__ = MagicMock(return_value=mock_client)
        mock_client.__exit__ = MagicMock(return_value=False)
        mock_client.post.side_effect = [resp_504, resp_504, resp_ok]
        mock_client_cls.return_value = mock_client

        from tests._llm_chat_helpers import call_llm
        call_llm([{"role": "user", "content": "test"}])

        # The third call (index 2) should use the fallback model
        third_call_payload = mock_client.post.call_args_list[2][1]["json"]
        self.assertEqual(third_call_payload["model"], "fallback-model")


class TestToolCallFullFlow(unittest.TestCase):
    """Test the complete tools/call path through handle_request."""

    @patch("tests._llm_chat_helpers.LLM_API_KEY", "test-key")
    @patch("httpx.Client")
    def test_tool_call_success(self, mock_client_cls):
        """Successful tool call should return content without isError."""
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = {
            "choices": [{"message": {"content": "Test response"}}]
        }
        mock_client = MagicMock()
        mock_client.__enter__ = MagicMock(return_value=mock_client)
        mock_client.__exit__ = MagicMock(return_value=False)
        mock_client.post.return_value = mock_response
        mock_client_cls.return_value = mock_client

        from tests._llm_chat_helpers import handle_request
        resp = handle_request({
            "jsonrpc": "2.0", "id": 10, "method": "tools/call",
            "params": {"name": "chat", "arguments": {"prompt": "Hello"}}
        })
        self.assertFalse(resp["result"].get("isError", False))
        self.assertEqual(resp["result"]["content"][0]["text"], "Test response")

    @patch("tests._llm_chat_helpers.LLM_API_KEY", "test-key")
    @patch("httpx.Client")
    def test_tool_call_with_system_prompt(self, mock_client_cls):
        """System prompt should be included as first message with role='system'."""
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = {
            "choices": [{"message": {"content": "OK"}}]
        }
        mock_client = MagicMock()
        mock_client.__enter__ = MagicMock(return_value=mock_client)
        mock_client.__exit__ = MagicMock(return_value=False)
        mock_client.post.return_value = mock_response
        mock_client_cls.return_value = mock_client

        from tests._llm_chat_helpers import handle_request
        handle_request({
            "jsonrpc": "2.0", "id": 11, "method": "tools/call",
            "params": {
                "name": "chat",
                "arguments": {
                    "prompt": "Review this",
                    "system": "You are a strict reviewer"
                }
            }
        })

        payload = mock_client.post.call_args[1]["json"]
        self.assertEqual(len(payload["messages"]), 2)
        self.assertEqual(payload["messages"][0]["role"], "system")
        self.assertEqual(payload["messages"][0]["content"], "You are a strict reviewer")
        self.assertEqual(payload["messages"][1]["role"], "user")

    @patch("tests._llm_chat_helpers.LLM_API_KEY", "test-key")
    @patch("httpx.Client")
    def test_tool_call_api_error_returns_is_error(self, mock_client_cls):
        """An API error should be surfaced as isError=True in the result."""
        mock_response = MagicMock()
        mock_response.status_code = 500
        mock_response.text = "Internal Server Error"
        mock_client = MagicMock()
        mock_client.__enter__ = MagicMock(return_value=mock_client)
        mock_client.__exit__ = MagicMock(return_value=False)
        mock_client.post.return_value = mock_response
        mock_client_cls.return_value = mock_client

        from tests._llm_chat_helpers import handle_request
        resp = handle_request({
            "jsonrpc": "2.0", "id": 12, "method": "tools/call",
            "params": {"name": "chat", "arguments": {"prompt": "test"}}
        })
        self.assertTrue(resp["result"]["isError"])
        self.assertIn("500", resp["result"]["content"][0]["text"])


class TestDefaultConfig(unittest.TestCase):
    """Test default configuration values."""

    def test_default_base_url(self):
        """Default base URL should be OpenAI API."""
        self.assertEqual(
            os.environ.get("LLM_BASE_URL", "https://api.openai.com/v1"),
            "https://api.openai.com/v1"
        )

    def test_default_model(self):
        """Default model should be gpt-4o."""
        self.assertEqual(
            os.environ.get("LLM_MODEL", "gpt-4o"),
            "gpt-4o"
        )

    def test_default_server_name(self):
        """Default server name should be llm-chat."""
        self.assertEqual(
            os.environ.get("LLM_SERVER_NAME", "llm-chat"),
            "llm-chat"
        )


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