"""Azure PIM (Privileged Identity Management) enumeration for eligible role assignments."""

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

from azure.core.credentials_async import AsyncTokenCredential
from azure.core.exceptions import HttpResponseError
import aiohttp

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

LOGGER = get_logger()


class PIMEnumerator:
    """Enumerate PIM eligible role assignments for both Entra ID and Azure Resources."""

    def __init__(self, credential: AsyncTokenCredential, graph_endpoint: str, resource_manager_endpoint: str):
        """Initialize PIM enumerator.

        Args:
            credential: Azure credential for authentication
            graph_endpoint: Microsoft Graph API endpoint (e.g., https://graph.microsoft.com)
            resource_manager_endpoint: ARM endpoint (e.g., https://management.azure.com)
        """
        self.credential = credential
        self.graph_endpoint = graph_endpoint.rstrip("/")
        self.resource_manager_endpoint = resource_manager_endpoint.rstrip("/")

    async def enumerate_entra_role_eligibilities(
        self,
        include_role_eligibility_schedules: bool = True,
        include_active_requests: bool = False,
    ) -> List[ResourceNode]:
        """Enumerate Entra ID role eligibility assignments via Microsoft Graph.

        This enumerates eligible role assignments for Entra ID (directory) roles.

        Args:
            include_role_eligibility_schedules: Include currently active eligibility schedules
            include_active_requests: Include pending/active eligibility requests

        Returns:
            List of PIM eligibility ResourceNode objects
        """
        eligibilities = []

        # Get access token for Microsoft Graph
        try:
            token = await self.credential.get_token("https://graph.microsoft.com/.default")
            headers = {
                "Authorization": f"Bearer {token.token}",
                "Content-Type": "application/json",
            }
        except Exception as exc:
            LOGGER.error(
                "Failed to acquire Microsoft Graph token for PIM enumeration. "
                "Ensure you are authenticated (az login or service principal credentials) "
                "and have the required Microsoft Graph permissions: "
                "RoleEligibilitySchedule.Read.Directory, RoleManagement.Read.Directory, or RoleManagement.Read.All.",
                extra={"context": {"error": str(exc), "error_type": type(exc).__name__}},
            )
            return eligibilities

        async with aiohttp.ClientSession() as session:
            # Enumerate role eligibility schedules (current/active eligibilities)
            if include_role_eligibility_schedules:
                try:
                    url = f"{self.graph_endpoint}/v1.0/roleManagement/directory/roleEligibilitySchedules"
                    eligibilities_from_schedules = await self._fetch_paged_graph_api(
                        session, url, headers, "roleEligibilitySchedules"
                    )
                    eligibilities.extend(eligibilities_from_schedules)
                except Exception as exc:
                    LOGGER.error(
                        "Failed to enumerate Entra role eligibility schedules",
                        extra={"context": {"error": str(exc)}},
                    )

            # Enumerate role eligibility schedule requests (pending/active requests)
            if include_active_requests:
                try:
                    url = f"{self.graph_endpoint}/v1.0/roleManagement/directory/roleEligibilityScheduleRequests"
                    eligibilities_from_requests = await self._fetch_paged_graph_api(
                        session, url, headers, "roleEligibilityScheduleRequests"
                    )
                    eligibilities.extend(eligibilities_from_requests)
                except Exception as exc:
                    LOGGER.error(
                        "Failed to enumerate Entra role eligibility schedule requests",
                        extra={"context": {"error": str(exc)}},
                    )

        LOGGER.info(
            "Entra PIM enumeration complete",
            extra={"context": {"total_eligibilities": len(eligibilities)}},
        )

        return eligibilities

    async def enumerate_azure_resource_role_eligibilities(
        self,
        subscription_ids: List[str],
        scope_filter: Optional[List[str]] = None,
    ) -> List[ResourceNode]:
        """Enumerate Azure Resource role eligibility assignments via ARM.

        This enumerates eligible role assignments for Azure resources (subscriptions,
        resource groups, individual resources).

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

        Returns:
            List of PIM eligibility ResourceNode objects for Azure resources
        """
        eligibilities = []

        # Get access token for Azure Resource Manager
        try:
            token = await self.credential.get_token("https://management.azure.com/.default")
            headers = {
                "Authorization": f"Bearer {token.token}",
                "Content-Type": "application/json",
            }
        except Exception as exc:
            LOGGER.error(
                "Failed to acquire ARM token for PIM resource enumeration. "
                "Ensure you are authenticated (az login or service principal credentials) "
                "and have Reader or 'Role Based Access Control Administrator (read-only)' "
                "role at subscription or management group scope.",
                extra={"context": {"error": str(exc), "error_type": type(exc).__name__}},
            )
            return eligibilities

        async with aiohttp.ClientSession() as session:
            for subscription_id in subscription_ids:
                try:
                    LOGGER.info(
                        "Enumerating Azure resource PIM eligibilities",
                        extra={"context": {"subscription_id": subscription_id}},
                    )

                    # Enumerate role eligibility schedules for subscription scope
                    scope = f"/subscriptions/{subscription_id}"
                    url = f"{self.resource_manager_endpoint}{scope}/providers/Microsoft.Authorization/roleEligibilitySchedules"

                    eligibilities_for_sub = await self._fetch_paged_arm_api(
                        session, url, headers, subscription_id
                    )

                    # Apply scope filter if provided
                    if scope_filter:
                        eligibilities_for_sub = [
                            e for e in eligibilities_for_sub
                            if self._matches_scope(e.properties.get("scope", ""), scope_filter)
                        ]

                    eligibilities.extend(eligibilities_for_sub)

                except Exception as exc:
                    LOGGER.error(
                        "Failed to enumerate Azure resource PIM eligibilities",
                        extra={
                            "context": {
                                "subscription_id": subscription_id,
                                "error": str(exc),
                            }
                        },
                    )

        LOGGER.info(
            "Azure resource PIM enumeration complete",
            extra={"context": {"total_eligibilities": len(eligibilities)}},
        )

        return eligibilities

    async def _fetch_paged_graph_api(
        self,
        session: aiohttp.ClientSession,
        url: str,
        headers: Dict[str, str],
        node_type: str,
    ) -> List[ResourceNode]:
        """Fetch paginated results from Microsoft Graph API.

        Args:
            session: aiohttp session
            url: API URL to fetch
            headers: HTTP headers including auth token
            node_type: Type of nodes being fetched (for logging)

        Returns:
            List of ResourceNode objects
        """
        nodes = []

        while url:
            try:
                async with session.get(url, headers=headers) as response:
                    if response.status == 200:
                        data = await response.json()

                        # Parse items
                        items = data.get("value", [])
                        for item in items:
                            node = self._parse_entra_eligibility(item, node_type)
                            if node:
                                nodes.append(node)

                        # Check for next page
                        url = data.get("@odata.nextLink")
                    elif response.status == 403:
                        error_text = await response.text()
                        LOGGER.warning(
                            "Insufficient permissions for PIM Entra enumeration. "
                            "Required Microsoft Graph permissions: RoleEligibilitySchedule.Read.Directory, "
                            "RoleManagement.Read.Directory, or RoleManagement.Read.All. "
                            "Also requires Entra ID P2 or Entra ID Governance licensing.",
                            extra={"context": {"status": response.status, "url": url, "error": error_text}},
                        )
                        break
                    else:
                        error_text = await response.text()
                        LOGGER.error(
                            "Failed to fetch from Graph API",
                            extra={
                                "context": {
                                    "status": response.status,
                                    "url": url,
                                    "error": error_text,
                                }
                            },
                        )
                        break
            except Exception as exc:
                LOGGER.error(
                    "Exception during Graph API pagination",
                    extra={"context": {"url": url, "error": str(exc)}},
                )
                break

        return nodes

    async def _fetch_paged_arm_api(
        self,
        session: aiohttp.ClientSession,
        url: str,
        headers: Dict[str, str],
        subscription_id: str,
    ) -> List[ResourceNode]:
        """Fetch paginated results from ARM API.

        Args:
            session: aiohttp session
            url: API URL to fetch
            headers: HTTP headers including auth token
            subscription_id: Subscription ID for the resources

        Returns:
            List of ResourceNode objects
        """
        nodes = []

        # Add api-version parameter
        separator = "&" if "?" in url else "?"
        url = f"{url}{separator}api-version=2020-10-01"

        while url:
            try:
                async with session.get(url, headers=headers) as response:
                    if response.status == 200:
                        data = await response.json()

                        # Parse items
                        items = data.get("value", [])
                        for item in items:
                            node = self._parse_azure_resource_eligibility(item, subscription_id)
                            if node:
                                nodes.append(node)

                        # Check for next page
                        url = data.get("nextLink")
                    elif response.status == 403:
                        error_text = await response.text()
                        LOGGER.warning(
                            "Insufficient permissions for PIM Azure resource enumeration. "
                            "Required Azure RBAC role: Reader or 'Role Based Access Control Administrator (read-only)' "
                            "at subscription or management group scope.",
                            extra={"context": {"status": response.status, "subscription_id": subscription_id, "error": error_text}},
                        )
                        break
                    elif response.status == 404:
                        LOGGER.warning(
                            "PIM not configured or enabled for this subscription. "
                            "Ensure Azure PIM for resources is set up in your tenant.",
                            extra={"context": {"status": response.status, "subscription_id": subscription_id}},
                        )
                        break
                    else:
                        error_text = await response.text()
                        LOGGER.error(
                            "Failed to fetch from ARM API",
                            extra={
                                "context": {
                                    "status": response.status,
                                    "url": url,
                                    "error": error_text,
                                }
                            },
                        )
                        break
            except Exception as exc:
                LOGGER.error(
                    "Exception during ARM API pagination",
                    extra={"context": {"url": url, "error": str(exc)}},
                )
                break

        return nodes

    def _parse_entra_eligibility(self, item: Dict[str, Any], node_type: str) -> Optional[ResourceNode]:
        """Parse an Entra PIM eligibility item into a ResourceNode.

        Args:
            item: Raw API response item
            node_type: Type identifier for the node

        Returns:
            ResourceNode or None if parsing fails
        """
        try:
            item_id = item.get("id")
            if not item_id:
                return None

            # Build properties from the eligibility data
            properties = {
                "principalId": item.get("principalId"),
                "roleDefinitionId": item.get("roleDefinitionId"),
                "directoryScopeId": item.get("directoryScopeId"),
                "memberType": item.get("memberType", "Direct"),
                "status": item.get("status"),
                "scheduleInfo": item.get("scheduleInfo"),
                "createdDateTime": item.get("createdDateTime"),
                "modifiedDateTime": item.get("modifiedDateTime"),
                "eligibilityType": "EntraRole",
            }

            # Add start/end dates if available
            if "scheduleInfo" in item:
                schedule = item["scheduleInfo"]
                properties["startDateTime"] = schedule.get("startDateTime")
                if "expiration" in schedule:
                    properties["endDateTime"] = schedule["expiration"].get("endDateTime")
                    properties["expirationType"] = schedule["expiration"].get("type")

            return ResourceNode(
                id=f"graph://pim/entra/{item_id}",
                name=item_id.split("/")[-1],
                type=f"Microsoft.Graph.PIM/{node_type}",
                subscription_id="",  # Entra roles are not subscription-scoped
                location=None,
                properties=properties,
                tags={},
            )
        except Exception as exc:
            LOGGER.debug(
                "Failed to parse Entra eligibility item",
                extra={"context": {"item_id": item.get("id"), "error": str(exc)}},
            )
            return None

    def _parse_azure_resource_eligibility(
        self, item: Dict[str, Any], subscription_id: str
    ) -> Optional[ResourceNode]:
        """Parse an Azure Resource PIM eligibility item into a ResourceNode.

        Args:
            item: Raw API response item
            subscription_id: Subscription ID

        Returns:
            ResourceNode or None if parsing fails
        """
        try:
            item_id = item.get("id")
            if not item_id:
                return None

            props = item.get("properties", {})

            # Build properties from the eligibility data
            properties = {
                "scope": props.get("scope"),
                "roleDefinitionId": props.get("roleDefinitionId"),
                "principalId": props.get("principalId"),
                "principalType": props.get("principalType"),
                "memberType": props.get("memberType", "Direct"),
                "status": props.get("status"),
                "startDateTime": props.get("startDateTime"),
                "endDateTime": props.get("endDateTime"),
                "createdOn": props.get("createdOn"),
                "updatedOn": props.get("updatedOn"),
                "eligibilityType": "AzureResource",
            }

            return ResourceNode(
                id=item_id,
                name=item.get("name", item_id.split("/")[-1]),
                type="Microsoft.Authorization/roleEligibilitySchedules",
                subscription_id=subscription_id,
                location=None,
                properties=properties,
                tags={},
            )
        except Exception as exc:
            LOGGER.debug(
                "Failed to parse Azure resource eligibility item",
                extra={"context": {"item_id": item.get("id"), "error": str(exc)}},
            )
            return None

    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
