import copy
from collections.abc import Mapping, Sequence
from datetime import datetime

from loguru import logger

from exo.shared.types.common import NodeId
from exo.shared.types.events import (
    ChunkGenerated,
    Event,
    IndexedEvent,
    InstanceCreated,
    InstanceDeleted,
    NodeDownloadProgress,
    NodeGatheredInfo,
    NodeTimedOut,
    RunnerStatusUpdated,
    TaskAcknowledged,
    TaskCreated,
    TaskDeleted,
    TaskFailed,
    TaskStatusUpdated,
    TestEvent,
    TopologyEdgeCreated,
    TopologyEdgeDeleted,
    TracesCollected,
    TracesMerged,
)
from exo.shared.types.profiling import (
    NodeIdentity,
    NodeNetworkInfo,
    NodeRdmaCtlStatus,
    NodeThunderboltInfo,
    ThunderboltBridgeStatus,
)
from exo.shared.types.state import State
from exo.shared.types.tasks import Task, TaskId, TaskStatus
from exo.shared.types.worker.downloads import DownloadProgress
from exo.shared.types.worker.instances import Instance, InstanceId
from exo.shared.types.worker.runners import (
    RunnerId,
    RunnerReady,
    RunnerShutdown,
    RunnerStatus,
)
from exo.utils.info_gatherer.info_gatherer import (
    MacmonMetrics,
    MacThunderboltConnections,
    MacThunderboltIdentifiers,
    MemoryUsage,
    MiscData,
    NodeBackends,
    NodeConfig,
    NodeDiskUsage,
    NodeNetworkInterfaces,
    RdmaCtlStatus,
    StaticNodeInformation,
    ThunderboltBridgeInfo,
)


def event_apply(event: Event, state: State) -> State:
    """Apply an event to state."""
    match event:
        case (
            TestEvent()
            | ChunkGenerated()
            | TaskAcknowledged()
            | TracesCollected()
            | TracesMerged()
        ):  # Pass-through events that don't modify state
            return state
        case InstanceCreated():
            return apply_instance_created(event, state)
        case InstanceDeleted():
            return apply_instance_deleted(event, state)
        case NodeTimedOut():
            return apply_node_timed_out(event, state)
        case NodeDownloadProgress():
            return apply_node_download_progress(event, state)
        case NodeGatheredInfo():
            return apply_node_gathered_info(event, state)
        case RunnerStatusUpdated():
            return apply_runner_status_updated(event, state)
        case TaskCreated():
            return apply_task_created(event, state)
        case TaskDeleted():
            return apply_task_deleted(event, state)
        case TaskFailed():
            return apply_task_failed(event, state)
        case TaskStatusUpdated():
            return apply_task_status_updated(event, state)
        case TopologyEdgeCreated():
            return apply_topology_edge_created(event, state)
        case TopologyEdgeDeleted():
            return apply_topology_edge_deleted(event, state)


def apply(state: State, event: IndexedEvent) -> State:
    # Just to test that events are only applied in correct order
    if state.last_event_applied_idx != event.idx - 1:
        logger.warning(
            f"Expected event {state.last_event_applied_idx + 1} but received {event.idx}"
        )
    assert state.last_event_applied_idx == event.idx - 1
    new_state: State = event_apply(event.event, state)
    return new_state.model_copy(update={"last_event_applied_idx": event.idx})


def apply_node_download_progress(event: NodeDownloadProgress, state: State) -> State:
    """
    Update or add a node download progress to state.
    """
    dp = event.download_progress
    node_id = dp.node_id

    current = list(state.downloads.get(node_id, ()))

    replaced = False
    for i, existing_dp in enumerate(current):
        # TODO(ciaran): deduplicate by model_id for now. Will need to use
        # shard_metadata again when pipeline and tensor downloads differ.
        # For now this is fine
        if (
            existing_dp.shard_metadata.model_card.model_id
            == dp.shard_metadata.model_card.model_id
        ):
            current[i] = dp
            replaced = True
            break

    if not replaced:
        current.append(dp)

    new_downloads: Mapping[NodeId, Sequence[DownloadProgress]] = {
        **state.downloads,
        node_id: current,
    }
    return state.model_copy(update={"downloads": new_downloads})


def apply_task_created(event: TaskCreated, state: State) -> State:
    new_tasks: Mapping[TaskId, Task] = {**state.tasks, event.task_id: event.task}
    return state.model_copy(update={"tasks": new_tasks})


def apply_task_deleted(event: TaskDeleted, state: State) -> State:
    new_tasks: Mapping[TaskId, Task] = {
        tid: task for tid, task in state.tasks.items() if tid != event.task_id
    }
    return state.model_copy(update={"tasks": new_tasks})


def apply_task_status_updated(event: TaskStatusUpdated, state: State) -> State:
    if event.task_id not in state.tasks:
        # maybe should raise
        return state

    update: dict[str, TaskStatus | None] = {
        "task_status": event.task_status,
    }
    if event.task_status != TaskStatus.Failed:
        update["error_type"] = None
        update["error_message"] = None

    updated_task = state.tasks[event.task_id].model_copy(update=update)
    new_tasks: Mapping[TaskId, Task] = {**state.tasks, event.task_id: updated_task}
    return state.model_copy(update={"tasks": new_tasks})


def apply_task_failed(event: TaskFailed, state: State) -> State:
    if event.task_id not in state.tasks:
        # maybe should raise
        return state

    updated_task = state.tasks[event.task_id].model_copy(
        update={"error_type": event.error_type, "error_message": event.error_message}
    )
    new_tasks: Mapping[TaskId, Task] = {**state.tasks, event.task_id: updated_task}
    return state.model_copy(update={"tasks": new_tasks})


def apply_instance_created(event: InstanceCreated, state: State) -> State:
    instance = event.instance
    new_instances: Mapping[InstanceId, Instance] = {
        **state.instances,
        instance.instance_id: instance,
    }
    return state.model_copy(update={"instances": new_instances})


def apply_instance_deleted(event: InstanceDeleted, state: State) -> State:
    new_instances: Mapping[InstanceId, Instance] = {
        iid: inst for iid, inst in state.instances.items() if iid != event.instance_id
    }
    return state.model_copy(update={"instances": new_instances})


def apply_runner_status_updated(event: RunnerStatusUpdated, state: State) -> State:
    if isinstance(event.runner_status, RunnerShutdown):
        new_runners: Mapping[RunnerId, RunnerStatus] = {
            rid: rs for rid, rs in state.runners.items() if rid != event.runner_id
        }
        new_ports: Mapping[RunnerId, int] = {
            rid: p
            for rid, p in state.prefill_server_ports.items()
            if rid != event.runner_id
        }
        return state.model_copy(
            update={"runners": new_runners, "prefill_server_ports": new_ports}
        )
    new_runners = {
        **state.runners,
        event.runner_id: event.runner_status,
    }
    update: dict[str, object] = {"runners": new_runners}
    if (
        isinstance(event.runner_status, RunnerReady)
        and event.runner_status.prefill_server_port is not None
    ):
        update["prefill_server_ports"] = {
            **state.prefill_server_ports,
            event.runner_id: event.runner_status.prefill_server_port,
        }
    return state.model_copy(update=update)


def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
    topology = copy.deepcopy(state.topology)
    topology.remove_node(event.node_id)
    last_seen = {
        key: value for key, value in state.last_seen.items() if key != event.node_id
    }
    downloads = {
        key: value for key, value in state.downloads.items() if key != event.node_id
    }
    # Clean up all granular node mappings
    node_memory = {
        key: value for key, value in state.node_memory.items() if key != event.node_id
    }
    node_disk = {
        key: value for key, value in state.node_disk.items() if key != event.node_id
    }
    node_system = {
        key: value for key, value in state.node_system.items() if key != event.node_id
    }
    node_network = {
        key: value for key, value in state.node_network.items() if key != event.node_id
    }
    node_thunderbolt = {
        key: value
        for key, value in state.node_thunderbolt.items()
        if key != event.node_id
    }
    node_thunderbolt_bridge = {
        key: value
        for key, value in state.node_thunderbolt_bridge.items()
        if key != event.node_id
    }
    node_rdma_ctl = {
        key: value for key, value in state.node_rdma_ctl.items() if key != event.node_id
    }
    # Only recompute cycles if the leaving node had TB bridge enabled
    leaving_node_status = state.node_thunderbolt_bridge.get(event.node_id)
    leaving_node_had_tb_enabled = (
        leaving_node_status is not None and leaving_node_status.enabled
    )
    thunderbolt_bridge_cycles = (
        topology.get_thunderbolt_bridge_cycles(node_thunderbolt_bridge, node_network)
        if leaving_node_had_tb_enabled
        else [list(cycle) for cycle in state.thunderbolt_bridge_cycles]
    )
    return state.model_copy(
        update={
            "downloads": downloads,
            "topology": topology,
            "last_seen": last_seen,
            "node_memory": node_memory,
            "node_disk": node_disk,
            "node_system": node_system,
            "node_network": node_network,
            "node_thunderbolt": node_thunderbolt,
            "node_thunderbolt_bridge": node_thunderbolt_bridge,
            "node_rdma_ctl": node_rdma_ctl,
            "thunderbolt_bridge_cycles": thunderbolt_bridge_cycles,
        }
    )


def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
    topology = copy.deepcopy(state.topology)
    topology.add_node(event.node_id)
    info = event.info

    # Build update dict with only the mappings that change
    update: dict[str, object] = {
        "last_seen": {
            **state.last_seen,
            event.node_id: datetime.fromisoformat(event.when),
        },
        "topology": topology,
    }

    match info:
        case MacmonMetrics():
            update["node_system"] = {
                **state.node_system,
                event.node_id: info.system_profile,
            }
            update["node_memory"] = {**state.node_memory, event.node_id: info.memory}
        case MemoryUsage():
            update["node_memory"] = {**state.node_memory, event.node_id: info}
        case NodeDiskUsage():
            update["node_disk"] = {**state.node_disk, event.node_id: info.disk_usage}
        case NodeConfig():
            pass
        case MiscData():
            current_identity = state.node_identities.get(event.node_id, NodeIdentity())
            new_identity = current_identity.model_copy(
                update={"friendly_name": info.friendly_name}
            )
            update["node_identities"] = {
                **state.node_identities,
                event.node_id: new_identity,
            }
        case StaticNodeInformation():
            current_identity = state.node_identities.get(event.node_id, NodeIdentity())
            new_identity = current_identity.model_copy(
                update={
                    "model_id": info.model,
                    "chip_id": info.chip,
                    "os_version": info.os_version,
                    "os_build_version": info.os_build_version,
                }
            )
            update["node_identities"] = {
                **state.node_identities,
                event.node_id: new_identity,
            }
        case NodeNetworkInterfaces():
            update["node_network"] = {
                **state.node_network,
                event.node_id: NodeNetworkInfo(interfaces=info.ifaces),
            }
        case MacThunderboltIdentifiers():
            update["node_thunderbolt"] = {
                **state.node_thunderbolt,
                event.node_id: NodeThunderboltInfo(interfaces=info.idents),
            }
        case MacThunderboltConnections():
            update["node_thunderbolt_connections"] = {
                **state.node_thunderbolt_connections,
                event.node_id: info,
            }
        case ThunderboltBridgeInfo():
            new_tb_bridge: dict[NodeId, ThunderboltBridgeStatus] = {
                **state.node_thunderbolt_bridge,
                event.node_id: info.status,
            }
            update["node_thunderbolt_bridge"] = new_tb_bridge
            update["thunderbolt_bridge_cycles"] = (
                topology.get_thunderbolt_bridge_cycles(
                    new_tb_bridge, state.node_network
                )
            )
        case RdmaCtlStatus():
            update["node_rdma_ctl"] = {
                **state.node_rdma_ctl,
                event.node_id: NodeRdmaCtlStatus(enabled=info.enabled),
            }
        case NodeBackends():
            update["node_backends"] = {
                **state.node_backends,
                event.node_id: info.backends,
            }

    return state.model_copy(update=update)


def apply_topology_edge_created(event: TopologyEdgeCreated, state: State) -> State:
    source_connections = state.node_socket_connections.get(event.conn.source, {})
    sink_connections = source_connections.get(event.conn.sink, [])

    update = {
        "node_socket_connections": {
            **state.node_socket_connections,
            event.conn.source: {
                **source_connections,
                event.conn.sink: sink_connections
                if event.conn.edge in sink_connections
                else [*sink_connections, event.conn.edge],
            },
        }
    }
    return state.model_copy(update=update)


def apply_topology_edge_deleted(event: TopologyEdgeDeleted, state: State) -> State:
    inner_update = {
        sink: final_edges
        for sink, edges in state.node_socket_connections.get(
            event.conn.source, {}
        ).items()
        if (
            final_edges := [
                edge
                for edge in edges
                if (edge != event.conn.edge or sink != event.conn.sink)
            ]
        )
    }
    update = {
        "node_socket_connections": {
            source: maps
            for source, maps in {
                **state.node_socket_connections,
                event.conn.source: inner_update,
            }.items()
            if maps
        }
    }
    return state.model_copy(update=update)
