"""Discovery orchestrator."""

from __future__ import annotations

import time

from .adt_types import (
    AzureDiscoveryRequest,
    AzureDiscoveryResponse,
    DiscoveryMergeMetrics,
    DiscoveryMetrics,
    DiscoveryPhaseMetrics,
    ResourceNode,
    ResourceRelationship,
)
from .enumerators import enumerate_azure_resources, enumerate_entra_resources
from .enumerators.rbac import RBACEnumerator
from .enumerators.pim import PIMEnumerator
from .enumerators.defender import DefenderEnumerator
from .reporting import emit_console_summary, render_visualization
from .utils import build_environment_config, get_credential
from .utils.graph_helpers import (
    ensure_unique_nodes,
    ensure_unique_relationships,
    filter_relationships_to_known_nodes,
)
from .utils.rbac_helpers import build_rbac_relationships
from .utils.pim_helpers import build_pim_relationships
from .utils.defender_helpers import build_defender_relationships


def _stub_node_for_relationship_endpoint(node_id: str) -> ResourceNode | None:
    if not node_id.startswith("graph://"):
        return None

    # Best-effort parsing: graph://<kind>/<objectId>
    remainder = node_id.removeprefix("graph://")
    kind, _, object_id = remainder.partition("/")
    name = object_id or node_id

    tags = {}
    if object_id:
        tags["graph_id"] = object_id

    return ResourceNode(
        id=node_id,
        name=name,
        type="Microsoft.Graph/DirectoryObject" if kind else "Microsoft.Graph/DirectoryObject",
        subscription_id="Tenant",
        tags=tags,
    )


def _materialize_missing_endpoints(
    nodes: list[ResourceNode],
    relationships: list[ResourceRelationship],
) -> list[ResourceNode]:
    node_ids = {n.id for n in nodes}
    stubs: list[ResourceNode] = []
    for r in relationships:
        for endpoint in (r.source_id, r.target_id):
            if endpoint in node_ids:
                continue
            stub = _stub_node_for_relationship_endpoint(endpoint)
            if not stub:
                continue
            node_ids.add(stub.id)
            stubs.append(stub)

    if not stubs:
        return nodes
    return ensure_unique_nodes([*nodes, *stubs])


def _merge_results(
    base: AzureDiscoveryResponse,
    nodes: list[ResourceNode],
    relationships: list[ResourceRelationship],
    materialize_missing_endpoints: bool,
) -> tuple[AzureDiscoveryResponse, DiscoveryMergeMetrics]:
    all_nodes = ensure_unique_nodes([*base.nodes, *nodes])

    combined_relationships = [*base.relationships, *relationships]

    if materialize_missing_endpoints:
        all_nodes = _materialize_missing_endpoints(all_nodes, combined_relationships)

    filtered_relationships = filter_relationships_to_known_nodes(
        all_nodes,
        combined_relationships,
    )
    dropped_dangling = len(combined_relationships) - len(filtered_relationships)

    all_relationships = ensure_unique_relationships(filtered_relationships)
    deduped_edges = len(filtered_relationships) - len(all_relationships)

    merged = base.model_copy(
        update={
            "nodes": all_nodes,
            "relationships": all_relationships,
            "total_resources": len(all_nodes),
        }
    )

    return merged, DiscoveryMergeMetrics(
        dropped_dangling_relationships=dropped_dangling,
        deduped_relationships=deduped_edges,
    )


async def run_discovery(request: AzureDiscoveryRequest) -> AzureDiscoveryResponse:
    """Coordinate enumeration and visualization."""

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

    phases: list[DiscoveryPhaseMetrics] = []
    merge_metrics = DiscoveryMergeMetrics()

    environment_config = build_environment_config(request.environment)
    credential = get_credential(request, environment_config)

    t0 = time.perf_counter()
    discovery_response = await enumerate_azure_resources(request, credential=credential)
    phases.append(
        DiscoveryPhaseMetrics(
            name="azure.arm",
            elapsed_seconds=time.perf_counter() - t0,
            nodes_added=len(discovery_response.nodes),
            relationships_added=len(discovery_response.relationships),
        )
    )

    if request.include_entra:
        pre_nodes = len(discovery_response.nodes)
        pre_relationships = len(discovery_response.relationships)

        t0 = time.perf_counter()
        entra_nodes, entra_relationships = await enumerate_entra_resources(
            request,
            credential=credential,
        )
        if entra_nodes or entra_relationships:
            discovery_response, metrics = _merge_results(
                discovery_response,
                entra_nodes,
                entra_relationships,
                materialize_missing_endpoints=request.materialize_missing_endpoints,
            )
            merge_metrics = DiscoveryMergeMetrics(
                dropped_dangling_relationships=
                merge_metrics.dropped_dangling_relationships + metrics.dropped_dangling_relationships,
                deduped_relationships=merge_metrics.deduped_relationships + metrics.deduped_relationships,
            )

        phases.append(
            DiscoveryPhaseMetrics(
                name="entra.graph",
                elapsed_seconds=time.perf_counter() - t0,
                nodes_added=max(0, len(discovery_response.nodes) - pre_nodes),
                relationships_added=max(0, len(discovery_response.relationships) - pre_relationships),
            )
        )

    # Phase 3 (optional): RBAC enumeration
    if request.include_rbac_assignments or request.include_rbac_definitions:
        pre_nodes = len(discovery_response.nodes)
        pre_relationships = len(discovery_response.relationships)

        t0 = time.perf_counter()
        rbac_enumerator = RBACEnumerator(credential, environment_config.resource_manager)

        rbac_nodes: list[ResourceNode] = []
        rbac_relationships: list[ResourceRelationship] = []

        if request.include_rbac_assignments:
            assignments = await rbac_enumerator.enumerate_role_assignments(
                subscription_ids=discovery_response.discovered_subscriptions,
                scope_filter=request.rbac_scope_filter,
            )
            rbac_nodes.extend(assignments)

            # Build relationships if Entra principals are available
            if request.include_entra and request.include_relationships:
                entra_principals = [
                    n
                    for n in discovery_response.nodes
                    if n.type.startswith("Microsoft.Graph/")
                ]
                rbac_rels = build_rbac_relationships(
                    assignments,
                    discovery_response.nodes,
                    entra_principals,
                )
                rbac_relationships.extend(rbac_rels)

        if request.include_rbac_definitions:
            definitions = await rbac_enumerator.enumerate_role_definitions(
                subscription_ids=discovery_response.discovered_subscriptions,
            )
            rbac_nodes.extend(definitions)

        if rbac_nodes or rbac_relationships:
            discovery_response, metrics = _merge_results(
                discovery_response,
                rbac_nodes,
                rbac_relationships,
                materialize_missing_endpoints=request.materialize_missing_endpoints,
            )
            merge_metrics = DiscoveryMergeMetrics(
                dropped_dangling_relationships=(
                    merge_metrics.dropped_dangling_relationships
                    + metrics.dropped_dangling_relationships
                ),
                deduped_relationships=(
                    merge_metrics.deduped_relationships + metrics.deduped_relationships
                ),
            )

        phases.append(
            DiscoveryPhaseMetrics(
                name="rbac.enumeration",
                elapsed_seconds=time.perf_counter() - t0,
                nodes_added=max(0, len(discovery_response.nodes) - pre_nodes),
                relationships_added=max(
                    0, len(discovery_response.relationships) - pre_relationships
                ),
            )
        )

    # Phase 3.5 (optional): PIM (Privileged Identity Management) enumeration
    if request.include_pim:
        pre_nodes = len(discovery_response.nodes)
        pre_relationships = len(discovery_response.relationships)

        t0 = time.perf_counter()
        pim_enumerator = PIMEnumerator(
            credential,
            environment_config.graph_endpoint,
            environment_config.resource_manager,
        )

        pim_nodes: list[ResourceNode] = []
        pim_relationships: list[ResourceRelationship] = []

        # Enumerate Entra ID role eligibilities
        if request.pim_config.include_entra_role_eligibilities:
            entra_eligibilities = await pim_enumerator.enumerate_entra_role_eligibilities(
                include_role_eligibility_schedules=True,
                include_active_requests=request.pim_config.include_entra_role_eligibility_requests,
            )
            pim_nodes.extend(entra_eligibilities)

        # Enumerate Azure Resource role eligibilities
        if request.pim_config.include_azure_resource_eligibilities:
            resource_eligibilities = await pim_enumerator.enumerate_azure_resource_role_eligibilities(
                subscription_ids=discovery_response.discovered_subscriptions,
                scope_filter=request.pim_scope_filter,
            )
            pim_nodes.extend(resource_eligibilities)

        # Build relationships if Entra principals and resources are available
        if request.include_relationships and pim_nodes:
            entra_principals = [
                n
                for n in discovery_response.nodes
                if n.type.startswith("Microsoft.Graph/")
            ]
            role_definitions = [
                n
                for n in discovery_response.nodes
                if n.type == "Microsoft.Authorization/roleDefinitions"
            ]
            pim_rels = build_pim_relationships(
                pim_nodes,
                discovery_response.nodes,
                entra_principals,
                role_definitions,
            )
            pim_relationships.extend(pim_rels)

        if pim_nodes or pim_relationships:
            discovery_response, metrics = _merge_results(
                discovery_response,
                pim_nodes,
                pim_relationships,
                materialize_missing_endpoints=request.materialize_missing_endpoints,
            )
            merge_metrics = DiscoveryMergeMetrics(
                dropped_dangling_relationships=(
                    merge_metrics.dropped_dangling_relationships
                    + metrics.dropped_dangling_relationships
                ),
                deduped_relationships=(
                    merge_metrics.deduped_relationships + metrics.deduped_relationships
                ),
            )

        phases.append(
            DiscoveryPhaseMetrics(
                name="pim.enumeration",
                elapsed_seconds=time.perf_counter() - t0,
                nodes_added=max(0, len(discovery_response.nodes) - pre_nodes),
                relationships_added=max(
                    0, len(discovery_response.relationships) - pre_relationships
                ),
            )
        )

    # Phase 4 (optional): Defender for Cloud enumeration
    if request.include_defender_cloud:
        pre_nodes = len(discovery_response.nodes)
        pre_relationships = len(discovery_response.relationships)

        t0 = time.perf_counter()
        defender_enumerator = DefenderEnumerator(
            credential, environment_config.resource_manager
        )

        defender_nodes: list[ResourceNode] = []
        defender_relationships: list[ResourceRelationship] = []

        # Enumerate security alerts
        if request.defender_config.include_security_alerts:
            alerts = await defender_enumerator.enumerate_security_alerts(
                subscription_ids=discovery_response.discovered_subscriptions,
                severity_filter=request.defender_config.alert_severity_filter,
                status_filter=request.defender_config.alert_status_filter,
            )
            defender_nodes.extend(alerts)

        # Enumerate security assessments
        if request.defender_config.include_security_assessments:
            assessments = await defender_enumerator.enumerate_security_assessments(
                subscription_ids=discovery_response.discovered_subscriptions,
                severity_filter=request.defender_config.assessment_severity_filter,
                status_filter=request.defender_config.assessment_status_filter,
            )
            defender_nodes.extend(assessments)

        # Enumerate secure scores
        if request.defender_config.include_secure_scores:
            scores = await defender_enumerator.enumerate_secure_scores(
                subscription_ids=discovery_response.discovered_subscriptions,
            )
            defender_nodes.extend(scores)

        # Build relationships from security findings to affected resources
        if request.include_relationships and defender_nodes:
            alerts_for_relationships = [
                n for n in defender_nodes if n.type == "Microsoft.Security/alerts"
            ]
            assessments_for_relationships = [
                n for n in defender_nodes if n.type == "Microsoft.Security/assessments"
            ]

            defender_rels = build_defender_relationships(
                alerts=alerts_for_relationships,
                assessments=assessments_for_relationships,
                all_nodes=discovery_response.nodes,
                materialize_missing=request.materialize_missing_endpoints,
            )
            defender_relationships.extend(defender_rels)

        if defender_nodes or defender_relationships:
            discovery_response, metrics = _merge_results(
                discovery_response,
                defender_nodes,
                defender_relationships,
                materialize_missing_endpoints=request.materialize_missing_endpoints,
            )
            merge_metrics = DiscoveryMergeMetrics(
                dropped_dangling_relationships=(
                    merge_metrics.dropped_dangling_relationships
                    + metrics.dropped_dangling_relationships
                ),
                deduped_relationships=(
                    merge_metrics.deduped_relationships + metrics.deduped_relationships
                ),
            )

        phases.append(
            DiscoveryPhaseMetrics(
                name="defender.enumeration",
                elapsed_seconds=time.perf_counter() - t0,
                nodes_added=max(0, len(discovery_response.nodes) - pre_nodes),
                relationships_added=max(
                    0, len(discovery_response.relationships) - pre_relationships
                ),
            )
        )

    t0 = time.perf_counter()
    visualization = render_visualization(
        response=discovery_response,
        options=request.visualization,
    )
    phases.append(
        DiscoveryPhaseMetrics(
            name="report.render_html",
            elapsed_seconds=time.perf_counter() - t0,
        )
    )

    enriched_response = discovery_response.model_copy(
        update={"html_report_path": visualization.html_path}
    )

    if request.include_metrics:
        enriched_response = enriched_response.model_copy(
            update={"metrics": DiscoveryMetrics(phases=phases, merge=merge_metrics)}
        )

    emit_console_summary(enriched_response)
    return enriched_response
