"""Helper functions for Microsoft Purview resource processing.

This module provides utilities for:
- Filtering Purview resources by type
- Building relationships between Purview resources and other entities
- Grouping Purview resources for analysis
"""

import logging
from typing import Dict, List, Tuple

from azure_discovery.adt_types.models import ResourceNode, ResourceRelationship

logger = logging.getLogger(__name__)


# Resource type filtering helpers

def get_sensitivity_labels(nodes: List[ResourceNode]) -> List[ResourceNode]:
    """Filter nodes to only sensitivity labels."""
    return [n for n in nodes if n.type == "Microsoft.Purview/sensitivityLabels"]


def get_dlp_policies(nodes: List[ResourceNode]) -> List[ResourceNode]:
    """Filter nodes to only DLP policies."""
    return [n for n in nodes if n.type == "Microsoft.Purview/dlpPolicies"]


def get_retention_policies(nodes: List[ResourceNode]) -> List[ResourceNode]:
    """Filter nodes to only retention policies."""
    return [n for n in nodes if n.type == "Microsoft.Purview/retentionPolicies"]


def get_retention_labels(nodes: List[ResourceNode]) -> List[ResourceNode]:
    """Filter nodes to only retention labels."""
    return [n for n in nodes if n.type == "Microsoft.Purview/retentionLabels"]


def get_information_protection_policies(nodes: List[ResourceNode]) -> List[ResourceNode]:
    """Filter nodes to only information protection policies."""
    return [n for n in nodes if n.type == "Microsoft.Purview/informationProtectionPolicies"]


def filter_policies_by_location_type(
    policies: List[ResourceNode],
    location_type: str,
) -> List[ResourceNode]:
    """Filter DLP or retention policy nodes to those that apply to the given location type.

    Locations are read from node properties or tags (list of dicts with
    'locationType' key). Matching is case-insensitive.

    Args:
        policies: List of Purview policy nodes (DLP or retention).
        location_type: Location type to match (e.g. 'SharePoint', 'Teams', 'OneDrive').

    Returns:
        List of policies that have at least one location with matching locationType.
    """
    if not policies or not location_type:
        return []
    target = location_type.strip().lower()
    if not target:
        return []
    result: List[ResourceNode] = []
    for node in policies:
        locations = node.properties.get("locations", []) or node.tags.get("locations", [])
        if not isinstance(locations, list):
            continue
        for loc in locations:
            if isinstance(loc, dict):
                lt = loc.get("locationType") or loc.get("location_type")
                if lt and str(lt).strip().lower() == target:
                    result.append(node)
                    break
    return result


def group_purview_by_type(nodes: List[ResourceNode]) -> Dict[str, List[ResourceNode]]:
    """Group Purview nodes by resource type.
    
    Args:
        nodes: List of Purview resource nodes
        
    Returns:
        Dictionary mapping resource type to list of nodes
    """
    grouped: Dict[str, List[ResourceNode]] = {}
    
    for node in nodes:
        if node.type.startswith("Microsoft.Purview/"):
            if node.type not in grouped:
                grouped[node.type] = []
            grouped[node.type].append(node)
    
    return grouped


# Relationship building

def build_purview_relationships(
    purview_nodes: List[ResourceNode],
    all_nodes: List[ResourceNode] | None = None,
    m365_nodes: List[ResourceNode] | None = None,
    materialize_missing: bool = False,
) -> Tuple[List[ResourceRelationship], List[ResourceNode]]:
    """Build relationships between Purview resources and other entities.

    Creates edges for:
    - Sensitivity Label -> File/Document (when labels are assigned)
    - DLP Policy -> M365 Location (Exchange, SharePoint, OneDrive, Teams)
    - Retention Policy -> M365 Location
    - Retention Label -> File/Document (when labels are assigned)

    Args:
        purview_nodes: List of Purview resource nodes
        all_nodes: List of all nodes in the discovery graph (used when m365_nodes not provided)
        m365_nodes: Optional list of M365 nodes only (tests/convenience); used instead of all_nodes when provided
        materialize_missing: If True, create placeholder nodes for missing resources

    Returns:
        Tuple of (relationships, materialized_nodes)
    """
    nodes = m365_nodes if m365_nodes is not None else (all_nodes or [])
    relationships: List[ResourceRelationship] = []
    materialized: List[ResourceNode] = []

    # Build index of existing nodes by ID for fast lookup
    existing_nodes_by_id = {node.id: node for node in nodes}

    # Build index of existing nodes by type and tags
    existing_m365_sites = {
        node.tags.get("siteId"): node
        for node in nodes
        if node.type == "Microsoft.SharePoint/sites" and "siteId" in node.tags
    }
    existing_m365_teams = {
        node.tags.get("teamId"): node
        for node in nodes
        if node.type == "Microsoft.Teams/teams" and "teamId" in node.tags
    }
    existing_m365_drives = {
        node.tags.get("driveId"): node
        for node in nodes
        if node.type == "Microsoft.OneDrive/drives" and "driveId" in node.tags
    }
    
    # Process each Purview node type
    for node in purview_nodes:
        if node.type == "Microsoft.Purview/dlpPolicies":
            # Build DLP Policy -> M365 Location edges
            relationships_batch = _build_dlp_policy_relationships(
                node,
                existing_m365_sites,
                existing_m365_teams,
                existing_m365_drives,
            )
            relationships.extend(relationships_batch)
            
        elif node.type == "Microsoft.Purview/retentionPolicies":
            # Build Retention Policy -> M365 Location edges
            relationships_batch = _build_retention_policy_relationships(
                node,
                existing_m365_sites,
                existing_m365_teams,
                existing_m365_drives,
            )
            relationships.extend(relationships_batch)
            
        elif node.type == "Microsoft.Purview/sensitivityLabels":
            # Sensitivity labels are typically assigned to documents, not enumerable via Graph
            # These would be discovered through document-level API calls (not implemented)
            pass
            
        elif node.type == "Microsoft.Purview/retentionLabels":
            # Retention labels are typically assigned to documents, not enumerable via Graph
            # These would be discovered through document-level API calls (not implemented)
            pass
    
    logger.info(
        f"Built {len(relationships)} Purview relationships "
        f"({len(materialized)} materialized nodes)"
    )
    
    return relationships, materialized


def _build_dlp_policy_relationships(
    policy_node: ResourceNode,
    existing_sites: Dict[str, ResourceNode],
    existing_teams: Dict[str, ResourceNode],
    existing_drives: Dict[str, ResourceNode],
) -> List[ResourceRelationship]:
    """Build relationships from DLP policy to M365 locations.
    
    Args:
        policy_node: DLP policy node
        existing_sites: Index of SharePoint sites by site ID
        existing_teams: Index of Teams by team ID
        existing_drives: Index of OneDrive drives by drive ID
        
    Returns:
        List of relationships
    """
    relationships: List[ResourceRelationship] = []
    locations = policy_node.properties.get("locations", []) or policy_node.tags.get(
        "locations", []
    )
    if not isinstance(locations, list):
        return relationships

    for location in locations:
        if not isinstance(location, dict):
            continue
        location_type = (
            location.get("type") or location.get("locationType") or ""
        ).lower()
        location_id = location.get("id")
        location_url = location.get("location")
        if not location_type:
            continue

        # Normalize locationType (e.g. SharePoint -> sharepointonline)
        if location_type == "sharepoint":
            location_type = "sharepointonline"
        elif location_type == "teams":
            location_type = "teamsforbusiness"
        elif location_type == "onedrive":
            location_type = "onedriveforbusiness"

        targets: List[ResourceNode] = []
        if location_type == "sharepointonline":
            if location_id and location_id in existing_sites:
                targets = [existing_sites[location_id]]
            elif location_url == "All":
                targets = list(existing_sites.values())
            elif location_url:
                for node in existing_sites.values():
                    if node.properties.get("webUrl") == location_url:
                        targets = [node]
                        break
        elif location_type == "teamsforbusiness":
            if location_id and location_id in existing_teams:
                targets = [existing_teams[location_id]]
            elif location_url == "All":
                targets = list(existing_teams.values())
        elif location_type == "onedriveforbusiness":
            if location_id and location_id in existing_drives:
                targets = [existing_drives[location_id]]
            elif location_url == "All":
                targets = list(existing_drives.values())

        for target_node in targets:
            relationships.append(
                ResourceRelationship(
                    source_id=policy_node.id,
                    target_id=target_node.id,
                    relation_type="applies_to",
                    weight=1.0,
                )
            )
    return relationships


def _build_retention_policy_relationships(
    policy_node: ResourceNode,
    existing_sites: Dict[str, ResourceNode],
    existing_teams: Dict[str, ResourceNode],
    existing_drives: Dict[str, ResourceNode],
) -> List[ResourceRelationship]:
    """Build relationships from retention policy to M365 locations.
    
    Similar to DLP policy relationships but for retention policies.
    
    Args:
        policy_node: Retention policy node
        existing_sites: Index of SharePoint sites by site ID
        existing_teams: Index of Teams by team ID
        existing_drives: Index of OneDrive drives by drive ID
        
    Returns:
        List of relationships
    """
    relationships: List[ResourceRelationship] = []
    locations = policy_node.properties.get("locations", []) or policy_node.tags.get(
        "locations", []
    )
    if not isinstance(locations, list):
        return relationships

    for location in locations:
        if not isinstance(location, dict):
            continue
        location_type = (
            location.get("type") or location.get("locationType") or ""
        ).lower()
        location_id = location.get("id")
        location_url = location.get("location")
        if not location_type:
            continue

        if location_type == "sharepoint":
            location_type = "sharepointonline"
        elif location_type == "teams":
            location_type = "teamsforbusiness"
        elif location_type == "onedrive":
            location_type = "onedriveforbusiness"

        targets: List[ResourceNode] = []
        if location_type == "sharepointonline":
            if location_id and location_id in existing_sites:
                targets = [existing_sites[location_id]]
            elif location_url == "All":
                targets = list(existing_sites.values())
            elif location_url:
                for node in existing_sites.values():
                    if node.properties.get("webUrl") == location_url:
                        targets = [node]
                        break
        elif location_type == "teamsforbusiness":
            if location_id and location_id in existing_teams:
                targets = [existing_teams[location_id]]
            elif location_url == "All":
                targets = list(existing_teams.values())
        elif location_type == "onedriveforbusiness":
            if location_id and location_id in existing_drives:
                targets = [existing_drives[location_id]]
            elif location_url == "All":
                targets = list(existing_drives.values())

        for target_node in targets:
            relationships.append(
                ResourceRelationship(
                    source_id=policy_node.id,
                    target_id=target_node.id,
                    relation_type="retains",
                    weight=1.0,
                )
            )
    return relationships
