"""Unit tests for Purview enumerator."""

from unittest.mock import AsyncMock, MagicMock

import pytest

from azure_discovery.enumerators.purview import PurviewEnumerator
from azure_discovery.adt_types.models import PurviewConfig, ResourceNode


@pytest.fixture
def mock_credential():
    """Mock Azure credential."""
    return AsyncMock()


@pytest.fixture
def purview_config():
    """Create PurviewConfig instance."""
    return PurviewConfig(
        include_sensitivity_labels=True,
        include_dlp_policies=True,
        include_retention_policies=True,
        include_retention_labels=True,
        include_information_protection_policies=True,
        sensitivity_max_labels=1000,
        dlp_max_policies=1000,
        retention_max_policies=1000,
        retention_max_labels=1000,
        dlp_include_disabled=False,
        retention_include_disabled=False,
    )


@pytest.fixture
def purview_enumerator(mock_credential, purview_config):
    """Create PurviewEnumerator instance."""
    return PurviewEnumerator(
        tenant_id="tenant-123",
        credential=mock_credential,
        config=purview_config,
    )


@pytest.fixture
def mock_sensitivity_label():
    """Create a mock sensitivity label."""
    label = MagicMock()
    label.id = "label-guid-123"
    label.name = "Confidential"
    label.description = "Confidential information"
    label.color = "#FF0000"
    label.sensitivity = 3
    label.is_active = True
    label.tooltip = "Apply to confidential documents"
    label.additional_data = {}
    return label


@pytest.fixture
def mock_dlp_policy():
    """Create a mock DLP policy."""
    policy = MagicMock()
    policy.id = "dlp-policy-guid-456"
    policy.name = "GDPR Compliance Policy"
    policy.description = "Protect GDPR sensitive data"
    policy.state = "enabled"
    policy.mode = "Enforce"
    policy.priority = 1
    
    # Mock locations
    location1 = MagicMock()
    location1.location = "https://contoso.sharepoint.com/sites/hr"
    location1.additional_data = {"locationType": "SharePoint"}
    
    location2 = MagicMock()
    location2.location = "All"
    location2.additional_data = {"locationType": "Teams"}
    
    policy.locations = [location1, location2]
    policy.additional_data = {}
    return policy


@pytest.fixture
def mock_retention_policy():
    """Create a mock retention policy."""
    policy = MagicMock()
    policy.id = "retention-policy-guid-789"
    policy.display_name = "7 Year Retention Policy"
    policy.description = "Retain records for 7 years"
    policy.is_enabled = True
    policy.retention_duration = 2555  # days
    
    # Mock locations
    location = MagicMock()
    location.location = "All"
    location.additional_data = {"locationType": "OneDrive"}
    
    policy.locations = [location]
    policy.additional_data = {}
    return policy


@pytest.fixture
def mock_retention_label():
    """Create a mock retention label."""
    label = MagicMock()
    label.id = "retention-label-guid-abc"
    label.display_name = "Permanent Record"
    label.description = "Keep permanently"
    label.is_in_use = True
    label.retention_duration = "Unlimited"
    label.additional_data = {}
    return label


@pytest.fixture
def mock_ip_policy():
    """Create a mock information protection policy."""
    policy = MagicMock()
    policy.id = "ip-policy-guid-def"
    policy.display_name = "Default Information Protection Policy"
    policy.description = "Organization-wide IP policy"
    policy.additional_data = {}
    return policy


def _make_sensitivity_labels_mock_client(mock_response):
    """Build a mock Graph client that returns the given response from sensitivity labels get()."""
    mock_client = MagicMock()
    mock_client.security.data_security_and_governance.sensitivity_labels.get = AsyncMock(
        return_value=mock_response
    )
    return mock_client


@pytest.mark.asyncio
async def test_enumerate_sensitivity_labels_success(purview_enumerator, mock_sensitivity_label):
    """Test successful enumeration of sensitivity labels."""
    mock_response = MagicMock()
    mock_response.value = [mock_sensitivity_label]
    purview_enumerator.client = _make_sensitivity_labels_mock_client(mock_response)

    labels = await purview_enumerator.enumerate_sensitivity_labels()

    assert len(labels) == 1
    label_node = labels[0]
    assert isinstance(label_node, ResourceNode)
    assert label_node.type == "Microsoft.Purview/sensitivityLabels"
    assert label_node.name == "Confidential"
    assert label_node.properties["description"] == "Confidential information"
    assert label_node.properties["color"] == "#FF0000"
    assert label_node.properties["sensitivity"] == 3
    assert label_node.properties["isActive"] is True
    assert "labelId" in label_node.tags


@pytest.mark.asyncio
async def test_enumerate_sensitivity_labels_with_filter(purview_enumerator, mock_sensitivity_label):
    """Test sensitivity label enumeration with name filter."""
    purview_enumerator.config.sensitivity_label_filter = ["Confidential"]

    public_label = MagicMock()
    public_label.id = "label-public-123"
    public_label.name = "Public"
    public_label.description = "Public information"
    public_label.color = "#00FF00"
    public_label.sensitivity = 0
    public_label.is_active = True
    public_label.additional_data = {}

    mock_response = MagicMock()
    mock_response.value = [mock_sensitivity_label, public_label]
    purview_enumerator.client = _make_sensitivity_labels_mock_client(mock_response)

    labels = await purview_enumerator.enumerate_sensitivity_labels()

    assert len(labels) == 1
    assert labels[0].name == "Confidential"


@pytest.mark.asyncio
async def test_enumerate_sensitivity_labels_max_limit(purview_enumerator, mock_sensitivity_label):
    """Test sensitivity label enumeration respects max limit."""
    purview_enumerator.config.sensitivity_max_labels = 2

    mock_response = MagicMock()
    mock_response.value = [mock_sensitivity_label] * 5
    purview_enumerator.client = _make_sensitivity_labels_mock_client(mock_response)

    labels = await purview_enumerator.enumerate_sensitivity_labels()

    assert len(labels) == 2


@pytest.mark.asyncio
async def test_enumerate_dlp_policies_success(purview_enumerator, mock_dlp_policy):
    """Test successful enumeration of DLP policies."""
    mock_dlp_policy.is_enabled = True
    mock_response = MagicMock()
    mock_response.value = [mock_dlp_policy]
    mock_client = MagicMock()
    mock_client.information_protection.data_loss_prevention_policies.get = AsyncMock(
        return_value=mock_response
    )
    purview_enumerator.client = mock_client

    policies = await purview_enumerator.enumerate_dlp_policies()

    assert len(policies) == 1
    policy_node = policies[0]
    assert isinstance(policy_node, ResourceNode)
    assert policy_node.type == "Microsoft.Purview/dlpPolicies"
    assert policy_node.name == "GDPR Compliance Policy"
    assert policy_node.properties["description"] == "Protect GDPR sensitive data"
    assert policy_node.properties["isEnabled"] is True
    assert policy_node.properties["mode"] == "Enforce"
    assert policy_node.properties["priority"] == 1
    assert "policyId" in policy_node.tags


@pytest.mark.asyncio
async def test_enumerate_dlp_policies_with_filter(purview_enumerator, mock_dlp_policy):
    """Test DLP policy enumeration with name filter."""
    purview_enumerator.config.dlp_policy_filter = ["GDPR"]

    hipaa_policy = MagicMock()
    hipaa_policy.id = "dlp-hipaa-123"
    hipaa_policy.name = "HIPAA Compliance Policy"
    hipaa_policy.description = "Protect healthcare data"
    hipaa_policy.is_enabled = True
    hipaa_policy.locations = []
    hipaa_policy.additional_data = {}

    mock_response = MagicMock()
    mock_response.value = [mock_dlp_policy, hipaa_policy]
    mock_client = MagicMock()
    mock_client.information_protection.data_loss_prevention_policies.get = AsyncMock(
        return_value=mock_response
    )
    purview_enumerator.client = mock_client

    policies = await purview_enumerator.enumerate_dlp_policies()

    assert len(policies) == 1
    assert policies[0].name == "GDPR Compliance Policy"


@pytest.mark.asyncio
async def test_enumerate_dlp_policies_exclude_disabled(purview_enumerator, mock_dlp_policy):
    """Test DLP policy enumeration excludes disabled policies."""
    mock_dlp_policy.is_enabled = True
    disabled_policy = MagicMock()
    disabled_policy.id = "dlp-disabled-123"
    disabled_policy.name = "Disabled Policy"
    disabled_policy.description = "This is disabled"
    disabled_policy.is_enabled = False
    disabled_policy.locations = []
    disabled_policy.additional_data = {}

    mock_response = MagicMock()
    mock_response.value = [mock_dlp_policy, disabled_policy]
    mock_client = MagicMock()
    mock_client.information_protection.data_loss_prevention_policies.get = AsyncMock(
        return_value=mock_response
    )
    purview_enumerator.client = mock_client

    policies = await purview_enumerator.enumerate_dlp_policies()

    assert len(policies) == 1
    assert policies[0].properties["isEnabled"] is True


@pytest.mark.asyncio
async def test_enumerate_dlp_policies_include_disabled(purview_enumerator, mock_dlp_policy):
    """Test DLP policy enumeration includes disabled policies when configured."""
    purview_enumerator.config.dlp_include_disabled = True
    mock_dlp_policy.is_enabled = True

    disabled_policy = MagicMock()
    disabled_policy.id = "dlp-disabled-123"
    disabled_policy.name = "Disabled Policy"
    disabled_policy.description = "This is disabled"
    disabled_policy.is_enabled = False
    disabled_policy.locations = []
    disabled_policy.additional_data = {}

    mock_response = MagicMock()
    mock_response.value = [mock_dlp_policy, disabled_policy]
    mock_client = MagicMock()
    mock_client.information_protection.data_loss_prevention_policies.get = AsyncMock(
        return_value=mock_response
    )
    purview_enumerator.client = mock_client

    policies = await purview_enumerator.enumerate_dlp_policies()

    assert len(policies) == 2


@pytest.mark.asyncio
async def test_enumerate_retention_policies_success(purview_enumerator, mock_retention_policy):
    """Test successful enumeration of retention policies."""
    mock_response = MagicMock()
    mock_response.value = [mock_retention_policy]
    mock_client = MagicMock()
    mock_client.security.retention_policies.get = AsyncMock(return_value=mock_response)
    purview_enumerator.client = mock_client

    policies = await purview_enumerator.enumerate_retention_policies()

    assert len(policies) == 1
    policy_node = policies[0]
    assert isinstance(policy_node, ResourceNode)
    assert policy_node.type == "Microsoft.Purview/retentionPolicies"
    assert policy_node.name == "7 Year Retention Policy"
    assert policy_node.properties["description"] == "Retain records for 7 years"
    assert policy_node.properties["isEnabled"] is True
    assert policy_node.properties["retentionDuration"] == 2555
    assert "policyId" in policy_node.tags


@pytest.mark.asyncio
async def test_enumerate_retention_policies_with_filter(purview_enumerator, mock_retention_policy):
    """Test retention policy enumeration with name filter."""
    purview_enumerator.config.retention_policy_filter = ["7 Year"]

    short_policy = MagicMock()
    short_policy.id = "retention-short-123"
    short_policy.display_name = "30 Day Retention"
    short_policy.description = "Short retention"
    short_policy.is_enabled = True
    short_policy.locations = []
    short_policy.additional_data = {}

    mock_response = MagicMock()
    mock_response.value = [mock_retention_policy, short_policy]
    mock_client = MagicMock()
    mock_client.security.retention_policies.get = AsyncMock(return_value=mock_response)
    purview_enumerator.client = mock_client

    policies = await purview_enumerator.enumerate_retention_policies()

    assert len(policies) == 1
    assert policies[0].name == "7 Year Retention Policy"


@pytest.mark.asyncio
async def test_enumerate_retention_labels_success(purview_enumerator, mock_retention_label):
    """Test successful enumeration of retention labels."""
    mock_response = MagicMock()
    mock_response.value = [mock_retention_label]
    mock_client = MagicMock()
    mock_client.security.labels.retention_labels.get = AsyncMock(return_value=mock_response)
    purview_enumerator.client = mock_client

    labels = await purview_enumerator.enumerate_retention_labels()

    assert len(labels) == 1
    label_node = labels[0]
    assert isinstance(label_node, ResourceNode)
    assert label_node.type == "Microsoft.Purview/retentionLabels"
    assert label_node.name == "Permanent Record"
    assert label_node.properties["description"] == "Keep permanently"
    assert label_node.properties["isInUse"] is True
    assert label_node.properties["retentionDuration"] == "Unlimited"
    assert "labelId" in label_node.tags


@pytest.mark.asyncio
async def test_enumerate_retention_labels_with_filter(purview_enumerator, mock_retention_label):
    """Test retention label enumeration with name filter."""
    purview_enumerator.config.retention_label_filter = ["Permanent"]

    temp_label = MagicMock()
    temp_label.id = "retention-temp-123"
    temp_label.display_name = "Temporary"
    temp_label.description = "Delete after use"
    temp_label.is_in_use = False
    temp_label.additional_data = {}

    mock_response = MagicMock()
    mock_response.value = [mock_retention_label, temp_label]
    mock_client = MagicMock()
    mock_client.security.labels.retention_labels.get = AsyncMock(return_value=mock_response)
    purview_enumerator.client = mock_client

    labels = await purview_enumerator.enumerate_retention_labels()

    assert len(labels) == 1
    assert labels[0].name == "Permanent Record"


@pytest.mark.asyncio
async def test_enumerate_information_protection_policies_success(purview_enumerator, mock_ip_policy):
    """Test successful enumeration of information protection policies."""
    mock_client = MagicMock()
    mock_client.information_protection.policy.get = AsyncMock(return_value=mock_ip_policy)
    purview_enumerator.client = mock_client

    policies = await purview_enumerator.enumerate_information_protection_policies()

    assert len(policies) == 1
    policy_node = policies[0]
    assert isinstance(policy_node, ResourceNode)
    assert policy_node.type == "Microsoft.Purview/informationProtectionPolicies"
    assert policy_node.name == "Default Information Protection Policy"
    assert policy_node.properties["description"] == "Organization-wide IP policy"
    assert "policyId" in policy_node.tags


@pytest.mark.asyncio
async def test_enumerate_empty_response(purview_enumerator):
    """Test enumeration with empty API response."""
    mock_response = MagicMock()
    mock_response.value = []
    purview_enumerator.client = _make_sensitivity_labels_mock_client(mock_response)

    labels = await purview_enumerator.enumerate_sensitivity_labels()
    assert len(labels) == 0


@pytest.mark.asyncio
async def test_enumerate_sensitivity_labels_api_error(purview_enumerator):
    """Test sensitivity label enumeration handles API errors gracefully."""
    mock_client = MagicMock()
    mock_client.security.data_security_and_governance.sensitivity_labels.get = AsyncMock(
        side_effect=Exception("Graph API error")
    )
    purview_enumerator.client = mock_client

    labels = await purview_enumerator.enumerate_sensitivity_labels()
    assert len(labels) == 0


@pytest.mark.asyncio
async def test_enumerate_dlp_policies_api_error(purview_enumerator):
    """Test DLP policy enumeration handles API errors gracefully."""
    mock_client = MagicMock()
    mock_client.information_protection.data_loss_prevention_policies.get = AsyncMock(
        side_effect=Exception("Graph API error")
    )
    purview_enumerator.client = mock_client

    policies = await purview_enumerator.enumerate_dlp_policies()
    assert len(policies) == 0


@pytest.mark.asyncio
async def test_enumerate_retention_policies_api_error(purview_enumerator):
    """Test retention policy enumeration handles API errors gracefully."""
    mock_client = MagicMock()
    mock_client.security.retention_policies.get = AsyncMock(
        side_effect=Exception("Graph API error")
    )
    purview_enumerator.client = mock_client

    policies = await purview_enumerator.enumerate_retention_policies()
    assert len(policies) == 0


@pytest.mark.asyncio
async def test_purview_config_defaults():
    """Test PurviewConfig default values."""
    config = PurviewConfig()
    
    assert config.include_sensitivity_labels is True
    assert config.include_dlp_policies is True
    assert config.include_retention_policies is True
    assert config.include_retention_labels is True
    assert config.include_information_protection_policies is True
    assert config.sensitivity_max_labels == 1000
    assert config.dlp_max_policies == 1000
    assert config.retention_max_policies == 1000
    assert config.retention_max_labels == 1000
    assert config.dlp_include_disabled is False
    assert config.retention_include_disabled is False
    assert config.include_label_assignments is True
    assert config.include_policy_assignments is True
