"""Azure RBAC enumeration for role assignments and definitions."""

from typing import List, Optional, Any
from datetime import datetime

from azure.mgmt.authorization.aio import AuthorizationManagementClient
from azure.core.credentials_async import AsyncTokenCredential
from azure.core.exceptions import HttpResponseError

from ..adt_types.models import ResourceNode
from ..utils.logging import get_logger

LOGGER = get_logger()


class RBACEnumerator:
    """Enumerate Azure RBAC assignments and definitions."""

    def __init__(self, credential: AsyncTokenCredential, base_url: str):
        self.credential = credential
        self.base_url = base_url

    async def enumerate_role_assignments(
        self,
        subscription_ids: List[str],
        scope_filter: Optional[List[str]] = None,
    ) -> List[ResourceNode]:
        """Enumerate role assignments across subscriptions.

        Args:
            subscription_ids: List of subscription IDs to enumerate
            scope_filter: Optional list of scope prefixes to filter by

        Returns:
            List of role assignment ResourceNode objects
        """
        assignments = []

        for subscription_id in subscription_ids:
            try:
                client = AuthorizationManagementClient(
                    credential=self.credential,
                    subscription_id=subscription_id,
                    base_url=self.base_url,
                )

                LOGGER.info(
                    "Enumerating role assignments",
                    extra={"context": {"subscription_id": subscription_id}},
                )

                async for assignment in client.role_assignments.list_for_subscription():
                    # Apply scope filter if provided
                    if scope_filter and not self._matches_scope(
                        assignment.scope, scope_filter
                    ):
                        continue

                    # Resolve role definition name
                    role_def = await self._get_role_definition(
                        client, assignment.properties.role_definition_id
                    )

                    # Build properties dict
                    properties = {
                        "scope": assignment.properties.scope,
                        "roleDefinitionId": assignment.properties.role_definition_id,
                        "roleDefinitionName": (
                            role_def.properties.role_name if role_def else "Unknown"
                        ),
                        "principalId": assignment.properties.principal_id,
                        "principalType": assignment.properties.principal_type,
                        "createdOn": (
                            assignment.properties.created_on.isoformat()
                            if assignment.properties.created_on
                            else None
                        ),
                        "createdBy": assignment.properties.created_by,
                    }

                    assignments.append(
                        ResourceNode(
                            id=assignment.id,
                            name=assignment.name,
                            type="Microsoft.Authorization/roleAssignments",
                            subscription_id=subscription_id,
                            location=None,
                            properties=properties,
                            tags={},
                        )
                    )

                await client.close()

            except HttpResponseError as exc:
                LOGGER.error(
                    "Failed to enumerate role assignments",
                    extra={
                        "context": {
                            "subscription_id": subscription_id,
                            "error": str(exc),
                        }
                    },
                )
            except Exception as exc:
                LOGGER.error(
                    "Unexpected error enumerating role assignments",
                    extra={
                        "context": {
                            "subscription_id": subscription_id,
                            "error": str(exc),
                        }
                    },
                )

        LOGGER.info(
            "Role assignment enumeration complete",
            extra={"context": {"total_assignments": len(assignments)}},
        )

        return assignments

    async def enumerate_role_definitions(
        self,
        subscription_ids: List[str],
        include_built_in: bool = True,
        include_custom: bool = True,
    ) -> List[ResourceNode]:
        """Enumerate role definitions (built-in and custom).

        Args:
            subscription_ids: List of subscription IDs to enumerate
            include_built_in: Include built-in role definitions
            include_custom: Include custom role definitions

        Returns:
            List of role definition ResourceNode objects
        """
        definitions = []
        seen_ids = set()  # De-duplicate built-in roles across subscriptions

        for subscription_id in subscription_ids:
            try:
                client = AuthorizationManagementClient(
                    credential=self.credential,
                    subscription_id=subscription_id,
                    base_url=self.base_url,
                )

                LOGGER.info(
                    "Enumerating role definitions",
                    extra={"context": {"subscription_id": subscription_id}},
                )

                async for role_def in client.role_definitions.list(
                    scope=f"/subscriptions/{subscription_id}"
                ):
                    # Filter by role type
                    if (
                        role_def.properties.role_type == "BuiltInRole"
                        and not include_built_in
                    ):
                        continue
                    if (
                        role_def.properties.role_type == "CustomRole"
                        and not include_custom
                    ):
                        continue

                    # Skip duplicates for built-in roles
                    if (
                        role_def.properties.role_type == "BuiltInRole"
                        and role_def.id in seen_ids
                    ):
                        continue
                    seen_ids.add(role_def.id)

                    # Extract permissions
                    permissions = []
                    for perm in role_def.properties.permissions or []:
                        permissions.append(
                            {
                                "actions": perm.actions or [],
                                "notActions": perm.not_actions or [],
                                "dataActions": perm.data_actions or [],
                                "notDataActions": perm.not_data_actions or [],
                            }
                        )

                    properties = {
                        "roleName": role_def.properties.role_name,
                        "roleType": role_def.properties.role_type,
                        "description": role_def.properties.description,
                        "assignableScopes": role_def.properties.assignable_scopes or [],
                        "permissions": permissions,
                    }

                    definitions.append(
                        ResourceNode(
                            id=role_def.id,
                            name=role_def.name,
                            type="Microsoft.Authorization/roleDefinitions",
                            subscription_id=subscription_id,
                            location=None,
                            properties=properties,
                            tags={},
                        )
                    )

                await client.close()

            except HttpResponseError as exc:
                LOGGER.error(
                    "Failed to enumerate role definitions",
                    extra={
                        "context": {
                            "subscription_id": subscription_id,
                            "error": str(exc),
                        }
                    },
                )
            except Exception as exc:
                LOGGER.error(
                    "Unexpected error enumerating role definitions",
                    extra={
                        "context": {
                            "subscription_id": subscription_id,
                            "error": str(exc),
                        }
                    },
                )

        LOGGER.info(
            "Role definition enumeration complete",
            extra={"context": {"total_definitions": len(definitions)}},
        )

        return definitions

    def _matches_scope(self, assignment_scope: str, scope_filter: List[str]) -> bool:
        """Check if assignment scope matches any filter pattern.

        Args:
            assignment_scope: Scope of the assignment
            scope_filter: List of scope prefixes to match

        Returns:
            True if scope matches any filter
        """
        for filter_scope in scope_filter:
            if assignment_scope.lower().startswith(filter_scope.lower()):
                return True
        return False

    async def _get_role_definition(
        self, client: AuthorizationManagementClient, role_def_id: str
    ) -> Optional[Any]:
        """Fetch role definition details.

        Args:
            client: Authorization management client
            role_def_id: Role definition ID

        Returns:
            Role definition object or None if not found
        """
        try:
            return await client.role_definitions.get_by_id(role_def_id)
        except Exception as exc:
            LOGGER.debug(
                "Could not fetch role definition",
                extra={"context": {"role_def_id": role_def_id, "error": str(exc)}},
            )
            return None
