import contextlib
import json
import queue
import threading
import time
from typing import Any
from unittest.mock import Mock, patch

import httpx
import pytest

from core.mcp import types
from core.mcp.client.sse_client import sse_client
from core.mcp.error import MCPAuthError, MCPConnectionError

SERVER_NAME = "test_server_for_SSE"


def test_sse_message_id_coercion():
    """Test that string message IDs that look like integers are parsed as integers.

    See <https://github.com/modelcontextprotocol/python-sdk/pull/851> for more details.
    """
    json_message = '{"jsonrpc": "2.0", "id": "123", "method": "ping", "params": null}'
    msg = types.JSONRPCMessage.model_validate_json(json_message)
    expected = types.JSONRPCMessage(root=types.JSONRPCRequest(method="ping", jsonrpc="2.0", id=123))

    # Check if both are JSONRPCRequest instances
    assert isinstance(msg.root, types.JSONRPCRequest)
    assert isinstance(expected.root, types.JSONRPCRequest)

    assert msg.root.id == expected.root.id
    assert msg.root.method == expected.root.method
    assert msg.root.jsonrpc == expected.root.jsonrpc


class MockSSEClient:
    """Mock SSE client for testing."""

    def __init__(self, url: str, headers: dict[str, Any] | None = None):
        self.url = url
        self.headers = headers or {}
        self.connected = False
        self.read_queue: queue.Queue = queue.Queue()
        self.write_queue: queue.Queue = queue.Queue()

    def connect(self):
        """Simulate connection establishment."""
        self.connected = True

        # Send endpoint event
        endpoint_data = "/messages/?session_id=test-session-123"
        self.read_queue.put(("endpoint", endpoint_data))

        return self.read_queue, self.write_queue

    def send_initialize_response(self):
        """Send a mock initialize response."""
        response = {
            "jsonrpc": "2.0",
            "id": 1,
            "result": {
                "protocolVersion": types.LATEST_PROTOCOL_VERSION,
                "capabilities": {
                    "logging": None,
                    "resources": None,
                    "tools": None,
                    "experimental": None,
                    "prompts": None,
                },
                "serverInfo": {"name": SERVER_NAME, "version": "0.1.0"},
                "instructions": "Test server instructions.",
            },
        }
        self.read_queue.put(("message", json.dumps(response)))


def test_sse_client_message_id_handling():
    """Test SSE client properly handles message ID coercion."""
    mock_client = MockSSEClient("http://test.example/sse")
    read_queue, write_queue = mock_client.connect()

    # Send a message with string ID that should be coerced to int
    message_data = {
        "jsonrpc": "2.0",
        "id": "456",  # String ID
        "result": {"test": "data"},
    }
    read_queue.put(("message", json.dumps(message_data)))
    read_queue.get(timeout=1.0)
    # Get the message from queue
    event_type, data = read_queue.get(timeout=1.0)
    assert event_type == "message"

    # Parse the message
    parsed_message = types.JSONRPCMessage.model_validate_json(data)
    # Check that it's a JSONRPCResponse and verify the ID
    assert isinstance(parsed_message.root, types.JSONRPCResponse)
    assert parsed_message.root.id == 456  # Should be converted to int


def test_sse_client_connection_validation():
    """Test SSE client validates endpoint URLs properly."""
    test_url = "http://test.example/sse"

    with patch("core.mcp.client.sse_client.create_ssrf_proxy_mcp_http_client") as mock_client_factory:
        with patch("core.mcp.client.sse_client.ssrf_proxy_sse_connect") as mock_sse_connect:
            # Mock the HTTP client
            mock_client = Mock()
            mock_client_factory.return_value.__enter__.return_value = mock_client

            # Mock the SSE connection
            mock_event_source = Mock()
            mock_event_source.response.raise_for_status.return_value = None
            mock_sse_connect.return_value.__enter__.return_value = mock_event_source

            # Mock SSE events
            class MockSSEEvent:
                def __init__(self, event_type: str, data: str):
                    self.event = event_type
                    self.data = data

            # Simulate endpoint event
            endpoint_event = MockSSEEvent("endpoint", "/messages/?session_id=test-123")
            mock_event_source.iter_sse.return_value = [endpoint_event]

            # Test connection
            with contextlib.suppress(Exception):
                with sse_client(test_url) as (read_queue, write_queue):
                    assert read_queue is not None
                    assert write_queue is not None


def test_sse_client_error_handling():
    """Test SSE client properly handles various error conditions."""
    test_url = "http://test.example/sse"

    # Test 401 error handling
    with patch("core.mcp.client.sse_client.create_ssrf_proxy_mcp_http_client") as mock_client_factory:
        with patch("core.mcp.client.sse_client.ssrf_proxy_sse_connect") as mock_sse_connect:
            # Mock 401 HTTP error
            mock_response = Mock(status_code=401)
            mock_response.headers = {"WWW-Authenticate": 'Bearer realm="example"'}
            mock_error = httpx.HTTPStatusError("Unauthorized", request=Mock(), response=mock_response)
            mock_sse_connect.side_effect = mock_error

            with pytest.raises(MCPAuthError):
                with sse_client(test_url):
                    pass

    # Test other HTTP errors
    with patch("core.mcp.client.sse_client.create_ssrf_proxy_mcp_http_client") as mock_client_factory:
        with patch("core.mcp.client.sse_client.ssrf_proxy_sse_connect") as mock_sse_connect:
            # Mock other HTTP error
            mock_response = Mock(status_code=500)
            mock_response.headers = {}
            mock_error = httpx.HTTPStatusError("Server Error", request=Mock(), response=mock_response)
            mock_sse_connect.side_effect = mock_error

            with pytest.raises(MCPConnectionError):
                with sse_client(test_url):
                    pass


def test_sse_client_timeout_configuration():
    """Test SSE client timeout configuration."""
    test_url = "http://test.example/sse"
    custom_timeout = 10.0
    custom_sse_timeout = 300.0
    custom_headers = {"Authorization": "Bearer test-token"}

    with patch("core.mcp.client.sse_client.create_ssrf_proxy_mcp_http_client") as mock_client_factory:
        with patch("core.mcp.client.sse_client.ssrf_proxy_sse_connect") as mock_sse_connect:
            # Mock successful connection
            mock_client = Mock()
            mock_client_factory.return_value.__enter__.return_value = mock_client

            mock_event_source = Mock()
            mock_event_source.response.raise_for_status.return_value = None
            mock_event_source.iter_sse.return_value = []
            mock_sse_connect.return_value.__enter__.return_value = mock_event_source

            with contextlib.suppress(Exception):
                with sse_client(
                    test_url, headers=custom_headers, timeout=custom_timeout, sse_read_timeout=custom_sse_timeout
                ) as (read_queue, write_queue):
                    # Verify the configuration was passed correctly
                    mock_client_factory.assert_called_with(headers=custom_headers)

                    # Check that timeout was configured
                    call_args = mock_sse_connect.call_args
                    assert call_args is not None
                    timeout_arg = call_args[1]["timeout"]
                    assert timeout_arg.read == custom_sse_timeout


def test_sse_transport_endpoint_validation():
    """Test SSE transport validates endpoint URLs correctly."""
    from core.mcp.client.sse_client import SSETransport

    transport = SSETransport("http://example.com/sse")

    # Valid endpoint (same origin)
    valid_endpoint = "http://example.com/messages/session123"
    assert transport._validate_endpoint_url(valid_endpoint) == True

    # Invalid endpoint (different origin)
    invalid_endpoint = "http://malicious.com/messages/session123"
    assert transport._validate_endpoint_url(invalid_endpoint) == False

    # Invalid endpoint (different scheme)
    invalid_scheme = "https://example.com/messages/session123"
    assert transport._validate_endpoint_url(invalid_scheme) == False


def test_sse_transport_message_parsing():
    """Test SSE transport properly parses different message types."""
    from core.mcp.client.sse_client import SSETransport

    transport = SSETransport("http://example.com/sse")
    read_queue: queue.Queue = queue.Queue()

    # Test valid JSON-RPC message
    valid_message = '{"jsonrpc": "2.0", "id": 1, "method": "ping"}'
    transport._handle_message_event(valid_message, read_queue)

    # Should have a SessionMessage in the queue
    message = read_queue.get(timeout=1.0)
    assert message is not None
    assert hasattr(message, "message")

    # Test invalid JSON
    invalid_json = '{"invalid": json}'
    transport._handle_message_event(invalid_json, read_queue)

    # Should have an exception in the queue
    error = read_queue.get(timeout=1.0)
    assert isinstance(error, Exception)


def test_sse_client_queue_cleanup():
    """Test that SSE client properly cleans up queues on exit."""
    test_url = "http://test.example/sse"

    read_queue = None
    write_queue = None

    with patch("core.mcp.client.sse_client.create_ssrf_proxy_mcp_http_client") as mock_client_factory:
        with patch("core.mcp.client.sse_client.ssrf_proxy_sse_connect") as mock_sse_connect:
            # Mock connection that raises an exception
            mock_sse_connect.side_effect = Exception("Connection failed")

            with contextlib.suppress(Exception):
                with sse_client(test_url) as (rq, wq):
                    read_queue = rq
                    write_queue = wq

            # Queues should be cleaned up even on exception
            # Note: In real implementation, cleanup should put None to signal shutdown


def test_sse_client_headers_propagation():
    """Test that custom headers are properly propagated in SSE client."""
    test_url = "http://test.example/sse"
    custom_headers = {
        "Authorization": "Bearer test-token",
        "X-Custom-Header": "test-value",
        "User-Agent": "test-client/1.0",
    }

    with patch("core.mcp.client.sse_client.create_ssrf_proxy_mcp_http_client") as mock_client_factory:
        with patch("core.mcp.client.sse_client.ssrf_proxy_sse_connect") as mock_sse_connect:
            # Mock the client factory to capture headers
            mock_client = Mock()
            mock_client_factory.return_value.__enter__.return_value = mock_client

            # Mock the SSE connection
            mock_event_source = Mock()
            mock_event_source.response.raise_for_status.return_value = None
            mock_event_source.iter_sse.return_value = []
            mock_sse_connect.return_value.__enter__.return_value = mock_event_source

            with contextlib.suppress(Exception):
                with sse_client(test_url, headers=custom_headers):
                    pass

            # Verify headers were passed to client factory
            mock_client_factory.assert_called_with(headers=custom_headers)


def test_sse_client_concurrent_access():
    """Test SSE client behavior with concurrent queue access."""
    test_read_queue: queue.Queue = queue.Queue()

    # Simulate concurrent producers and consumers
    def producer():
        for i in range(10):
            test_read_queue.put(f"message_{i}")
            time.sleep(0.01)  # Small delay to simulate real conditions

    def consumer():
        received = []
        for _ in range(10):
            try:
                msg = test_read_queue.get(timeout=2.0)
                received.append(msg)
            except queue.Empty:
                break
        return received

    # Start producer in separate thread
    producer_thread = threading.Thread(target=producer, daemon=True)
    producer_thread.start()

    # Consume messages
    received_messages = consumer()

    # Wait for producer to finish
    producer_thread.join(timeout=5.0)

    # Verify all messages were received
    assert len(received_messages) == 10
    for i in range(10):
        assert f"message_{i}" in received_messages
