"""Entra ID discovery via Microsoft Graph."""

from __future__ import annotations

import asyncio
from typing import Any, Iterable, List, Optional, Tuple

from azure.core.credentials import TokenCredential

from ..adt_types import AzureDiscoveryRequest, ResourceNode, ResourceRelationship
from ..utils import build_environment_config
from ..utils.graph_relationships import build_graph_appid_edges
from ..utils.graph_mappers import (
    application_to_node,
    conditional_access_policy_to_node,
    domain_to_node,
    group_to_node,
    organization_to_node,
    risky_user_to_node,
    service_principal_to_node,
    user_to_node,
)
from ..utils.logging import get_logger

LOGGER = get_logger()


def _graph_id(kind: str, object_id: str) -> str:
    return f"graph://{kind}/{object_id}" if object_id else f"graph://{kind}/unknown"


def _safe_str(value: Any) -> str | None:
    if value is None:
        return None
    return str(value)


def _infer_kind_and_type(obj: Any) -> tuple[str, str]:
    """Best-effort inference of Graph object kind for ID prefixing."""

    odata_type = (
        getattr(obj, "odata_type", None)
        or getattr(obj, "odataType", None)
        or getattr(obj, "@odata.type", None)
    )
    if isinstance(odata_type, str):
        lowered = odata_type.lower()
        if "user" in lowered:
            return "user", "Microsoft.Graph/User"
        if "group" in lowered:
            return "group", "Microsoft.Graph/Group"
        if "serviceprincipal" in lowered:
            return "servicePrincipal", "Microsoft.Graph/ServicePrincipal"
        if "application" in lowered:
            return "application", "Microsoft.Graph/Application"

    # Fall back to attribute heuristics
    if getattr(obj, "user_principal_name", None) or getattr(obj, "userPrincipalName", None):
        return "user", "Microsoft.Graph/User"
    if getattr(obj, "group_types", None) or getattr(obj, "groupTypes", None):
        return "group", "Microsoft.Graph/Group"
    if getattr(obj, "app_id", None) or getattr(obj, "appId", None):
        # appId exists on both applications and service principals; if it looks like an SP,
        # callers should pass the right kind.
        return "directoryObject", "Microsoft.Graph/DirectoryObject"

    return "directoryObject", "Microsoft.Graph/DirectoryObject"


def _directory_object_to_node(obj: Any) -> ResourceNode:
    """Map a directoryObject-ish thing to a node.

    We use graph_mappers when possible, but gracefully degrade to a stub node.
    """

    kind, node_type = _infer_kind_and_type(obj)

    if node_type == "Microsoft.Graph/User":
        return user_to_node(obj)
    if node_type == "Microsoft.Graph/Group":
        return group_to_node(obj)
    if node_type == "Microsoft.Graph/ServicePrincipal":
        return service_principal_to_node(obj)
    if node_type == "Microsoft.Graph/Application":
        return application_to_node(obj)

    object_id = _safe_str(getattr(obj, "id", None)) or "unknown"
    name = (
        _safe_str(getattr(obj, "display_name", None))
        or _safe_str(getattr(obj, "displayName", None))
        or object_id
    )

    return ResourceNode(
        id=_graph_id(kind, object_id),
        name=name,
        type=node_type,
        subscription_id="Tenant",
        tags={"graph_id": object_id},
    )


async def get_all_pages(
    client: Any,
    request_builder: Any,
    *,
    operation_name: str,
    max_objects: int = 0,
    throttle_delay_seconds: float = 0.0,
) -> list[Any]:
    """Fetch all pages from a Graph collection.

    This is intentionally defined at module scope so unit tests can patch it.

    If `max_objects` is non-zero, results are truncated to that count.
    """

    from msgraph_core.tasks.page_iterator import PageIterator  # noqa: PLC0415

    results: list[Any] = []
    try:
        response = await request_builder.get()
        if response:
            iterator = PageIterator(
                response,
                client.request_adapter,
                lambda item: results.append(item) or True,
            )
            await iterator.iterate_async()
    except Exception as exc:  # noqa: BLE001
        LOGGER.warning(
            "Failed to fetch Graph collection",
            extra={
                "context": {
                    "operation": operation_name,
                    "error": str(exc),
                }
            },
        )

    if throttle_delay_seconds:
        await asyncio.sleep(float(throttle_delay_seconds))

    if max_objects:
        return results[:max_objects]
    return results


def _extend_unique_nodes(out: list[ResourceNode], new_nodes: Iterable[ResourceNode]) -> None:
    existing = {n.id for n in out}
    for n in new_nodes:
        if n.id in existing:
            continue
        existing.add(n.id)
        out.append(n)


async def enumerate_entra_resources(
    request: AzureDiscoveryRequest,
    credential: Optional[TokenCredential] = None,
) -> Tuple[List[ResourceNode], List[ResourceRelationship]]:
    """Enumerate Entra ID resources via Microsoft Graph."""

    if not request:
        raise ValueError("request cannot be None")

    if credential is None:
        raise ValueError("credential is required")

    from msgraph import GraphServiceClient  # noqa: PLC0415

    env_cfg = build_environment_config(request.environment)
    scopes = [f"{env_cfg.graph_endpoint}/.default"]
    client = GraphServiceClient(credentials=credential, scopes=scopes)

    nodes: List[ResourceNode] = []
    relationships: List[ResourceRelationship] = []

    remaining_total = request.graph_total_max_objects

    def _cap_for_collection(per_collection: int) -> int:
        if remaining_total and per_collection:
            return min(per_collection, remaining_total)
        if remaining_total:
            return remaining_total
        return per_collection

    # Organization (tenant root)
    org_node: ResourceNode | None = None
    if request.entra_include_organization and (not remaining_total or remaining_total > 0):
        org_items = await get_all_pages(
            client,
            client.organization,
            operation_name="entra.organization",
            max_objects=_cap_for_collection(1),
            throttle_delay_seconds=request.throttle_delay_seconds,
        )
        if org_items:
            org_node = organization_to_node(org_items[0], tenant_id=request.tenant_id)
            nodes.append(org_node)
            if remaining_total:
                remaining_total = max(0, remaining_total - 1)

    # Domains
    if request.entra_include_domains and (not remaining_total or remaining_total > 0):
        domain_items = await get_all_pages(
            client,
            client.domains,
            operation_name="entra.domains",
            max_objects=_cap_for_collection(request.entra_max_objects),
            throttle_delay_seconds=request.throttle_delay_seconds,
        )
        domain_nodes = [domain_to_node(d) for d in domain_items]
        _extend_unique_nodes(nodes, domain_nodes)
        if remaining_total:
            remaining_total = max(0, remaining_total - len(domain_items))

        if request.include_relationships and org_node:
            for dn in domain_nodes:
                relationships.append(
                    ResourceRelationship(
                        source_id=org_node.id,
                        target_id=dn.id,
                        relation_type="has_domain",
                        weight=1.0,
                    )
                )

    # Users
    if request.entra_include_users and (not remaining_total or remaining_total > 0):
        user_items = await get_all_pages(
            client,
            client.users,
            operation_name="entra.users",
            max_objects=_cap_for_collection(request.entra_max_objects),
            throttle_delay_seconds=request.throttle_delay_seconds,
        )
        _extend_unique_nodes(nodes, (user_to_node(u) for u in user_items))
        if remaining_total:
            remaining_total = max(0, remaining_total - len(user_items))

    # Groups
    group_items: list[Any] = []
    group_nodes: list[ResourceNode] = []
    if request.entra_include_groups and (not remaining_total or remaining_total > 0):
        group_items = await get_all_pages(
            client,
            client.groups,
            operation_name="entra.groups",
            max_objects=_cap_for_collection(request.entra_max_objects),
            throttle_delay_seconds=request.throttle_delay_seconds,
        )
        group_nodes = [group_to_node(g) for g in group_items]
        _extend_unique_nodes(nodes, group_nodes)
        if remaining_total:
            remaining_total = max(0, remaining_total - len(group_items))

    # Applications
    app_items: list[Any] = []
    app_nodes: list[ResourceNode] = []
    if request.entra_include_applications and (not remaining_total or remaining_total > 0):
        app_items = await get_all_pages(
            client,
            client.applications,
            operation_name="entra.applications",
            max_objects=_cap_for_collection(request.entra_max_objects),
            throttle_delay_seconds=request.throttle_delay_seconds,
        )
        app_nodes = [application_to_node(a) for a in app_items]
        _extend_unique_nodes(nodes, app_nodes)
        if remaining_total:
            remaining_total = max(0, remaining_total - len(app_items))

    # Service principals (needed for appId correlation edges; guarded by SP expansion caps)
    sp_items: list[Any] = []
    if (
        request.entra_sp_ownership_max_sps
        and request.entra_sp_ownership_max_owners_per_sp
        and (not remaining_total or remaining_total > 0)
    ):
        sp_items = await get_all_pages(
            client,
            client.service_principals,
            operation_name="entra.service_principals",
            max_objects=_cap_for_collection(request.entra_max_objects),
            throttle_delay_seconds=request.throttle_delay_seconds,
        )
        _extend_unique_nodes(nodes, (service_principal_to_node(sp) for sp in sp_items))
        if remaining_total:
            remaining_total = max(0, remaining_total - len(sp_items))

    # Conditional access policies
    if request.entra_include_conditional_access_policies and (not remaining_total or remaining_total > 0):
        policy_items = await get_all_pages(
            client,
            client.identity.conditional_access.policies,
            operation_name="entra.conditional_access_policies",
            max_objects=_cap_for_collection(request.entra_max_objects),
            throttle_delay_seconds=request.throttle_delay_seconds,
        )
        _extend_unique_nodes(nodes, (conditional_access_policy_to_node(p) for p in policy_items))
        if remaining_total:
            remaining_total = max(0, remaining_total - len(policy_items))

    # Risky users
    if request.entra_include_risky_users and (not remaining_total or remaining_total > 0):
        risky_items = await get_all_pages(
            client,
            client.identity_protection.risky_users,
            operation_name="entra.risky_users",
            max_objects=_cap_for_collection(request.entra_max_objects),
            throttle_delay_seconds=request.throttle_delay_seconds,
        )
        _extend_unique_nodes(nodes, (risky_user_to_node(u) for u in risky_items))
        if remaining_total:
            remaining_total = max(0, remaining_total - len(risky_items))

    # Relationship expansions (bounded)
    if request.include_relationships:
        # 1) Group membership expansion
        if request.entra_group_membership_max_groups and request.entra_group_membership_max_members_per_group:
            for g in group_items[: request.entra_group_membership_max_groups]:
                group_id = _safe_str(getattr(g, "id", None))
                if not group_id:
                    continue
                members = await get_all_pages(
                    client,
                    client.groups.by_group_id(group_id).members,
                    operation_name="entra.group_members",
                    max_objects=request.entra_group_membership_max_members_per_group,
                    throttle_delay_seconds=request.throttle_delay_seconds,
                )
                member_nodes = [_directory_object_to_node(m) for m in members if getattr(m, "id", None)]
                _extend_unique_nodes(nodes, member_nodes)
                for mn in member_nodes:
                    relationships.append(
                        ResourceRelationship(
                            source_id=_graph_id("group", group_id),
                            target_id=mn.id,
                            relation_type="has_member",
                            weight=1.0,
                        )
                    )

        # 2) Application ownership expansion
        if request.entra_ownership_max_apps and request.entra_ownership_max_owners_per_app:
            for a in app_items[: request.entra_ownership_max_apps]:
                app_object_id = _safe_str(getattr(a, "id", None))
                if not app_object_id:
                    continue
                owners = await get_all_pages(
                    client,
                    client.applications.by_application_id(app_object_id).owners,
                    operation_name="entra.application_owners",
                    max_objects=request.entra_ownership_max_owners_per_app,
                    throttle_delay_seconds=request.throttle_delay_seconds,
                )
                owner_nodes = [_directory_object_to_node(o) for o in owners if getattr(o, "id", None)]
                _extend_unique_nodes(nodes, owner_nodes)
                for on in owner_nodes:
                    relationships.append(
                        ResourceRelationship(
                            source_id=_graph_id("application", app_object_id),
                            target_id=on.id,
                            relation_type="has_owner",
                            weight=1.0,
                        )
                    )

        # 3) Service principal ownership expansion
        if request.entra_sp_ownership_max_sps and request.entra_sp_ownership_max_owners_per_sp:
            for sp in sp_items[: request.entra_sp_ownership_max_sps]:
                sp_object_id = _safe_str(getattr(sp, "id", None))
                if not sp_object_id:
                    continue
                owners = await get_all_pages(
                    client,
                    client.service_principals.by_service_principal_id(sp_object_id).owners,
                    operation_name="entra.service_principal_owners",
                    max_objects=request.entra_sp_ownership_max_owners_per_sp,
                    throttle_delay_seconds=request.throttle_delay_seconds,
                )
                owner_nodes = [_directory_object_to_node(o) for o in owners if getattr(o, "id", None)]
                _extend_unique_nodes(nodes, owner_nodes)
                for on in owner_nodes:
                    relationships.append(
                        ResourceRelationship(
                            source_id=_graph_id("servicePrincipal", sp_object_id),
                            target_id=on.id,
                            relation_type="has_owner",
                            weight=1.0,
                        )
                    )

        # Minimal inferred relationships: appId correlation between service principals and applications.
        relationships.extend(build_graph_appid_edges(nodes))

    LOGGER.info(
        "Entra enumeration completed",
        extra={
            "context": {
                "tenant_id": request.tenant_id,
                "nodes": len(nodes),
                "relationships": len(relationships),
            }
        },
    )

    return nodes, relationships
