"""Unit tests for MCP utils module."""

import json
from collections.abc import Generator
from unittest.mock import MagicMock, Mock, patch

import httpx
import httpx_sse
import pytest

from core.mcp.utils import (
    STATUS_FORCELIST,
    create_mcp_error_response,
    create_ssrf_proxy_mcp_http_client,
    ssrf_proxy_sse_connect,
)


class TestConstants:
    """Test module constants."""

    def test_status_forcelist(self):
        """Test STATUS_FORCELIST contains expected HTTP status codes."""
        assert STATUS_FORCELIST == [429, 500, 502, 503, 504]
        assert 429 in STATUS_FORCELIST  # Too Many Requests
        assert 500 in STATUS_FORCELIST  # Internal Server Error
        assert 502 in STATUS_FORCELIST  # Bad Gateway
        assert 503 in STATUS_FORCELIST  # Service Unavailable
        assert 504 in STATUS_FORCELIST  # Gateway Timeout


class TestCreateSSRFProxyMCPHTTPClient:
    """Test create_ssrf_proxy_mcp_http_client function."""

    @patch("core.mcp.utils.dify_config", autospec=True)
    def test_create_client_with_all_url_proxy(self, mock_config):
        """Test client creation with SSRF_PROXY_ALL_URL configured."""
        mock_config.SSRF_PROXY_ALL_URL = "http://proxy.example.com:8080"
        mock_config.HTTP_REQUEST_NODE_SSL_VERIFY = True

        client = create_ssrf_proxy_mcp_http_client(
            headers={"Authorization": "Bearer token"}, timeout=httpx.Timeout(30.0)
        )

        assert isinstance(client, httpx.Client)
        assert client.headers["Authorization"] == "Bearer token"
        assert client.timeout.connect == 30.0
        assert client.follow_redirects is True

        # Clean up
        client.close()

    @patch("core.mcp.utils.dify_config", autospec=True)
    def test_create_client_with_http_https_proxies(self, mock_config):
        """Test client creation with separate HTTP/HTTPS proxies."""
        mock_config.SSRF_PROXY_ALL_URL = None
        mock_config.SSRF_PROXY_HTTP_URL = "http://http-proxy.example.com:8080"
        mock_config.SSRF_PROXY_HTTPS_URL = "http://https-proxy.example.com:8443"
        mock_config.HTTP_REQUEST_NODE_SSL_VERIFY = False

        client = create_ssrf_proxy_mcp_http_client()

        assert isinstance(client, httpx.Client)
        assert client.follow_redirects is True

        # Clean up
        client.close()

    @patch("core.mcp.utils.dify_config", autospec=True)
    def test_create_client_without_proxy(self, mock_config):
        """Test client creation without proxy configuration."""
        mock_config.SSRF_PROXY_ALL_URL = None
        mock_config.SSRF_PROXY_HTTP_URL = None
        mock_config.SSRF_PROXY_HTTPS_URL = None
        mock_config.HTTP_REQUEST_NODE_SSL_VERIFY = True

        headers = {"X-Custom-Header": "value"}
        timeout = httpx.Timeout(timeout=30.0, connect=5.0, read=10.0, write=30.0)

        client = create_ssrf_proxy_mcp_http_client(headers=headers, timeout=timeout)

        assert isinstance(client, httpx.Client)
        assert client.headers["X-Custom-Header"] == "value"
        assert client.timeout.connect == 5.0
        assert client.timeout.read == 10.0
        assert client.follow_redirects is True

        # Clean up
        client.close()

    @patch("core.mcp.utils.dify_config", autospec=True)
    def test_create_client_default_params(self, mock_config):
        """Test client creation with default parameters."""
        mock_config.SSRF_PROXY_ALL_URL = None
        mock_config.SSRF_PROXY_HTTP_URL = None
        mock_config.SSRF_PROXY_HTTPS_URL = None
        mock_config.HTTP_REQUEST_NODE_SSL_VERIFY = True

        client = create_ssrf_proxy_mcp_http_client()

        assert isinstance(client, httpx.Client)
        # httpx.Client adds default headers, so we just check it's a Headers object
        assert isinstance(client.headers, httpx.Headers)
        # When no timeout is provided, httpx uses its default timeout
        assert client.timeout is not None

        # Clean up
        client.close()


class TestSSRFProxySSEConnect:
    """Test ssrf_proxy_sse_connect function."""

    @patch("core.mcp.utils.connect_sse", autospec=True)
    @patch("core.mcp.utils.create_ssrf_proxy_mcp_http_client", autospec=True)
    def test_sse_connect_with_provided_client(self, mock_create_client, mock_connect_sse):
        """Test SSE connection with pre-configured client."""
        # Setup mocks
        mock_client = Mock(spec=httpx.Client)
        mock_event_source = Mock(spec=httpx_sse.EventSource)
        mock_context = MagicMock()
        mock_context.__enter__.return_value = mock_event_source
        mock_connect_sse.return_value = mock_context

        # Call with provided client
        result = ssrf_proxy_sse_connect(
            "http://example.com/sse", client=mock_client, method="POST", headers={"Authorization": "Bearer token"}
        )

        # Verify client creation was not called
        mock_create_client.assert_not_called()

        # Verify connect_sse was called correctly
        mock_connect_sse.assert_called_once_with(
            mock_client, "POST", "http://example.com/sse", headers={"Authorization": "Bearer token"}
        )

        # Verify result
        assert result == mock_context

    @patch("core.mcp.utils.connect_sse", autospec=True)
    @patch("core.mcp.utils.create_ssrf_proxy_mcp_http_client", autospec=True)
    @patch("core.mcp.utils.dify_config", autospec=True)
    def test_sse_connect_without_client(self, mock_config, mock_create_client, mock_connect_sse):
        """Test SSE connection without pre-configured client."""
        # Setup config
        mock_config.SSRF_DEFAULT_TIME_OUT = 30.0
        mock_config.SSRF_DEFAULT_CONNECT_TIME_OUT = 10.0
        mock_config.SSRF_DEFAULT_READ_TIME_OUT = 60.0
        mock_config.SSRF_DEFAULT_WRITE_TIME_OUT = 30.0

        # Setup mocks
        mock_client = Mock(spec=httpx.Client)
        mock_create_client.return_value = mock_client

        mock_event_source = Mock(spec=httpx_sse.EventSource)
        mock_context = MagicMock()
        mock_context.__enter__.return_value = mock_event_source
        mock_connect_sse.return_value = mock_context

        # Call without client
        result = ssrf_proxy_sse_connect("http://example.com/sse", headers={"X-Custom": "value"})

        # Verify client was created
        mock_create_client.assert_called_once()
        call_args = mock_create_client.call_args
        assert call_args[1]["headers"] == {"X-Custom": "value"}

        timeout = call_args[1]["timeout"]
        # httpx.Timeout object has these attributes
        assert isinstance(timeout, httpx.Timeout)
        assert timeout.connect == 10.0
        assert timeout.read == 60.0
        assert timeout.write == 30.0

        # Verify connect_sse was called
        mock_connect_sse.assert_called_once_with(
            mock_client,
            "GET",  # Default method
            "http://example.com/sse",
        )

        # Verify result
        assert result == mock_context

    @patch("core.mcp.utils.connect_sse", autospec=True)
    @patch("core.mcp.utils.create_ssrf_proxy_mcp_http_client", autospec=True)
    def test_sse_connect_with_custom_timeout(self, mock_create_client, mock_connect_sse):
        """Test SSE connection with custom timeout."""
        # Setup mocks
        mock_client = Mock(spec=httpx.Client)
        mock_create_client.return_value = mock_client

        mock_event_source = Mock(spec=httpx_sse.EventSource)
        mock_context = MagicMock()
        mock_context.__enter__.return_value = mock_event_source
        mock_connect_sse.return_value = mock_context

        custom_timeout = httpx.Timeout(timeout=60.0, read=120.0)

        # Call with custom timeout
        result = ssrf_proxy_sse_connect("http://example.com/sse", timeout=custom_timeout)

        # Verify client was created with custom timeout
        mock_create_client.assert_called_once()
        call_args = mock_create_client.call_args
        assert call_args[1]["timeout"] == custom_timeout

        # Verify result
        assert result == mock_context

    @patch("core.mcp.utils.connect_sse", autospec=True)
    @patch("core.mcp.utils.create_ssrf_proxy_mcp_http_client", autospec=True)
    def test_sse_connect_error_cleanup(self, mock_create_client, mock_connect_sse):
        """Test SSE connection cleans up client on error."""
        # Setup mocks
        mock_client = Mock(spec=httpx.Client)
        mock_create_client.return_value = mock_client

        # Make connect_sse raise an exception
        mock_connect_sse.side_effect = httpx.ConnectError("Connection failed")

        # Call should raise the exception
        with pytest.raises(httpx.ConnectError):
            ssrf_proxy_sse_connect("http://example.com/sse")

        # Verify client was cleaned up
        mock_client.close.assert_called_once()

    @patch("core.mcp.utils.connect_sse", autospec=True)
    def test_sse_connect_error_no_cleanup_with_provided_client(self, mock_connect_sse):
        """Test SSE connection doesn't clean up provided client on error."""
        # Setup mocks
        mock_client = Mock(spec=httpx.Client)

        # Make connect_sse raise an exception
        mock_connect_sse.side_effect = httpx.ConnectError("Connection failed")

        # Call should raise the exception
        with pytest.raises(httpx.ConnectError):
            ssrf_proxy_sse_connect("http://example.com/sse", client=mock_client)

        # Verify client was NOT cleaned up (because it was provided)
        mock_client.close.assert_not_called()


class TestCreateMCPErrorResponse:
    """Test create_mcp_error_response function."""

    def test_create_error_response_basic(self):
        """Test creating basic error response."""
        generator = create_mcp_error_response(request_id="req-123", code=-32600, message="Invalid Request")

        # Generator should yield bytes
        assert isinstance(generator, Generator)

        # Get the response
        response_bytes = next(generator)
        assert isinstance(response_bytes, bytes)

        # Parse the response
        response_str = response_bytes.decode("utf-8")
        response_json = json.loads(response_str)

        assert response_json["jsonrpc"] == "2.0"
        assert response_json["id"] == "req-123"
        assert response_json["error"]["code"] == -32600
        assert response_json["error"]["message"] == "Invalid Request"
        assert response_json["error"]["data"] is None

        # Generator should be exhausted
        with pytest.raises(StopIteration):
            next(generator)

    def test_create_error_response_with_data(self):
        """Test creating error response with additional data."""
        error_data = {"field": "username", "reason": "required"}

        generator = create_mcp_error_response(
            request_id=456,  # Numeric ID
            code=-32602,
            message="Invalid params",
            data=error_data,
        )

        response_bytes = next(generator)
        response_json = json.loads(response_bytes.decode("utf-8"))

        assert response_json["id"] == 456
        assert response_json["error"]["code"] == -32602
        assert response_json["error"]["message"] == "Invalid params"
        assert response_json["error"]["data"] == error_data

    def test_create_error_response_without_request_id(self):
        """Test creating error response without request ID."""
        generator = create_mcp_error_response(request_id=None, code=-32700, message="Parse error")

        response_bytes = next(generator)
        response_json = json.loads(response_bytes.decode("utf-8"))

        # Should default to ID 1
        assert response_json["id"] == 1
        assert response_json["error"]["code"] == -32700
        assert response_json["error"]["message"] == "Parse error"

    def test_create_error_response_with_complex_data(self):
        """Test creating error response with complex error data."""
        complex_data = {
            "errors": [{"field": "name", "message": "Too short"}, {"field": "email", "message": "Invalid format"}],
            "timestamp": "2024-01-01T00:00:00Z",
        }

        generator = create_mcp_error_response(
            request_id="complex-req", code=-32602, message="Validation failed", data=complex_data
        )

        response_bytes = next(generator)
        response_json = json.loads(response_bytes.decode("utf-8"))

        assert response_json["error"]["data"] == complex_data
        assert len(response_json["error"]["data"]["errors"]) == 2

    def test_create_error_response_encoding(self):
        """Test error response with non-ASCII characters."""
        generator = create_mcp_error_response(
            request_id="unicode-req",
            code=-32603,
            message="内部错误",  # Chinese characters
            data={"details": "エラー詳細"},  # Japanese characters
        )

        response_bytes = next(generator)

        # Should be valid UTF-8
        response_str = response_bytes.decode("utf-8")
        response_json = json.loads(response_str)

        assert response_json["error"]["message"] == "内部错误"
        assert response_json["error"]["data"]["details"] == "エラー詳細"

    def test_create_error_response_yields_once(self):
        """Test that error response generator yields exactly once."""
        generator = create_mcp_error_response(request_id="test", code=-32600, message="Test")

        # First yield should work
        first_yield = next(generator)
        assert isinstance(first_yield, bytes)

        # Second yield should raise StopIteration
        with pytest.raises(StopIteration):
            next(generator)

        # Subsequent calls should also raise
        with pytest.raises(StopIteration):
            next(generator)
