"""Tests for header propagation to MCP tools."""

from google.adk.tools.mcp_tool import SseConnectionParams, StreamableHTTPConnectionParams

from kagent.adk.types import (
    HttpMcpServerConfig,
    SseMcpServerConfig,
    create_header_provider,
)


class MockReadonlyContext:
    """Mock ReadonlyContext for testing."""

    def __init__(self, state: dict):
        self._state = state

    @property
    def state(self):
        return self._state


class TestCreateHeaderProvider:
    """Tests for create_header_provider function."""

    def test_returns_none_when_no_config(self):
        """Test that None is returned when no headers are configured."""
        result = create_header_provider()
        assert result is None

    def test_returns_none_when_empty_allowed_headers(self):
        """Test that None is returned when allowed_headers is empty list."""
        result = create_header_provider(allowed_headers=[])
        assert result is None

    def test_propagates_allowed_headers(self):
        """Test that allowed headers are propagated from context state."""
        provider = create_header_provider(allowed_headers=["x-user-email", "x-tenant-id"])
        assert provider is not None

        context = MockReadonlyContext(
            state={
                "headers": {
                    "x-user-email": "test@example.com",
                    "x-tenant-id": "tenant-123",
                    "x-secret-header": "should-not-propagate",
                }
            }
        )

        headers = provider(context)
        assert headers == {
            "x-user-email": "test@example.com",
            "x-tenant-id": "tenant-123",
        }
        assert "x-secret-header" not in headers

    def test_case_insensitive_header_matching(self):
        """Test that header matching is case-insensitive."""
        provider = create_header_provider(allowed_headers=["X-User-Email"])
        assert provider is not None

        context = MockReadonlyContext(
            state={
                "headers": {
                    "x-user-email": "test@example.com",
                }
            }
        )

        headers = provider(context)
        assert headers == {"x-user-email": "test@example.com"}

    def test_combines_sts_and_allowed_headers(self):
        """Test that STS headers and allowed headers are combined."""

        def mock_sts_provider(ctx):
            return {"Authorization": "Bearer token123"}

        provider = create_header_provider(
            allowed_headers=["x-user-email"],
            sts_header_provider=mock_sts_provider,
        )
        assert provider is not None

        context = MockReadonlyContext(
            state={
                "headers": {
                    "x-user-email": "test@example.com",
                }
            }
        )

        headers = provider(context)
        assert headers == {
            "Authorization": "Bearer token123",
            "x-user-email": "test@example.com",
        }

    def test_sts_headers_take_precedence_over_allowed_headers(self):
        """Test that STS headers take precedence when there's a conflict.

        This is a security measure to prevent allowed headers from overwriting
        authentication tokens generated by STS. The test uses different cases
        for header names to verify case-insensitive handling.
        """

        def mock_sts_provider(ctx):
            # STS returns "Authorization" with capital A
            return {"Authorization": "Bearer sts-token"}

        provider = create_header_provider(
            allowed_headers=["authorization", "x-user-email"],
            sts_header_provider=mock_sts_provider,
        )
        assert provider is not None

        context = MockReadonlyContext(
            state={
                "headers": {
                    # Request uses lowercase "authorization"
                    "authorization": "Bearer malicious-token",
                    "x-user-email": "test@example.com",
                }
            }
        )

        headers = provider(context)
        # STS token should take precedence, and lowercase "authorization" should be removed
        assert "authorization" not in headers  # lowercase version should be removed
        assert headers["Authorization"] == "Bearer sts-token"  # STS version should be present
        assert headers["x-user-email"] == "test@example.com"
        assert len(headers) == 2  # Only two headers should remain

    def test_sts_only_when_no_allowed_headers(self):
        """Test that only STS headers are returned when no allowed headers."""

        def mock_sts_provider(ctx):
            return {"Authorization": "Bearer token123"}

        provider = create_header_provider(sts_header_provider=mock_sts_provider)
        assert provider is not None

        context = MockReadonlyContext(state={"headers": {}})
        headers = provider(context)
        assert headers == {"Authorization": "Bearer token123"}

    def test_handles_missing_headers_in_state(self):
        """Test that missing headers in state doesn't cause errors."""
        provider = create_header_provider(allowed_headers=["x-user-email"])
        assert provider is not None

        context = MockReadonlyContext(state={})  # No headers key
        headers = provider(context)
        assert headers == {}

    def test_handles_none_context(self):
        """Test that None context doesn't cause errors."""
        provider = create_header_provider(allowed_headers=["x-user-email"])
        assert provider is not None

        headers = provider(None)
        assert headers == {}


class TestMcpServerConfigAllowedHeaders:
    """Tests for allowed_headers field in MCP server configs."""

    def test_http_mcp_config_has_allowed_headers(self):
        """Test that HttpMcpServerConfig has allowed_headers field."""
        config = HttpMcpServerConfig(
            params=StreamableHTTPConnectionParams(url="http://localhost:8080"),
            allowed_headers=["x-user-email", "x-tenant-id"],
        )
        assert config.allowed_headers == ["x-user-email", "x-tenant-id"]

    def test_http_mcp_config_allowed_headers_default_none(self):
        """Test that allowed_headers defaults to None."""
        config = HttpMcpServerConfig(
            params=StreamableHTTPConnectionParams(url="http://localhost:8080"),
        )
        assert config.allowed_headers is None

    def test_sse_mcp_config_has_allowed_headers(self):
        """Test that SseMcpServerConfig has allowed_headers field."""
        config = SseMcpServerConfig(
            params=SseConnectionParams(url="http://localhost:8080"),
            allowed_headers=["x-user-email"],
        )
        assert config.allowed_headers == ["x-user-email"]

    def test_sse_mcp_config_allowed_headers_default_none(self):
        """Test that allowed_headers defaults to None."""
        config = SseMcpServerConfig(
            params=SseConnectionParams(url="http://localhost:8080"),
        )
        assert config.allowed_headers is None
