from __future__ import annotations

from typing import Any

import pytest

from azure_discovery.adt_types import AzureDiscoveryRequest, ResourceNode
from azure_discovery.enumerators import policy_compliance


def test_extract_controls_from_compliance_metadata_dicts() -> None:
    metadata: dict[str, Any] = {
        "compliance": {
            "complianceType": "NIST SP 800-171 Rev. 2",
            "controls": [
                {
                    "controlId": "3.1.1",
                    "controlName": "Limit system access",
                    "description": "Limit system access to authorized users.",
                }
            ],
        }
    }

    controls = policy_compliance._extract_compliance_controls(metadata)
    assert len(controls) == 1
    assert controls[0].standard == "NIST SP 800-171 Rev. 2"
    assert controls[0].control_id == "3.1.1"
    assert controls[0].control_name == "Limit system access"
    assert controls[0].description == "Limit system access to authorized users."


def test_build_mapping_expands_initiative_member_policies_with_reference_id() -> None:
    initiative_id = "/providers/Microsoft.Authorization/policySetDefinitions/nist800171r2"
    policy_id = "/providers/Microsoft.Authorization/policyDefinitions/policy1"

    definitions: dict[str, dict[str, Any]] = {
        initiative_id.lower(): {
            "name": "nist800171r2",
            "type": "microsoft.authorization/policysetdefinitions",
            "properties": {
                "displayName": "NIST SP 800-171 Rev. 2",
                "description": "NIST 800-171 R2 initiative",
                "metadata": {"category": "Regulatory Compliance"},
                "policyDefinitions": [
                    {
                        "policyDefinitionId": policy_id,
                        "policyDefinitionReferenceId": "ref-3.1.1",
                    }
                ],
            },
        },
        policy_id.lower(): {
            "name": "policy1",
            "type": "microsoft.authorization/policydefinitions",
            "properties": {
                "displayName": "Example policy",
                "metadata": {
                    "compliance": {
                        "complianceType": "NIST SP 800-171 Rev. 2",
                        "controls": ["3.1.1"],
                    }
                },
            },
        },
    }

    assignment: dict[str, Any] = {
        "id": "/subscriptions/0000/providers/Microsoft.Authorization/policyAssignments/a1",
        "name": "a1",
        "properties": {
            "policyDefinitionId": initiative_id,
            "displayName": "Assign NIST 800-171",
            "scope": "/subscriptions/0000",
        },
    }

    mapping = policy_compliance._build_compliance_mapping(assignment, definitions)
    assert mapping is not None
    assert mapping.initiative_id == initiative_id.lower()
    assert mapping.initiative_name == "NIST SP 800-171 Rev. 2"

    assert any(c.control_id == "3.1.1" for c in mapping.controls)
    assert any(
        e.control_id == "3.1.1"
        and e.policy_definition_id == policy_id.lower()
        and e.policy_definition_reference_id == "ref-3.1.1"
        for e in mapping.control_evidence
    )


@pytest.mark.asyncio
async def test_enrich_sets_flattened_noncompliance_reference_ids(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    async def fake_assignments(*_: Any, **__: Any) -> list[dict[str, Any]]:
        return [
            {
                "id": "/subscriptions/0000/providers/Microsoft.Authorization/policyAssignments/a1",
                "name": "a1",
                "scope": "/subscriptions/0000",
                "properties": {
                    "policyDefinitionId": "/providers/Microsoft.Authorization/policyDefinitions/policy1",
                    "displayName": "Assign policy1",
                    "scope": "/subscriptions/0000",
                },
            }
        ]

    async def fake_definitions(*_: Any, **__: Any) -> dict[str, dict[str, Any]]:
        policy_id = "/providers/microsoft.authorization/policydefinitions/policy1"
        return {
            policy_id: {
                "name": "policy1",
                "type": "microsoft.authorization/policydefinitions",
                "properties": {
                    "displayName": "Example policy",
                    "metadata": {
                        "compliance": {
                            "complianceType": "NIST SP 800-171 Rev. 2",
                            "controls": ["3.1.1"],
                        }
                    },
                },
            }
        }

    async def fake_states(*_: Any, **__: Any) -> list[dict[str, Any]]:
        return [
            {
                "resourceId": "/subscriptions/0000/resourcegroups/rg/providers/microsoft.storage/storageaccounts/sa1",
                "policyAssignmentId": "/subscriptions/0000/providers/microsoft.authorization/policyassignments/a1",
                "policyDefinitionId": "/providers/microsoft.authorization/policydefinitions/policy1",
                "policyDefinitionReferenceId": "ref-a",
                "complianceState": "NonCompliant",
            },
            {
                "resourceId": "/subscriptions/0000/resourcegroups/rg/providers/microsoft.storage/storageaccounts/sa1",
                "policyAssignmentId": "/subscriptions/0000/providers/microsoft.authorization/policyassignments/a1",
                "policyDefinitionId": "/providers/microsoft.authorization/policydefinitions/policy1",
                "policyDefinitionReferenceId": "ref-b",
                "complianceState": "NonCompliant",
            },
        ]

    monkeypatch.setattr(policy_compliance, "_query_policy_assignments", fake_assignments)
    monkeypatch.setattr(policy_compliance, "_query_policy_definitions", fake_definitions)
    monkeypatch.setattr(policy_compliance, "_query_policy_states", fake_states)

    nodes = [
        ResourceNode(
            id="/subscriptions/0000/resourceGroups/rg/providers/Microsoft.Storage/storageAccounts/sa1",
            name="sa1",
            type="Microsoft.Storage/storageAccounts",
            subscription_id="0000",
            resource_group="rg",
        )
    ]

    request = AzureDiscoveryRequest(tenant_id="t")

    enriched = await policy_compliance.enrich_resources_with_compliance(
        nodes=nodes,
        request=request,
        credential=object(),
        subscriptions=["0000"],
        base_url="https://management.azure.com/",
    )

    assert len(enriched) == 1
    assert len(enriched[0].compliance_mappings) == 1

    mapping = enriched[0].compliance_mappings[0]
    assert mapping.compliance_state == "NonCompliant"
    assert len(mapping.noncompliance_evidence) == 2

    assert mapping.noncompliant_reference_ids == ["ref-a", "ref-b"]
    assert mapping.noncompliant_policy_definition_ids == [
        "/providers/microsoft.authorization/policydefinitions/policy1"
    ]


@pytest.mark.asyncio
async def test_list_policy_initiatives_uses_resource_graph(monkeypatch: pytest.MonkeyPatch) -> None:
    async def fake_query_resource_graph(*_: Any, **__: Any) -> list[dict[str, Any]]:
        return [
            {
                "id": "/providers/Microsoft.Authorization/policySetDefinitions/nist800171r2",
                "name": "nist800171r2",
                "displayName": "NIST SP 800-171 Rev. 2",
                "category": "Regulatory Compliance",
                "policyType": "BuiltIn",
            }
        ]

    monkeypatch.setattr(policy_compliance, "query_resource_graph", fake_query_resource_graph)

    initiatives = await policy_compliance.list_policy_initiatives(
        credential=object(),
        subscriptions=["0000"],
        base_url="https://management.azure.com/",
        batch_size=1000,
    )

    assert len(initiatives) == 1
    assert initiatives[0].display_name == "NIST SP 800-171 Rev. 2"
