# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

import os
from collections.abc import MutableSequence
from unittest.mock import MagicMock, patch

import pytest
from unit.mocks import get_audio_message_piece, get_sample_conversations

from pyrit.models import Message, MessagePiece, flatten_to_message_pieces
from pyrit.prompt_target import PromptShieldTarget


@pytest.fixture
def audio_message_piece() -> MessagePiece:
    return get_audio_message_piece()


@pytest.fixture
def sample_conversations() -> MutableSequence[MessagePiece]:
    conversations = get_sample_conversations()
    return flatten_to_message_pieces(conversations)


@pytest.fixture
def promptshield_target(sqlite_instance) -> PromptShieldTarget:
    return PromptShieldTarget(endpoint="mock", api_key="mock")


@pytest.fixture
def sample_delineated_prompt_as_str() -> str:
    sample: str = """
    Mock userPrompt
    <document>
    mock document
    </document>
    """
    return sample


@pytest.fixture
def sample_delineated_prompt_as_dict() -> dict:
    sample: dict = {"userPrompt": "\n    Mock userPrompt\n    ", "documents": ["\n    mock document\n    "]}
    return sample


@pytest.fixture
def sample_conversation_piece(sample_delineated_prompt_as_str: str) -> MessagePiece:
    return MessagePiece(role="user", original_value=sample_delineated_prompt_as_str)


def test_promptshield_init(promptshield_target: PromptShieldTarget):
    assert promptshield_target


async def test_prompt_shield_validate_request_length(promptshield_target: PromptShieldTarget):
    request = Message(
        message_pieces=[
            MessagePiece(role="user", conversation_id="123", original_value="test1"),
            MessagePiece(role="user", conversation_id="123", original_value="test2"),
        ]
    )
    with pytest.raises(
        ValueError,
        match="This target only supports a single message piece.*If your target does support this, set the"
        " custom_configuration parameter accordingly",
    ):
        await promptshield_target.send_prompt_async(message=request)


async def test_prompt_shield_reject_non_text(
    promptshield_target: PromptShieldTarget, audio_message_piece: MessagePiece
):
    with pytest.raises(ValueError):
        await promptshield_target.send_prompt_async(message=Message(message_pieces=[audio_message_piece]))


async def test_prompt_shield_document_parsing(
    promptshield_target: PromptShieldTarget, sample_delineated_prompt_as_str: str, sample_delineated_prompt_as_dict
):
    result = promptshield_target._input_parser(sample_delineated_prompt_as_str)

    assert result == sample_delineated_prompt_as_dict


async def test_prompt_shield_response_validation(promptshield_target: PromptShieldTarget):
    # This tests handling both an empty request and an empty response
    promptshield_target._validate_response(request_body={}, response_body={})


def test_api_key_authentication():
    """Test that API key authentication works correctly."""
    target = PromptShieldTarget(endpoint="https://test.endpoint.com", api_key="test_key")

    # Verify target was created successfully with API key
    assert target is not None
    assert target._api_key == "test_key"


def test_token_provider_authentication():
    """Test that token provider (callable) authentication works correctly."""
    token_provider = MagicMock(return_value="test_token")
    target = PromptShieldTarget(endpoint="https://test.endpoint.com", api_key=token_provider)

    # Verify target was created successfully with token provider
    assert target is not None
    assert target._api_key == token_provider
    assert callable(target._api_key)


def test_add_auth_header_with_callable_api_key():
    """Test that _add_auth_param_to_headers calls the token provider and sets Bearer token."""
    token_provider = MagicMock(return_value="test_token")
    target = PromptShieldTarget(endpoint="https://test.endpoint.com", api_key=token_provider)

    headers: dict[str, str] = {}
    target._add_auth_param_to_headers(headers)
    token_provider.assert_called_once()
    assert headers["Authorization"] == "Bearer test_token"


def test_add_auth_header_with_string_api_key():
    """Test that _add_auth_param_to_headers sets Ocp-Apim-Subscription-Key for string keys."""
    target = PromptShieldTarget(endpoint="https://test.endpoint.com", api_key="my_key")

    headers: dict[str, str] = {}
    target._add_auth_param_to_headers(headers)
    assert headers["Ocp-Apim-Subscription-Key"] == "my_key"


def test_init_raises_when_endpoint_none():
    """A missing endpoint raises ValueError."""
    with patch("pyrit.prompt_target.prompt_shield_target.default_values") as mock_dv:
        mock_dv.get_required_value = MagicMock(return_value=None)
        with pytest.raises(ValueError, match="Endpoint value is required"):
            PromptShieldTarget(endpoint=None, api_key="test_key")


def test_init_raises_when_no_api_key_and_non_azure_endpoint(sqlite_instance):
    """No key + a non-Azure endpoint raises (identity auth only works for Azure endpoints)."""
    with patch.dict(os.environ, {}, clear=False):
        os.environ.pop("AZURE_CONTENT_SAFETY_API_KEY", None)
        with pytest.raises(ValueError, match="API key is required for non-Azure"):
            PromptShieldTarget(endpoint="https://test.endpoint.com", api_key=None)


def test_init_uses_identity_token_provider_for_azure_endpoint(sqlite_instance):
    """No key + a recognized Azure Content Safety endpoint falls back to an Entra ID token provider."""
    token_provider = MagicMock(return_value="minted-token")
    with patch.dict(os.environ, {}, clear=False):
        os.environ.pop("AZURE_CONTENT_SAFETY_API_KEY", None)
        with patch(
            "pyrit.prompt_target.prompt_shield_target.get_azure_token_provider",
            return_value=token_provider,
        ) as mock_provider:
            target = PromptShieldTarget(endpoint="https://myresource.cognitiveservices.azure.com", api_key=None)

    mock_provider.assert_called_once_with("https://cognitiveservices.azure.com/.default")
    assert target._api_key is token_provider


def test_supported_auth_modes_includes_identity():
    """Prompt Shield advertises identity-based auth alongside api_key."""
    assert PromptShieldTarget.supported_auth_modes == ("api_key", "identity")