# Copyright (c) Microsoft. All rights reserved.

"""Tests for Purview client."""

from collections.abc import AsyncGenerator
from typing import Any, cast
from unittest.mock import AsyncMock, MagicMock, patch

import httpx
import pytest
from azure.core.credentials import AccessToken

from agent_framework_purview import PurviewSettings
from agent_framework_purview._client import PurviewClient
from agent_framework_purview._exceptions import (
    PurviewAuthenticationError,
    PurviewPaymentRequiredError,
    PurviewRateLimitError,
    PurviewRequestError,
    PurviewServiceError,
)
from agent_framework_purview._models import (
    ContentActivitiesRequest,
    ContentActivitiesResponse,
    PolicyLocation,
    ProcessContentRequest,
    ProtectionScopesRequest,
)


class TestPurviewClient:
    """Test PurviewClient functionality."""

    @pytest.fixture
    def mock_credential(self) -> MagicMock:
        """Create a mock async credential."""
        from azure.core.credentials_async import AsyncTokenCredential

        credential = MagicMock(spec=AsyncTokenCredential)
        mock_token = AccessToken("fake-token", 9999999999)

        async def mock_get_token(*args, **kwargs):
            return mock_token

        credential.get_token = mock_get_token
        return credential

    @pytest.fixture
    def settings(self) -> PurviewSettings:
        """Create test settings."""
        return PurviewSettings(app_name="Test App", tenant_id="test-tenant")

    @pytest.fixture
    async def client(
        self, mock_credential: MagicMock, settings: PurviewSettings
    ) -> AsyncGenerator[PurviewClient, None]:
        """Create a PurviewClient with mock credential."""
        client = PurviewClient(mock_credential, settings, timeout=10.0)
        yield client
        await client.close()

    async def test_client_initialization(self, mock_credential: MagicMock, settings: PurviewSettings) -> None:
        """Test PurviewClient initialization."""
        client = PurviewClient(mock_credential, settings)

        assert client._credential == mock_credential
        assert client._settings == settings
        assert client._graph_uri == "https://graph.microsoft.com/v1.0"
        assert client._timeout == 10.0

        await client.close()

    async def test_get_token_async_credential(self, client: PurviewClient, mock_credential: MagicMock) -> None:
        """Test _get_token with async credential."""
        token = await client._get_token(tenant_id="test-tenant")

        assert token == "fake-token"

    async def test_get_token_sync_credential(self, settings: PurviewSettings) -> None:
        """Test _get_token with sync credential."""
        sync_credential = MagicMock()
        sync_credential.get_token = MagicMock(return_value=AccessToken("sync-token", 9999999999))

        client = PurviewClient(sync_credential, settings)

        with patch("asyncio.get_running_loop") as mock_loop:
            mock_executor = AsyncMock()
            mock_executor.return_value = AccessToken("sync-token", 9999999999)
            mock_loop.return_value.run_in_executor = mock_executor

            token = await client._get_token(tenant_id="test-tenant")

            assert token == "sync-token"

        await client.close()

    async def test_get_user_info_from_token(self, client: PurviewClient) -> None:
        """Test get_user_info_from_token extracts user info."""
        import base64
        import json

        payload = {"tid": "test-tenant", "oid": "test-user", "idtyp": "user"}
        payload_str = json.dumps(payload)
        payload_bytes = payload_str.encode("utf-8")
        payload_b64 = base64.urlsafe_b64encode(payload_bytes).decode("utf-8").rstrip("=")
        fake_token = f"header.{payload_b64}.signature"

        with patch.object(client, "_get_token", return_value=fake_token):
            user_info = await client.get_user_info_from_token(tenant_id="test-tenant")

            assert user_info["tenant_id"] == "test-tenant"
            assert user_info["user_id"] == "test-user"

    @pytest.mark.parametrize(
        "status_code,exception_type",
        [
            (401, PurviewAuthenticationError),
            (403, PurviewAuthenticationError),
            (429, PurviewRateLimitError),
            (400, PurviewRequestError),
            (404, PurviewRequestError),
            (500, PurviewServiceError),
            (502, PurviewServiceError),
        ],
    )
    async def test_post_error_handling(
        self, client: PurviewClient, content_to_process_factory, status_code: int, exception_type: type[Exception]
    ) -> None:
        """Test _post method handles different HTTP errors correctly."""
        from agent_framework_purview._models import ProcessContentResponse

        content = content_to_process_factory()
        request = ProcessContentRequest(
            content_to_process=content,
            user_id="user-123",
            tenant_id="tenant-456",
        )

        mock_response = MagicMock(spec=httpx.Response)
        mock_response.status_code = status_code
        mock_response.text = "Error message"
        mock_response.raise_for_status.side_effect = httpx.HTTPStatusError(
            "Error", request=MagicMock(), response=mock_response
        )

        with patch.object(client._client, "post", return_value=mock_response), pytest.raises(exception_type):
            await client._post(
                "https://graph.microsoft.com/v1.0/test",
                request,
                ProcessContentResponse,
                "fake-token",
            )

    async def test_process_content_success(
        self, client: PurviewClient, content_to_process_factory, mock_credential: MagicMock
    ) -> None:
        """Test process_content method success path."""
        content = content_to_process_factory("Test message")
        request = ProcessContentRequest(
            content_to_process=content,
            user_id="user-123",
            tenant_id="tenant-456",
        )

        mock_response = MagicMock(spec=httpx.Response)
        mock_response.status_code = 200
        mock_response.headers = {}
        mock_response.json.return_value = {"id": "response-123", "protectionScopeState": "notModified"}

        with patch.object(client._client, "post", return_value=mock_response):
            response = await client.process_content(request)

            assert response.id == "response-123"
            assert response.protection_scope_state == "notModified"

    async def test_get_protection_scopes_success(self, client: PurviewClient) -> None:
        """Test get_protection_scopes method success path."""
        location = PolicyLocation(**{"@odata.type": "microsoft.graph.policyLocationApplication", "value": "app-id"})
        request = ProtectionScopesRequest(
            user_id="user-123", tenant_id="tenant-456", locations=[location], correlation_id="corr-789"
        )

        mock_response = MagicMock(spec=httpx.Response)
        mock_response.status_code = 200
        mock_response.headers = {}  # Add headers attribute
        mock_response.json.return_value = {"scopeIdentifier": "scope-123", "value": []}

        with patch.object(client._client, "post", return_value=mock_response):
            response = await client.get_protection_scopes(request)

            assert response.scope_identifier == "scope-123"
            assert response.scopes == []

    async def test_get_protection_scopes_uses_etag_header_when_present(self, client: PurviewClient) -> None:
        """Test that get_protection_scopes prefers the HTTP ETag header when present."""
        from agent_framework_purview._models import ProtectionScopesResponse

        location = PolicyLocation(**{"@odata.type": "microsoft.graph.policyLocationApplication", "value": "app-id"})
        request = ProtectionScopesRequest(
            user_id="user-123", tenant_id="tenant-456", locations=[location], correlation_id="corr-789"
        )

        response_obj = ProtectionScopesResponse(scope_identifier="scope-from-body", scopes=[])

        with patch.object(
            client,
            "_post",
            return_value=(response_obj, {"etag": '"etag-from-header"'}),
        ):
            response = await client.get_protection_scopes(request)

        assert response.scope_identifier == "etag-from-header"

    async def test_post_402_returns_empty_response_when_ignore_payment_required_enabled(
        self, mock_credential: MagicMock
    ) -> None:
        """Test that 402 is suppressed when ignore_payment_required=True."""
        from agent_framework_purview._models import ProcessContentResponse

        settings = PurviewSettings(app_name="Test App", ignore_payment_required=True)
        client = PurviewClient(mock_credential, settings)

        request = ProcessContentRequest(user_id="user-123", tenant_id="tenant-456", content_to_process=cast(Any, []))

        resp = httpx.Response(402, text="Payment required", request=httpx.Request("POST", "http://test"))

        with patch.object(client._client, "post", return_value=resp):
            result = await client._post("http://test", request, ProcessContentResponse, token="fake-token")

        assert isinstance(result, ProcessContentResponse)
        await client.close()

    async def test_post_sets_request_and_response_correlation_id(self, client: PurviewClient) -> None:
        """Test that correlation_id is injected into request headers and hydrated from response headers."""
        from agent_framework_purview._models import ProcessContentResponse

        # correlation_id is optional and should be auto-generated when empty
        request = ProcessContentRequest(user_id="user-123", tenant_id="tenant-456", content_to_process=cast(Any, []))
        request.correlation_id = ""  # force auto-generation branch

        captured_headers: dict[str, str] = {}

        async def fake_post(url: str, json=None, headers=None):
            nonlocal captured_headers
            captured_headers = dict(headers or {})
            return httpx.Response(
                200,
                json={"id": "resp-1", "protectionScopeState": "notModified"},
                headers={"client-request-id": "corr-from-response"},
                request=httpx.Request("POST", url),
            )

        with patch.object(client._client, "post", side_effect=fake_post):
            result_obj, result_headers = await client._post(
                "http://test",
                request,
                ProcessContentResponse,
                token="fake-token",
                return_response=True,
            )

        assert "client-request-id" in captured_headers
        assert captured_headers["client-request-id"]
        assert result_headers["client-request-id"] == "corr-from-response"
        assert result_obj.correlation_id == "corr-from-response"

    async def test_process_content_402_returns_empty_when_ignored(self, mock_credential: MagicMock) -> None:
        """Test that process_content returns an empty response (non-tuple path) when 402 is ignored."""
        from agent_framework_purview._models import ProcessContentResponse

        settings = PurviewSettings(app_name="Test App", ignore_payment_required=True)
        client = PurviewClient(mock_credential, settings)

        req = ProcessContentRequest(user_id="user-123", tenant_id="tenant-456", content_to_process=cast(Any, []))

        mock_response = MagicMock(spec=httpx.Response)
        mock_response.status_code = 402
        mock_response.text = "Payment required"

        with patch.object(client._client, "post", return_value=mock_response):
            response = await client.process_content(req)

        assert isinstance(response, ProcessContentResponse)
        await client.close()

    async def test_post_sets_correlation_id_attribute_on_recording_span(self, client: PurviewClient) -> None:
        """Test that correlation_id is added to the active span when recording is enabled."""
        from agent_framework_purview._models import ProcessContentResponse

        request = ProcessContentRequest(user_id="user-123", tenant_id="tenant-456", content_to_process=cast(Any, []))
        request.correlation_id = "corr-123"

        class RecordingSpan:
            def __init__(self) -> None:
                self.attributes: dict[str, str] = {}

            def is_recording(self) -> bool:
                return True

            def set_attribute(self, key: str, value: str) -> None:
                self.attributes[key] = value

        span = RecordingSpan()

        with (
            patch("agent_framework_purview._client.trace.get_current_span", return_value=span),
            patch.object(
                client._client,
                "post",
                return_value=httpx.Response(
                    200,
                    json={"id": "resp-1", "protectionScopeState": "notModified"},
                    headers={},
                    request=httpx.Request("POST", "http://test"),
                ),
            ),
        ):
            await client._post("http://test", request, ProcessContentResponse, token="fake-token")

        assert span.attributes["correlation_id"] == "corr-123"

    async def test_post_uses_constructor_when_response_type_has_no_model_validate(self, client: PurviewClient) -> None:
        """Test that _post falls back to the response type constructor when model_validate is absent."""

        class DummyResponse:
            def __init__(self, **data):
                self.data = data

        request = ProcessContentRequest(user_id="user-123", tenant_id="tenant-456", content_to_process=cast(Any, []))
        request.correlation_id = "corr-123"

        with patch.object(
            client._client,
            "post",
            return_value=httpx.Response(
                200,
                json={"hello": "world"},
                headers={},
                request=httpx.Request("POST", "http://test"),
            ),
        ):
            result = await client._post("http://test", request, DummyResponse, token="fake-token")

        assert isinstance(result, DummyResponse)
        assert result.data == {"hello": "world"}

    async def test_send_content_activities_success(self, client: PurviewClient, content_to_process_factory) -> None:
        """Test send_content_activities success path."""
        request = ContentActivitiesRequest(
            user_id="user-123",
            tenant_id="tenant-456",
            content_to_process=content_to_process_factory("hello"),
            correlation_id="corr-1",
        )

        mock_response = MagicMock(spec=httpx.Response)
        mock_response.status_code = 200
        mock_response.headers = {}
        mock_response.json.return_value = {"error": None}

        with patch.object(client._client, "post", return_value=mock_response):
            resp = await client.send_content_activities(request)

        assert isinstance(resp, ContentActivitiesResponse)

    async def test_post_handles_invalid_json_response_body(self, client: PurviewClient) -> None:
        """Test that invalid JSON bodies fall back to an empty dict."""
        request = ProcessContentRequest(user_id="user-123", tenant_id="tenant-456", content_to_process=cast(Any, []))
        request.correlation_id = "corr-123"

        mock_response = MagicMock(spec=httpx.Response)
        mock_response.status_code = 200
        mock_response.headers = {}
        mock_response.json.side_effect = ValueError("not json")

        with patch.object(client._client, "post", return_value=mock_response):
            result = await client._post("http://test", request, ContentActivitiesResponse, token="fake-token")

        assert isinstance(result, ContentActivitiesResponse)

    async def test_post_deserialization_failure_raises_purview_service_error(self, client: PurviewClient) -> None:
        """Test that response deserialization errors are wrapped as PurviewServiceError."""

        class BadResponseType:
            @classmethod
            def model_validate(cls, value):
                raise RuntimeError("boom")

        request = ProcessContentRequest(user_id="user-123", tenant_id="tenant-456", content_to_process=cast(Any, []))
        request.correlation_id = "corr-123"

        mock_response = MagicMock(spec=httpx.Response)
        mock_response.status_code = 200
        mock_response.headers = {}
        mock_response.json.return_value = {"any": "data"}

        with (
            patch.object(client._client, "post", return_value=mock_response),
            pytest.raises(PurviewServiceError, match="Failed to deserialize Purview response"),
        ):
            await client._post("http://test", request, BadResponseType, token="fake-token")

    async def test_client_close(self, mock_credential: AsyncMock, settings: PurviewSettings) -> None:
        """Test client properly closes HTTP client."""
        client = PurviewClient(mock_credential, settings)

        with patch.object(client._client, "aclose", new_callable=AsyncMock) as mock_close:
            await client.close()
            mock_close.assert_called_once()

    async def test_invalid_jwt_token_format(self, client: PurviewClient) -> None:
        """Test that invalid JWT token format raises ValueError."""
        with pytest.raises(ValueError, match="Invalid JWT token format"):
            client._extract_token_info("invalid-token-without-dots")

    async def test_rate_limit_error(self, client: PurviewClient) -> None:
        """Test that 429 status code raises PurviewRateLimitError."""
        request = ProcessContentRequest(
            user_id="test-user",
            tenant_id="test-tenant",
            content_to_process=cast(Any, []),
            correlation_id="test-correlation-id",
        )

        with (
            patch.object(client, "_get_token", return_value="fake-token"),
            patch.object(
                client._client,
                "post",
                return_value=httpx.Response(429, text="Rate limited", request=httpx.Request("POST", "http://test")),
            ),
            pytest.raises(PurviewRateLimitError, match="Rate limited"),
        ):
            await client.process_content(request)

    async def test_generic_request_error(self, client: PurviewClient) -> None:
        """Test that non-200/201/202 status codes raise PurviewRequestError."""
        request = ProcessContentRequest(
            user_id="test-user",
            tenant_id="test-tenant",
            content_to_process=cast(Any, []),
            correlation_id="test-correlation-id",
        )

        with (
            patch.object(client, "_get_token", return_value="fake-token"),
            patch.object(
                client._client,
                "post",
                return_value=httpx.Response(
                    500, text="Internal server error", request=httpx.Request("POST", "http://test")
                ),
            ),
            pytest.raises(PurviewRequestError, match="Purview request failed"),
        ):
            await client.process_content(request)

    async def test_prefer_header_sent_when_process_inline_true(
        self, client: PurviewClient, content_to_process_factory
    ) -> None:
        """Test that Prefer: evaluateInline header is sent when process_inline is True."""
        content = content_to_process_factory()
        request = ProcessContentRequest(
            content_to_process=content,
            user_id="user-123",
            tenant_id="tenant-456",
            process_inline=True,
        )

        posted_headers: dict[str, str] = {}
        mock_response = MagicMock(spec=httpx.Response)
        mock_response.status_code = 200
        mock_response.headers = {}
        mock_response.json.return_value = {}

        async def capture_post(url, json, headers):
            posted_headers.update(headers)
            return mock_response

        with patch.object(client._client, "post", side_effect=capture_post):
            await client.process_content(request)

            assert "Prefer" in posted_headers
            assert posted_headers["Prefer"] == "evaluateInline"

    async def test_prefer_header_not_sent_when_process_inline_false(
        self, client: PurviewClient, content_to_process_factory
    ) -> None:
        """Test that Prefer header is not sent when process_inline is False."""
        content = content_to_process_factory()
        request = ProcessContentRequest(
            content_to_process=content,
            user_id="user-123",
            tenant_id="tenant-456",
            process_inline=False,
        )

        posted_headers: dict[str, str] = {}
        mock_response = MagicMock(spec=httpx.Response)
        mock_response.status_code = 200
        mock_response.headers = {}
        mock_response.json.return_value = {}

        async def capture_post(url, json, headers):
            posted_headers.update(headers)
            return mock_response

        with patch.object(client._client, "post", side_effect=capture_post):
            await client.process_content(request)

            assert "Prefer" not in posted_headers

    async def test_prefer_header_not_sent_when_process_inline_none(
        self, client: PurviewClient, content_to_process_factory
    ) -> None:
        """Test that Prefer header is not sent when process_inline is None."""
        content = content_to_process_factory()
        request = ProcessContentRequest(
            content_to_process=content,
            user_id="user-123",
            tenant_id="tenant-456",
            process_inline=None,
        )

        posted_headers: dict[str, str] = {}
        mock_response = MagicMock(spec=httpx.Response)
        mock_response.status_code = 200
        mock_response.headers = {}
        mock_response.json.return_value = {}

        async def capture_post(url, json, headers):
            posted_headers.update(headers)
            return mock_response

        with patch.object(client._client, "post", side_effect=capture_post):
            await client.process_content(request)

            assert "Prefer" not in posted_headers

    async def test_scope_identifier_extraction_from_etag(self, client: PurviewClient) -> None:
        """Test that scope_identifier is extracted from ETag header."""
        mock_response = MagicMock(spec=httpx.Response)
        mock_response.status_code = 200
        mock_response.headers = {"etag": '"test-scope-id"'}
        mock_response.json.return_value = {"value": []}

        with patch.object(client._client, "post", return_value=mock_response):
            req = ProtectionScopesRequest(user_id="user1", tenant_id="tenant1")
            response = await client.get_protection_scopes(req)

            assert response.scope_identifier == "test-scope-id"

    async def test_scope_identifier_sent_as_if_none_match_header(
        self, client: PurviewClient, content_to_process_factory
    ) -> None:
        """Test that scope_identifier is sent as If-None-Match header."""
        content = content_to_process_factory()
        request = ProcessContentRequest(
            content_to_process=content,
            user_id="user-123",
            tenant_id="tenant-456",
            scope_identifier="test-scope-id",
        )

        posted_headers: dict[str, str] = {}
        mock_response = MagicMock(spec=httpx.Response)
        mock_response.status_code = 200
        mock_response.headers = {}
        mock_response.json.return_value = {}

        async def capture_post(url, json, headers):
            posted_headers.update(headers)
            return mock_response

        with patch.object(client._client, "post", side_effect=capture_post):
            await client.process_content(request)

            assert "If-None-Match" in posted_headers
            assert posted_headers["If-None-Match"] == "test-scope-id"

    async def test_402_payment_required_raises_exception_by_default(self, client: PurviewClient) -> None:
        """Test that 402 raises exception when ignore_payment_required is False."""
        mock_response = MagicMock(spec=httpx.Response)
        mock_response.status_code = 402
        mock_response.text = "Payment required"

        with patch.object(client._client, "post", return_value=mock_response):
            req = ProtectionScopesRequest(user_id="user1", tenant_id="tenant1")

            with pytest.raises(PurviewPaymentRequiredError):
                await client.get_protection_scopes(req)

    async def test_402_payment_required_returns_empty_when_ignored(self, mock_credential: MagicMock) -> None:
        """Test that 402 returns empty response when ignore_payment_required is True."""
        settings = PurviewSettings(app_name="Test App", ignore_payment_required=True)
        client = PurviewClient(mock_credential, settings)

        mock_response = MagicMock(spec=httpx.Response)
        mock_response.status_code = 402
        mock_response.text = "Payment required"

        with patch.object(client._client, "post", return_value=mock_response):
            req = ProtectionScopesRequest(user_id="user1", tenant_id="tenant1")
            response = await client.get_protection_scopes(req)

            # Should return empty response without raising
            assert response is not None
            assert response.scopes is None or response.scopes == []

        await client.close()
