"""Unit tests for Defender relationship builders."""

import pytest

from azure_discovery.adt_types.models import ResourceNode, Relationship
from azure_discovery.utils.defender_helpers import (
    build_defender_relationships,
    get_high_severity_alerts,
    get_active_alerts,
    get_unhealthy_assessments,
)


@pytest.fixture
def mock_alert_node():
    """Create a mock alert node."""
    return ResourceNode(
        id="/subscriptions/sub1/providers/Microsoft.Security/locations/centralus/alerts/alert1",
        name="alert1",
        type="Microsoft.Security/alerts",
        subscription_id="sub1",
        properties={
            "alertDisplayName": "Suspicious activity",
            "severity": "High",
            "status": "Active",
            "alertType": "VM_SuspiciousActivity",
            "affectedResources": [
                "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Compute/virtualMachines/vm1"
            ],
        },
    )


@pytest.fixture
def mock_assessment_node():
    """Create a mock assessment node."""
    return ResourceNode(
        id="/subscriptions/sub1/providers/Microsoft.Security/assessments/assessment1",
        name="assessment1",
        type="Microsoft.Security/assessments",
        subscription_id="sub1",
        properties={
            "displayName": "Disk encryption missing",
            "severity": "Medium",
            "statusCode": "Unhealthy",
            "assessmentType": "BuiltIn",
            "affectedResourceId": "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Compute/disks/disk1",
        },
    )


@pytest.fixture
def mock_vm_node():
    """Create a mock VM resource node."""
    return ResourceNode(
        id="/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Compute/virtualMachines/vm1",
        name="vm1",
        type="Microsoft.Compute/virtualMachines",
        subscription_id="sub1",
        location="eastus",
        properties={},
    )


@pytest.fixture
def mock_disk_node():
    """Create a mock disk resource node."""
    return ResourceNode(
        id="/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Compute/disks/disk1",
        name="disk1",
        type="Microsoft.Compute/disks",
        subscription_id="sub1",
        location="eastus",
        properties={},
    )


def test_build_defender_relationships_with_alerts(
    mock_alert_node, mock_vm_node
):
    """Test building relationships from alerts to affected resources."""
    all_nodes = [mock_alert_node, mock_vm_node]

    relationships = build_defender_relationships(
        alerts=[mock_alert_node],
        assessments=[],
        all_nodes=all_nodes,
        materialize_missing=False,
    )

    assert len(relationships) == 1
    rel = relationships[0]
    assert isinstance(rel, Relationship)
    assert rel.source_id == mock_alert_node.id
    assert rel.target_id == mock_vm_node.id
    assert rel.relation_type == "affects"


def test_build_defender_relationships_with_assessments(
    mock_assessment_node, mock_disk_node
):
    """Test building relationships from assessments to affected resources."""
    all_nodes = [mock_assessment_node, mock_disk_node]

    relationships = build_defender_relationships(
        alerts=[],
        assessments=[mock_assessment_node],
        all_nodes=all_nodes,
        materialize_missing=False,
    )

    assert len(relationships) == 1
    rel = relationships[0]
    assert isinstance(rel, Relationship)
    assert rel.source_id == mock_assessment_node.id
    assert rel.target_id == mock_disk_node.id
    assert rel.relation_type == "affects"


def test_build_defender_relationships_both_types(
    mock_alert_node, mock_assessment_node, mock_vm_node, mock_disk_node
):
    """Test building relationships from both alerts and assessments."""
    all_nodes = [mock_alert_node, mock_assessment_node, mock_vm_node, mock_disk_node]

    relationships = build_defender_relationships(
        alerts=[mock_alert_node],
        assessments=[mock_assessment_node],
        all_nodes=all_nodes,
        materialize_missing=False,
    )

    # Should have 2 relationships (1 alert→vm, 1 assessment→disk)
    assert len(relationships) == 2
    
    alert_rels = [r for r in relationships if r.source_id == mock_alert_node.id]
    assessment_rels = [r for r in relationships if r.source_id == mock_assessment_node.id]
    
    assert len(alert_rels) == 1
    assert len(assessment_rels) == 1


def test_build_defender_relationships_missing_target_no_materialize(
    mock_alert_node
):
    """Test that relationships are skipped when target is missing and materialize=False."""
    # Alert references a VM that doesn't exist in all_nodes
    all_nodes = [mock_alert_node]

    relationships = build_defender_relationships(
        alerts=[mock_alert_node],
        assessments=[],
        all_nodes=all_nodes,
        materialize_missing=False,
    )

    # Should not create relationship since target doesn't exist
    assert len(relationships) == 0


def test_build_defender_relationships_missing_target_with_materialize(
    mock_alert_node
):
    """Test that stub nodes are created when target is missing and materialize=True."""
    # Alert references a VM that doesn't exist in all_nodes
    all_nodes = [mock_alert_node]

    relationships = build_defender_relationships(
        alerts=[mock_alert_node],
        assessments=[],
        all_nodes=all_nodes,
        materialize_missing=True,
    )

    # Should create relationship and materialize stub node
    assert len(relationships) == 1
    
    # Check that a stub node was added to all_nodes
    stub_nodes = [
        n for n in all_nodes
        if n.properties.get("materialized") is True
    ]
    assert len(stub_nodes) == 1
    assert stub_nodes[0].properties["reason"] == "defender_alert"


def test_build_defender_relationships_assessment_without_resource(
    mock_assessment_node
):
    """Test assessment without affected resource (subscription-scoped)."""
    # Remove affectedResourceId to simulate subscription-scoped assessment
    mock_assessment_node.properties.pop("affectedResourceId")
    
    # Add a subscription node
    subscription_node = ResourceNode(
        id="/subscriptions/sub1",
        name="sub1",
        type="Microsoft.Resources/subscriptions",
        subscription_id="sub1",
        properties={},
    )
    
    all_nodes = [mock_assessment_node, subscription_node]

    relationships = build_defender_relationships(
        alerts=[],
        assessments=[mock_assessment_node],
        all_nodes=all_nodes,
        materialize_missing=False,
    )

    # Should create relationship to subscription
    assert len(relationships) == 1
    assert relationships[0].target_id == "/subscriptions/sub1"


def test_build_defender_relationships_multiple_affected_resources():
    """Test alert with multiple affected resources."""
    alert = ResourceNode(
        id="/subscriptions/sub1/providers/Microsoft.Security/locations/centralus/alerts/alert1",
        name="alert1",
        type="Microsoft.Security/alerts",
        subscription_id="sub1",
        properties={
            "alertDisplayName": "Multi-resource attack",
            "severity": "Critical",
            "status": "Active",
            "alertType": "MultiResourceAttack",
            "affectedResources": [
                "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Compute/virtualMachines/vm1",
                "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Compute/virtualMachines/vm2",
                "/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Storage/storageAccounts/sa1",
            ],
        },
    )
    
    vm1 = ResourceNode(
        id="/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Compute/virtualMachines/vm1",
        name="vm1",
        type="Microsoft.Compute/virtualMachines",
        subscription_id="sub1",
        properties={},
    )
    
    vm2 = ResourceNode(
        id="/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Compute/virtualMachines/vm2",
        name="vm2",
        type="Microsoft.Compute/virtualMachines",
        subscription_id="sub1",
        properties={},
    )
    
    sa1 = ResourceNode(
        id="/subscriptions/sub1/resourceGroups/rg1/providers/Microsoft.Storage/storageAccounts/sa1",
        name="sa1",
        type="Microsoft.Storage/storageAccounts",
        subscription_id="sub1",
        properties={},
    )
    
    all_nodes = [alert, vm1, vm2, sa1]

    relationships = build_defender_relationships(
        alerts=[alert],
        assessments=[],
        all_nodes=all_nodes,
        materialize_missing=False,
    )

    # Should create 3 relationships (alert→vm1, alert→vm2, alert→sa1)
    assert len(relationships) == 3
    target_ids = {r.target_id for r in relationships}
    assert vm1.id in target_ids
    assert vm2.id in target_ids
    assert sa1.id in target_ids


def test_get_high_severity_alerts():
    """Test filtering alerts by high severity."""
    high_alert = ResourceNode(
        id="alert1",
        name="alert1",
        type="Microsoft.Security/alerts",
        subscription_id="sub1",
        properties={"severity": "High"},
    )
    
    critical_alert = ResourceNode(
        id="alert2",
        name="alert2",
        type="Microsoft.Security/alerts",
        subscription_id="sub1",
        properties={"severity": "Critical"},
    )
    
    medium_alert = ResourceNode(
        id="alert3",
        name="alert3",
        type="Microsoft.Security/alerts",
        subscription_id="sub1",
        properties={"severity": "Medium"},
    )
    
    alerts = [high_alert, critical_alert, medium_alert]
    high_severity_alerts = get_high_severity_alerts(alerts)
    
    assert len(high_severity_alerts) == 2
    severities = {a.properties["severity"] for a in high_severity_alerts}
    assert "High" in severities
    assert "Critical" in severities
    assert "Medium" not in severities


def test_get_active_alerts():
    """Test filtering alerts by active status."""
    active_alert = ResourceNode(
        id="alert1",
        name="alert1",
        type="Microsoft.Security/alerts",
        subscription_id="sub1",
        properties={"status": "Active"},
    )
    
    resolved_alert = ResourceNode(
        id="alert2",
        name="alert2",
        type="Microsoft.Security/alerts",
        subscription_id="sub1",
        properties={"status": "Resolved"},
    )
    
    alerts = [active_alert, resolved_alert]
    active_alerts = get_active_alerts(alerts)
    
    assert len(active_alerts) == 1
    assert active_alerts[0].properties["status"] == "Active"


def test_get_unhealthy_assessments():
    """Test filtering assessments by unhealthy status."""
    unhealthy_assessment = ResourceNode(
        id="assessment1",
        name="assessment1",
        type="Microsoft.Security/assessments",
        subscription_id="sub1",
        properties={"statusCode": "Unhealthy"},
    )
    
    healthy_assessment = ResourceNode(
        id="assessment2",
        name="assessment2",
        type="Microsoft.Security/assessments",
        subscription_id="sub1",
        properties={"statusCode": "Healthy"},
    )
    
    not_applicable_assessment = ResourceNode(
        id="assessment3",
        name="assessment3",
        type="Microsoft.Security/assessments",
        subscription_id="sub1",
        properties={"statusCode": "NotApplicable"},
    )
    
    assessments = [unhealthy_assessment, healthy_assessment, not_applicable_assessment]
    unhealthy_assessments = get_unhealthy_assessments(assessments)
    
    assert len(unhealthy_assessments) == 1
    assert unhealthy_assessments[0].properties["statusCode"] == "Unhealthy"
