"""Azure Policy compliance mapping logic."""

from __future__ import annotations

import re
from typing import Any, Dict, List, Tuple

from azure.core.credentials import TokenCredential

from ..adt_types import (
    AzureDiscoveryRequest,
    ComplianceControl,
    ComplianceMapping,
    PolicyControlEvidence,
    PolicyInitiativeInfo,
    PolicyStateEvidence,
    ResourceNode,
)
from ..utils import query_resource_graph
from ..utils.logging import get_logger

LOGGER = get_logger()

# Known compliance standards and their patterns in policy metadata
COMPLIANCE_STANDARDS = {
    "NIST SP 800-171": r"NIST\s*SP\s*800[-\s]*171",
    "NIST SP 800-53": r"NIST\s*SP\s*800[-\s]*53",
    "CMMC": r"CMMC",
    "ISO 27001": r"ISO\s*27001",
    "PCI DSS": r"PCI\s*DSS",
    "HIPAA": r"HIPAA",
    "SOC 2": r"SOC\s*2",
    "FedRAMP": r"FedRAMP",
    "CIS": r"CIS\s*(Benchmark|Controls)",
    "Azure Security Benchmark": r"Azure\s*Security\s*Benchmark",
}


def _safe_dict(value: Any) -> Dict[str, Any]:
    if isinstance(value, dict):
        return value
    return {}


def _safe_list(value: Any) -> List[Any]:
    if isinstance(value, list):
        return value
    return []


def _dedupe_controls(controls: List[ComplianceControl]) -> List[ComplianceControl]:
    seen: dict[Tuple[str, str], ComplianceControl] = {}
    ordered: list[ComplianceControl] = []

    for control in controls:
        key = (control.standard, control.control_id)
        if key not in seen:
            seen[key] = control
            ordered.append(control)
            continue

        existing = seen[key]
        merged = existing.model_copy(
            update={
                "control_name": existing.control_name or control.control_name,
                "description": existing.description or control.description,
            }
        )
        seen[key] = merged
        for idx, item in enumerate(ordered):
            if item.standard == existing.standard and item.control_id == existing.control_id:
                ordered[idx] = merged
                break

    return ordered


def _dedupe_control_evidence(evidence: List[PolicyControlEvidence]) -> List[PolicyControlEvidence]:
    seen: set[tuple[str, str, str, str | None]] = set()
    ordered: list[PolicyControlEvidence] = []

    for item in evidence:
        key = (
            item.standard,
            item.control_id,
            item.policy_definition_id,
            item.policy_definition_reference_id,
        )
        if key in seen:
            continue
        seen.add(key)
        ordered.append(item)

    return ordered


def _extract_compliance_controls(
    policy_metadata: Dict[str, Any],
    *,
    policy_display_name: str | None = None,
    policy_category: str | None = None,
) -> List[ComplianceControl]:
    """Extract compliance controls from policy metadata."""

    controls: List[ComplianceControl] = []

    if not policy_metadata:
        return controls

    compliance_metadata = _safe_dict(policy_metadata.get("compliance"))

    # Handle Azure built-in regulatory compliance format
    if "complianceType" in compliance_metadata:
        standard = str(compliance_metadata.get("complianceType") or "Unknown")
        raw_controls = compliance_metadata.get("controls")

        for raw in _safe_list(raw_controls):
            if isinstance(raw, str):
                controls.append(ComplianceControl(control_id=raw, standard=standard))
                continue

            if isinstance(raw, dict):
                control_id = str(
                    raw.get("id")
                    or raw.get("controlId")
                    or raw.get("control_id")
                    or raw.get("control")
                    or "Unknown"
                )
                controls.append(
                    ComplianceControl(
                        control_id=control_id,
                        standard=standard,
                        control_name=raw.get("name") or raw.get("controlName"),
                        description=raw.get("description") or raw.get("details"),
                    )
                )

    # Handle NIST 800-171 format (best-effort)
    if "nist800171" in policy_metadata:
        nist_controls = policy_metadata.get("nist800171")
        for raw in _safe_list(nist_controls):
            if isinstance(raw, str):
                controls.append(
                    ComplianceControl(
                        control_id=raw,
                        standard="NIST SP 800-171 Rev. 2",
                    )
                )
            elif isinstance(raw, dict):
                controls.append(
                    ComplianceControl(
                        control_id=raw.get("id") or raw.get("controlId") or "Unknown",
                        standard="NIST SP 800-171 Rev. 2",
                        control_name=raw.get("name") or raw.get("controlName"),
                        description=raw.get("description") or raw.get("details"),
                    )
                )

    # Handle NIST 800-53 format
    if "nist80053" in policy_metadata:
        nist_controls = policy_metadata.get("nist80053")
        for raw in _safe_list(nist_controls):
            if isinstance(raw, str):
                controls.append(
                    ComplianceControl(
                        control_id=raw,
                        standard="NIST SP 800-53 Rev. 5",
                    )
                )
            elif isinstance(raw, dict):
                controls.append(
                    ComplianceControl(
                        control_id=raw.get("id") or raw.get("controlId") or "Unknown",
                        standard="NIST SP 800-53 Rev. 5",
                        control_name=raw.get("name") or raw.get("controlName"),
                        description=raw.get("description") or raw.get("details"),
                    )
                )

    # Handle CMMC format
    if "cmmc" in policy_metadata:
        cmmc_controls = policy_metadata.get("cmmc")
        for raw in _safe_list(cmmc_controls):
            if isinstance(raw, str):
                controls.append(
                    ComplianceControl(
                        control_id=raw,
                        standard="CMMC 2.0",
                    )
                )
            elif isinstance(raw, dict):
                controls.append(
                    ComplianceControl(
                        control_id=raw.get("id") or raw.get("controlId") or "Unknown",
                        standard="CMMC 2.0",
                        control_name=raw.get("name") or raw.get("controlName"),
                        description=raw.get("description") or raw.get("details"),
                    )
                )

    # Handle generic compliance controls array
    if isinstance(policy_metadata.get("controls"), list):
        for control in policy_metadata["controls"]:
            if not isinstance(control, dict):
                continue

            control_id = str(control.get("id") or control.get("controlId") or "Unknown")
            standard = str(control.get("standard") or control.get("complianceStandard") or "Unknown")

            controls.append(
                ComplianceControl(
                    control_id=control_id,
                    standard=standard,
                    control_name=control.get("name") or control.get("controlName"),
                    description=control.get("description") or control.get("details"),
                )
            )

    # Fallback: infer standard + potential control IDs from category/displayName
    if not controls:
        category = policy_category or str(policy_metadata.get("category") or "")
        display_name = policy_display_name or str(policy_metadata.get("displayName") or "")

        for standard_name, pattern in COMPLIANCE_STANDARDS.items():
            if re.search(pattern, category, re.IGNORECASE) or re.search(pattern, display_name, re.IGNORECASE):
                # Common control id shapes: 3.1.1, AC-2, IA 5.1, etc.
                control_pattern = r"(\d+\.\d+\.\d+|[A-Z]{2,}[-\s]?\d+(?:\.\d+)*)"
                matches = re.findall(control_pattern, display_name)

                if matches:
                    for match in matches:
                        controls.append(
                            ComplianceControl(
                                control_id=match,
                                standard=standard_name,
                            )
                        )
                else:
                    controls.append(
                        ComplianceControl(
                            control_id="General",
                            standard=standard_name,
                        )
                    )
                break

    return _dedupe_controls(controls)


async def list_policy_initiatives(
    credential: TokenCredential,
    subscriptions: List[str],
    base_url: str,
    batch_size: int,
    name_contains: str | None = None,
) -> List[PolicyInitiativeInfo]:
    """List policy initiatives (policy set definitions) visible in Resource Graph."""

    where_filter = ""
    if name_contains and name_contains.strip():
        # Kusto string literal
        safe = name_contains.replace("'", "''")
        where_filter = f"| where tostring(properties.displayName) contains '{safe}'"

    query = f"""
    policyresources
    | where type == 'microsoft.authorization/policysetdefinitions'
    {where_filter}
    | project
        id = tolower(id),
        name,
        displayName = tostring(properties.displayName),
        category = tostring(properties.metadata.category),
        policyType = tostring(properties.policyType)
    | order by displayName asc
    """

    results = await query_resource_graph(
        credential=credential,
        query=query,
        subscriptions=subscriptions,
        batch_size=batch_size,
        base_url=base_url,
    )

    initiatives: list[PolicyInitiativeInfo] = []
    for item in results:
        initiatives.append(
            PolicyInitiativeInfo(
                id=str(item.get("id") or ""),
                name=str(item.get("name") or ""),
                display_name=str(item.get("displayName") or "") or None,
                category=str(item.get("category") or "") or None,
                policy_type=str(item.get("policyType") or "") or None,
            )
        )

    return initiatives


async def _query_policy_assignments(
    credential: TokenCredential,
    subscriptions: List[str],
    base_url: str,
    batch_size: int,
) -> List[Dict[str, Any]]:
    """Query all policy assignments across subscriptions."""

    query = """
    policyresources
    | where type == 'microsoft.authorization/policyassignments'
    | project
        id,
        name,
        properties,
        subscriptionId,
        scope = tolower(properties.scope)
    """

    results = await query_resource_graph(
        credential=credential,
        query=query,
        subscriptions=subscriptions,
        batch_size=batch_size,
        base_url=base_url,
    )

    LOGGER.info(
        f"Retrieved {len(results)} policy assignments",
        extra={"context": {"count": len(results)}},
    )

    return results


async def _query_policy_states(
    credential: TokenCredential,
    subscriptions: List[str],
    base_url: str,
    batch_size: int,
) -> List[Dict[str, Any]]:
    """Query non-compliant policy states for resources."""

    # NOTE: keep row-level results so we can attach policyDefinitionReferenceId evidence.
    query = """
    policyresources
    | where type == 'microsoft.policyinsights/policystates'
    | where properties.complianceState != 'Compliant'
    | project
        resourceId = tolower(properties.resourceId),
        policyAssignmentId = tolower(properties.policyAssignmentId),
        policyDefinitionId = tolower(properties.policyDefinitionId),
        policyDefinitionReferenceId = tostring(properties.policyDefinitionReferenceId),
        complianceState = tostring(properties.complianceState)
    """

    try:
        results = await query_resource_graph(
            credential=credential,
            query=query,
            subscriptions=subscriptions,
            batch_size=batch_size,
            base_url=base_url,
        )

        LOGGER.info(
            f"Retrieved {len(results)} non-compliant policy state rows",
            extra={"context": {"count": len(results)}},
        )

        return results
    except Exception as exc:  # noqa: BLE001
        LOGGER.warning(
            f"Failed to query policy states: {exc}",
            extra={"context": {"error": str(exc)}},
        )
        return []


async def _query_policy_definitions(
    credential: TokenCredential,
    subscriptions: List[str],
    base_url: str,
    batch_size: int,
) -> Dict[str, Dict[str, Any]]:
    """Query policy definitions and initiatives to extract metadata."""

    query = """
    policyresources
    | where type in ('microsoft.authorization/policydefinitions', 'microsoft.authorization/policysetdefinitions')
    | project
        id = tolower(id),
        name,
        type,
        properties
    """

    results = await query_resource_graph(
        credential=credential,
        query=query,
        subscriptions=subscriptions,
        batch_size=batch_size,
        base_url=base_url,
    )

    definitions: Dict[str, Dict[str, Any]] = {}
    for item in results:
        item_id = str(item.get("id") or "").lower()
        if not item_id:
            continue
        definitions[item_id] = {
            "name": item.get("name", ""),
            "type": item.get("type", ""),
            "properties": item.get("properties", {}),
        }

    LOGGER.info(
        f"Retrieved {len(definitions)} policy definitions/initiatives",
        extra={"context": {"count": len(definitions)}},
    )

    return definitions


def _extract_controls_from_definition(definition: Dict[str, Any]) -> List[ComplianceControl]:
    props = _safe_dict(definition.get("properties"))
    metadata = _safe_dict(props.get("metadata"))
    display_name = props.get("displayName")

    category = None
    if "category" in metadata:
        category = metadata.get("category")
    else:
        category = _safe_dict(metadata.get("category")).get("name")

    return _extract_compliance_controls(
        metadata,
        policy_display_name=str(display_name) if display_name else None,
        policy_category=str(category) if category else None,
    )


def _extract_initiative_member_policies(
    initiative_definition: Dict[str, Any],
) -> List[tuple[str, str | None]]:
    props = _safe_dict(initiative_definition.get("properties"))
    member_defs = _safe_list(props.get("policyDefinitions"))

    member_policies: list[tuple[str, str | None]] = []
    for entry in member_defs:
        if not isinstance(entry, dict):
            continue
        raw_id = entry.get("policyDefinitionId")
        if not isinstance(raw_id, str) or not raw_id.strip():
            continue

        ref_id = entry.get("policyDefinitionReferenceId")
        ref_str = str(ref_id) if isinstance(ref_id, str) and ref_id.strip() else None
        member_policies.append((raw_id.lower(), ref_str))

    return member_policies


def _build_compliance_mapping(
    assignment: Dict[str, Any],
    definitions: Dict[str, Dict[str, Any]],
    compliance_state: str | None = None,
) -> ComplianceMapping | None:
    """Build a compliance mapping from a policy assignment."""

    properties = _safe_dict(assignment.get("properties"))

    policy_definition_id = str(properties.get("policyDefinitionId") or "").lower()
    assignment_id = str(assignment.get("id") or "")
    assignment_name = str(assignment.get("name") or "Unknown")

    if not policy_definition_id:
        return None

    is_initiative = "policysetdefinitions" in policy_definition_id

    initiative_id = policy_definition_id if is_initiative else None
    initiative_name = None

    controls: List[ComplianceControl] = []
    control_evidence: list[PolicyControlEvidence] = []

    policy_definition_display_name: str | None = None
    policy_definition_description: str | None = None

    definition = definitions.get(policy_definition_id)
    if definition:
        definition_props = _safe_dict(definition.get("properties"))
        policy_definition_display_name = definition_props.get("displayName")
        policy_definition_description = definition_props.get("description")

        if is_initiative:
            initiative_name = definition_props.get("displayName") or definition.get("name")

            # Expand the initiative and aggregate controls from included policy definitions.
            for member_id, reference_id in _extract_initiative_member_policies(definition):
                member_def = definitions.get(member_id)
                if not member_def:
                    continue

                member_props = _safe_dict(member_def.get("properties"))
                member_display_name = member_props.get("displayName")
                member_description = member_props.get("description")

                member_controls = _extract_controls_from_definition(member_def)
                controls.extend(member_controls)

                for c in member_controls:
                    control_evidence.append(
                        PolicyControlEvidence(
                            standard=c.standard,
                            control_id=c.control_id,
                            policy_definition_id=member_id,
                            policy_definition_reference_id=reference_id,
                            policy_definition_display_name=str(member_display_name)
                            if member_display_name
                            else None,
                            policy_definition_description=str(member_description)
                            if member_description
                            else None,
                        )
                    )

            # If the included policies didn't yield controls, fall back to initiative metadata.
            if not controls:
                metadata = _safe_dict(definition_props.get("metadata"))
                controls.extend(
                    _extract_compliance_controls(
                        metadata,
                        policy_display_name=str(policy_definition_display_name)
                        if policy_definition_display_name
                        else None,
                        policy_category=str(metadata.get("category") or "") or None,
                    )
                )
        else:
            member_controls = _extract_controls_from_definition(definition)
            controls.extend(member_controls)
            for c in member_controls:
                control_evidence.append(
                    PolicyControlEvidence(
                        standard=c.standard,
                        control_id=c.control_id,
                        policy_definition_id=policy_definition_id,
                        policy_definition_reference_id=None,
                        policy_definition_display_name=str(policy_definition_display_name)
                        if policy_definition_display_name
                        else None,
                        policy_definition_description=str(policy_definition_description)
                        if policy_definition_description
                        else None,
                    )
                )

    # If still no controls, try to infer from assignment display name.
    if not controls:
        display_name = properties.get("displayName") or assignment_name
        for standard_name, pattern in COMPLIANCE_STANDARDS.items():
            if re.search(pattern, str(display_name), re.IGNORECASE):
                controls.append(
                    ComplianceControl(
                        control_id="General",
                        standard=standard_name,
                    )
                )
                break

    controls = _dedupe_controls(controls)
    if not controls:
        return None

    # Keep evidence only for controls that survived dedupe.
    allowed = {(c.standard, c.control_id) for c in controls}
    control_evidence = [
        e for e in control_evidence if (e.standard, e.control_id) in allowed
    ]

    return ComplianceMapping(
        policy_assignment_id=assignment_id,
        policy_assignment_name=assignment_name,
        policy_definition_id=policy_definition_id,
        policy_definition_display_name=str(policy_definition_display_name)
        if policy_definition_display_name
        else None,
        policy_definition_description=str(policy_definition_description)
        if policy_definition_description
        else None,
        initiative_id=initiative_id,
        initiative_name=str(initiative_name) if initiative_name else None,
        controls=controls,
        control_evidence=_dedupe_control_evidence(control_evidence),
        compliance_state=compliance_state,
    )


async def enrich_resources_with_compliance(
    nodes: List[ResourceNode],
    request: AzureDiscoveryRequest,
    credential: TokenCredential,
    subscriptions: List[str],
    base_url: str,
) -> List[ResourceNode]:
    """Enrich resource nodes with compliance mappings from Azure Policy."""

    if not request.compliance_include_policy_assignments:
        LOGGER.info("Policy assignment enrichment disabled")
        return nodes

    assignments = await _query_policy_assignments(
        credential=credential,
        subscriptions=subscriptions,
        base_url=base_url,
        batch_size=request.max_batch_size,
    )

    definitions = await _query_policy_definitions(
        credential=credential,
        subscriptions=subscriptions,
        base_url=base_url,
        batch_size=request.max_batch_size,
    )

    policy_state_rows_by_resource: dict[str, list[dict[str, Any]]] = {}
    if request.compliance_include_policy_states:
        state_rows = await _query_policy_states(
            credential=credential,
            subscriptions=subscriptions,
            base_url=base_url,
            batch_size=request.max_batch_size,
        )

        for row in state_rows:
            resource_id = str(row.get("resourceId") or "").lower()
            if not resource_id:
                continue
            policy_state_rows_by_resource.setdefault(resource_id, []).append(row)

    scope_assignments: Dict[str, List[Dict[str, Any]]] = {}
    for assignment in assignments:
        scope = str(assignment.get("scope") or "").lower()
        if not scope:
            continue
        scope_assignments.setdefault(scope, []).append(assignment)

    enriched_nodes: List[ResourceNode] = []
    for node in nodes:
        # Skip Entra/Graph resources
        if node.id.startswith("graph://"):
            enriched_nodes.append(node)
            continue

        node_id_lower = node.id.lower()
        compliance_mappings: List[ComplianceMapping] = []

        def _append_scope(scope: str) -> None:
            if scope not in scope_assignments:
                return
            for assignment in scope_assignments[scope]:
                mapping = _build_compliance_mapping(assignment, definitions)
                if mapping:
                    compliance_mappings.append(mapping)

        # Direct resource-level assignments
        _append_scope(node_id_lower)

        # Resource group level assignments
        if node.resource_group and node.subscription_id:
            rg_scope = f"/subscriptions/{node.subscription_id}/resourcegroups/{node.resource_group}".lower()
            _append_scope(rg_scope)

        # Subscription level assignments
        if node.subscription_id:
            sub_scope = f"/subscriptions/{node.subscription_id}".lower()
            _append_scope(sub_scope)

        # Attach noncompliance evidence (per assignment) if available.
        state_rows = policy_state_rows_by_resource.get(node_id_lower, [])
        if state_rows:
            for mapping in compliance_mappings:
                assignment_id = mapping.policy_assignment_id.lower()
                relevant = [
                    r
                    for r in state_rows
                    if str(r.get("policyAssignmentId") or "").lower() == assignment_id
                ]
                if not relevant:
                    continue

                mapping.noncompliance_evidence = [
                    PolicyStateEvidence(
                        policy_assignment_id=mapping.policy_assignment_id,
                        policy_definition_id=str(r.get("policyDefinitionId") or "") or None,
                        policy_definition_reference_id=str(r.get("policyDefinitionReferenceId") or "")
                        or None,
                        compliance_state=str(r.get("complianceState") or "") or None,
                    )
                    for r in relevant
                ]

                mapping.noncompliant_reference_ids = sorted(
                    {
                        e.policy_definition_reference_id
                        for e in mapping.noncompliance_evidence
                        if e.policy_definition_reference_id
                    }
                )
                mapping.noncompliant_policy_definition_ids = sorted(
                    {
                        e.policy_definition_id
                        for e in mapping.noncompliance_evidence
                        if e.policy_definition_id
                    }
                )

                # If we have any noncompliance evidence, set a coarse compliance_state.
                if not mapping.compliance_state:
                    mapping.compliance_state = "NonCompliant"

        if request.compliance_max_assignments_per_scope > 0:
            compliance_mappings = compliance_mappings[: request.compliance_max_assignments_per_scope]

        enriched_nodes.append(node.model_copy(update={"compliance_mappings": compliance_mappings}))

    total_mappings = sum(len(n.compliance_mappings) for n in enriched_nodes)
    resources_with_mappings = sum(1 for n in enriched_nodes if n.compliance_mappings)

    LOGGER.info(
        "Compliance enrichment completed: %s resources with %s total mappings",
        resources_with_mappings,
        total_mappings,
        extra={
            "context": {
                "resources_with_mappings": resources_with_mappings,
                "total_mappings": total_mappings,
                "total_resources": len(enriched_nodes),
            }
        },
    )

    return enriched_nodes
