volli-manager 0.1.12

use super::{ManagerContext, mesh::PeerDialRequest};
use crate::{
    ConnectionState, ManagerPeerEntry, Message,
    keys::load_peers_for_gossip,
    keys::save_csk,
    mesh::should_dial_peer,
    peers::{AliveTable, set_peer_state, update_alive},
    util::short_id,
};
use eyre::{Report, eyre};
use std::sync::atomic::Ordering;
use tracing::{debug, info, warn};
use volli_core::token::{decode_token, verify_token};
use volli_transport::MessageTransportExt;
use volli_transport::Transport;

// Helper function to update self_meta with current health metrics
async fn update_self_meta_with_health(ctx: &ManagerContext) -> ManagerPeerEntry {
    let mut self_meta = ctx.state.self_meta.clone();

    // Collect current health metrics
    let health_metrics = ctx.health.collector.lock().await.collect_metrics().await;
    self_meta.health = Some(volli_core::HealthMetrics {
        health_score: health_metrics.health_score,
        load_percentage: health_metrics.load_percentage,
        max_workers: health_metrics.max_workers,
        current_workers: health_metrics.current_workers,
        avg_cpu: health_metrics.avg_cpu,
        avg_memory: health_metrics.avg_memory,
        last_health_update: health_metrics.last_health_update,
    });

    self_meta
}

// Helper function to clean up peer state on connection error
async fn cleanup_on_error(
    profile: &str,
    peers: &AliveTable,
    active_meta: &Option<ManagerPeerEntry>,
) {
    if let Some(meta) = active_meta {
        warn!(
            "Cleaning up failed connection to manager {} ({}:{})",
            meta.manager_name, meta.host, meta.quic_port
        );
        // Immediate lost peer log with before/after counts
        let before_active = {
            let tbl = peers.lock().await;
            tbl.values()
                .filter(|p| p.conn_state != ConnectionState::Inactive)
                .count()
        };
        set_peer_state(peers, meta, ConnectionState::Inactive).await;
        let after_active = {
            let tbl = peers.lock().await;
            tbl.values()
                .filter(|p| p.conn_state != ConnectionState::Inactive)
                .count()
        };
        if after_active < before_active {
            warn!(
                "Manager {} lost {} peer(s): {} -> {}",
                profile,
                before_active - after_active,
                before_active,
                after_active
            );
        }
    } else {
        warn!("Cleaning up connection with no active meta");
    }
}

// Helper function to process a peer update and return if gossip is needed
pub(crate) fn should_skip_peer(ctx: &ManagerContext, peer: &ManagerPeerEntry) -> bool {
    peer.manager_id == ctx.state.self_meta.manager_id
}

pub(crate) async fn process_peer_update(ctx: &ManagerContext, peer: ManagerPeerEntry) -> bool {
    if should_skip_peer(ctx, &peer) {
        return false;
    }

    let is_new = update_alive(
        &ctx.state.peers,
        &ctx.communication.alive_tx,
        &ctx.communication.profile,
        peer.clone(),
        &ctx.state.peer_version,
    )
    .await;

    if is_new {
        let _ = ctx.communication.dial_tx.send(PeerDialRequest {
            peer,
            gossip_hint: false,
        });
    }
    is_new
}

async fn send_initial_announcement(
    ctx: &ManagerContext,
    transport: &mut Box<dyn Transport>,
    peer_id: &str,
) -> Result<(), Report> {
    let last_version_sent = ctx.state.peer_version.load(Ordering::SeqCst);
    let peers_list = load_peers_for_gossip(&ctx.communication.profile).unwrap_or_default();

    transport
        .send(&Message::Announce {
            meta: Box::new(update_self_meta_with_health(ctx).await),
            version: last_version_sent,
            peers: peers_list,
            workers: Vec::new(),
        })
        .await
        .map_err(|e| {
            tracing::error!(target: "connection", peer=%peer_id, "Failed to send initial heartbeat: {}", e);
            e.into()
        })
}

// Unified connection loop for manager-to-manager communication
async fn handle_manager_connection_loop(
    ctx: &ManagerContext,
    mut transport: Box<dyn Transport>,
    peer_id: String,
    active_meta: &mut Option<ManagerPeerEntry>,
) -> Result<(), Report> {
    let mut rx = ctx.communication.alive_tx.subscribe();
    let mut last_version_sent = ctx.state.peer_version.load(Ordering::SeqCst);
    let mut hb_timer = tokio::time::interval(std::time::Duration::from_secs(30));

    loop {
        tokio::select! {
            msg = transport.recv() => {
                let msg_result = match msg {
                    Ok(m) => {
                        m
                    },
                    Err(e) => {
                        cleanup_on_error(&ctx.communication.profile, &ctx.state.peers, active_meta)
                            .await;
                        tracing::debug!(target: "connection", peer=%peer_id, "Transport receive error: {}", e);
                        return Err(e.into());
                    }
                };
                match msg_result {
                    Some(Message::Announce { meta, version: _, peers: peer_list, workers: _ }) => {

                        let meta = *meta;
                        let mut need_gossip = false;
                        debug!(target: "connection", peer=%peer_id, id=%short_id(&meta.manager_id), "heartbeat from {}", meta.manager_name);

                        // Process the announcing peer and set connection state
                        let changed = process_peer_update(ctx, meta.clone()).await;
                        if changed {
                            debug!("Peer state changed for {}: updating gossip", meta.manager_name);
                        }
                        need_gossip |= changed;

                        // Set connection state for the peer we're directly connected to
                        let current_state = {
                            let peers = ctx.state.peers.lock().await;
                            peers
                                .get(&format!(
                                    "{}:{}:{}",
                                    meta.tenant, meta.cluster, meta.manager_id
                                ))
                                .map(|p| p.conn_state)
                                .unwrap_or(ConnectionState::Inactive)
                        };

                        if current_state != ConnectionState::Client {
                            // If transitioning from Inactive to connected, log immediate gain
                            let before_active = {
                                let tbl = ctx.state.peers.lock().await;
                                tbl.values().filter(|p| p.conn_state != ConnectionState::Inactive).count()
                            };
                            set_peer_state(&ctx.state.peers, &meta, ConnectionState::Server).await;
                            if current_state == ConnectionState::Inactive {
                                let after_active = {
                                    let tbl = ctx.state.peers.lock().await;
                                    tbl.values().filter(|p| p.conn_state != ConnectionState::Inactive).count()
                                };
                                info!(
                                    "Manager {} gained {} peer(s): {} -> {}",
                                    ctx.communication.profile,
                                    1,
                                    before_active,
                                    after_active
                                );
                            }
                        }
                        *active_meta = Some(meta);

                        // Process all peers from the peer list (gossip data)
                        for peer in peer_list {
                            let changed = process_peer_update(ctx, peer.clone()).await;
                            need_gossip |= changed;

                            // Check for gossip hint acceleration
                            if should_dial_peer(&ctx.state.self_meta, &peer) {
                                let peer_key = format!("{}:{}:{}", peer.tenant, peer.cluster, peer.manager_id);
                                let is_dead = {
                                    let peers = ctx.state.peers.lock().await;
                                    peers.get(&peer_key).map(|p| p.dead).unwrap_or(false)
                                };

                                if is_dead {
                                    // Send gossip hint to accelerate resurrection attempt
                                    if let Err(e) = ctx.communication.dial_tx.send(crate::connection::mesh::PeerDialRequest {
                                        peer: peer.clone(),
                                        gossip_hint: true,
                                    }) {
                                        tracing::debug!("Failed to send gossip hint for peer {}: {}", peer.manager_id, e);
                                    } else {
                                        tracing::debug!(
                                            peer = %crate::util::PeerLog(&peer),
                                            "sent gossip hint to accelerate dead peer resurrection"
                                        );
                                    }
                                }
                            }
                        }
                        if need_gossip {
                            let ver = ctx.state.peer_version.load(Ordering::SeqCst);
                            ctx.communication.alive_tx.send(ver).ok();
                        }
                    }
                    Some(other) => {
                        tracing::warn!(target: "connection", peer=%peer_id, "Received unexpected message: {:?}", other);
                    }
                    None => {
                        cleanup_on_error(&ctx.communication.profile, &ctx.state.peers, active_meta)
                            .await;
                        info!(target: "connection", peer=%peer_id, "connection closed");
                        return Ok(());
                    }
                }
            }
            // Periodic health announce to keep metrics fresh across mesh
            _ = hb_timer.tick() => {
                // Send an Announce with only updated meta; keep version unchanged
                let meta = update_self_meta_with_health(ctx).await;
                if let Err(e) = transport
                    .send(&crate::Message::Announce {
                        meta: Box::new(meta),
                        version: last_version_sent,
                        peers: Vec::new(),
                        workers: Vec::new(),
                    })
                    .await
                {
                    cleanup_on_error(&ctx.communication.profile, &ctx.state.peers, active_meta)
                        .await;
                    tracing::error!(target: "connection", peer=%peer_id, "Failed to send periodic health announce: {}", e);
                    return Err(e.into());
                }
            }
            version_result = rx.recv() => {
                match version_result {
                    Ok(version) => {
                        if version > last_version_sent {
                            last_version_sent = version;
                            let peers_to_send = load_peers_for_gossip(&ctx.communication.profile).unwrap_or_default();
                            tracing::debug!(target: "connection", peer=%peer_id, "sending immediate update");
                            if let Err(e) = transport
                                .send(&Message::Announce {
                                    meta: Box::new(update_self_meta_with_health(ctx).await),
                                    version: last_version_sent,
                                    peers: peers_to_send,
                                    workers: Vec::new(),
                                })
                                .await
                            {
                                cleanup_on_error(
                                    &ctx.communication.profile,
                                    &ctx.state.peers,
                                    active_meta,
                                )
                                .await;
                                tracing::error!(target: "connection", peer=%peer_id, "Failed to send update: {}", e);
                                return Err(e.into());
                            }
                        }
                    }
                    Err(err) => {
                        // If we lagged behind on the broadcast channel, coalesce to the latest
                        // version and continue without tearing down the connection.
                        if let tokio::sync::broadcast::error::RecvError::Lagged(dropped) = err {
                            tracing::debug!(target: "connection", peer=%peer_id, dropped=%dropped, "broadcast lagged; coalescing updates");
                            // Try to send a single up-to-date announce
                            let current = ctx.state.peer_version.load(Ordering::SeqCst);
                            if current > last_version_sent {
                                last_version_sent = current;
                                let peers_to_send = load_peers_for_gossip(&ctx.communication.profile).unwrap_or_default();
                                if let Err(e) = transport
                                    .send(&Message::Announce {
                                        meta: Box::new(update_self_meta_with_health(ctx).await),
                                        version: last_version_sent,
                                        peers: peers_to_send,
                                        workers: Vec::new(),
                                    })
                                    .await
                                {
                                    cleanup_on_error(
                                        &ctx.communication.profile,
                                        &ctx.state.peers,
                                        active_meta,
                                    )
                                    .await;
                                    tracing::error!(target: "connection", peer=%peer_id, "Failed to send coalesced update: {}", e);
                                    return Err(e.into());
                                }
                            }
                            // Continue loop; keep connection alive
                        } else {
                            // Channel closed (server shutting down); end gracefully
                            cleanup_on_error(
                                &ctx.communication.profile,
                                &ctx.state.peers,
                                active_meta,
                            )
                            .await;
                            return Ok(());
                        }
                    }
                }
            }
        }
    }
}

pub(crate) async fn handle_manager_as_server(
    ctx: &ManagerContext,
    mut transport: Box<dyn Transport>,
    peer: std::net::SocketAddr,
) -> Result<(), Report> {
    tracing::debug!("Starting manager client handler for peer: {}", peer);

    let csk = *ctx.security.csk.read().await;
    let csk_ver = ctx.security.csk_ver.load(Ordering::SeqCst);
    let whitelist = ctx.network.manager_nets.clone();
    let mut active_meta: Option<ManagerPeerEntry> = None;

    // Handle different connection types explicitly
    match transport.recv().await? {
        Some(Message::Auth { token, .. }) => {
            // Regular mesh connection - authenticate and send ClusterKey
            let token_dec = match decode_token(&token) {
                Ok(t) => t,
                Err(e) => {
                    tracing::error!("Auth token decode failed for peer {}: {}", peer, e);
                    transport.send(&Message::AuthErr).await.ok();
                    return Ok(());
                }
            };

            if token_dec.payload.tenant != ctx.state.self_meta.tenant
                || token_dec.payload.cluster != ctx.state.self_meta.cluster
            {
                tracing::warn!(
                    %peer,
                    token_tenant = %token_dec.payload.tenant,
                    token_cluster = %token_dec.payload.cluster,
                    expected_tenant = %ctx.state.self_meta.tenant,
                    expected_cluster = %ctx.state.self_meta.cluster,
                    "auth token tenant/cluster mismatch"
                );
                transport.send(&Message::AuthErr).await.ok();
                return Ok(());
            }

            // Check whitelist
            if !whitelist.is_empty() {
                let peer_addr = peer.ip();

                if !whitelist.iter().any(|net| net.contains(&peer_addr)) {
                    tracing::warn!(%peer, "mesh connection from non-whitelisted address");
                    transport.send(&Message::AuthErr).await.ok();
                    return Ok(());
                }
            }

            if let Err(e) = verify_token(&token_dec, &csk) {
                tracing::error!("Auth token verification failed for peer {}: {}", peer, e);
                transport.send(&Message::AuthErr).await.ok();
                return Err(e);
            }

            info!(target: "connection", %peer, "mesh connection established");
            transport.send(&Message::AuthOk).await?;

            if let Err(e) = transport
                .send(&Message::ClusterKey { ver: csk_ver, csk })
                .await
            {
                tracing::error!("Failed to send cluster key to {}: {}", peer, e);
                return Err(e.into());
            }
        }
        Some(Message::Join { token }) => {
            return crate::connection::join::join_as_server(ctx, transport, peer, token).await;
        }
        Some(Message::ClientAuth { token }) => {
            // External client connection - authenticate and handle commands
            return handle_client_connection(ctx, transport, peer, token).await;
        }
        _ => {
            tracing::error!("Invalid initial message from peer: {}", peer);
            transport.send(&Message::AuthErr).await.ok();
            return Ok(());
        }
    }
    // Send initial heartbeat
    send_initial_announcement(ctx, &mut transport, &peer.to_string()).await?;
    tracing::debug!(
        "Initial heartbeat sent successfully to manager client: {}",
        peer
    );

    handle_manager_connection_loop(ctx, transport, peer.to_string(), &mut active_meta).await
}

pub(crate) async fn handle_manager_as_client(
    ctx: ManagerContext,
    mut transport: Box<dyn Transport>,
    peer: String,
) -> Result<(), Report> {
    let mut self_meta = ctx.state.self_meta.clone();

    // Handle authentication as client
    transport
        .send(&Message::Auth {
            token: ctx.worker_config().unwrap().token.clone(),
            worker_id: None,
            worker_name: None,
        })
        .await?;
    match transport.recv().await? {
        Some(Message::AuthOk) => {
            tracing::info!(target: "connection", %peer, "connection established");
            // Expect cluster key immediately after authentication
            match transport.recv().await? {
                Some(Message::ClusterKey { ver, csk }) => {
                    if let Err(e) = save_csk(&ctx.communication.profile, &csk, ver) {
                        tracing::error!("Failed to save cluster key: {}", e);
                    }
                    {
                        let mut guard = ctx.security.csk.write().await;
                        *guard = csk;
                    }
                    ctx.security.csk_ver.store(ver, Ordering::SeqCst);
                    self_meta.csk_ver = ver;
                }

                _ => {
                    tracing::error!("Missing cluster key from manager {}", peer);
                    return Err(eyre!("missing cluster key"));
                }
            }
        }
        _other => {
            return Err(eyre!("authentication failed"));
        }
    }

    let mut updated_ctx = ctx.clone();
    updated_ctx.state.self_meta = self_meta;

    send_initial_announcement(&updated_ctx, &mut transport, &peer).await?;
    tracing::debug!(target: "connection", %peer, "Initial heartbeat sent successfully");

    let mut active_meta = None;
    handle_manager_connection_loop(&updated_ctx, transport, peer, &mut active_meta).await
}

/// Handle external client connections for command execution
async fn handle_client_connection(
    ctx: &ManagerContext,
    mut transport: Box<dyn Transport>,
    peer: std::net::SocketAddr,
    token: String,
) -> Result<(), Report> {
    info!(target: "connection", %peer, "handling client connection");

    // Authenticate client using cluster key
    let token_dec = match decode_token(&token) {
        Ok(t) => t,
        Err(e) => {
            tracing::error!("Client auth token decode failed for {}: {}", peer, e);
            transport.send(&Message::AuthErr).await.ok();
            return Ok(());
        }
    };

    // Verify tenant/cluster match
    if token_dec.payload.tenant != ctx.state.self_meta.tenant
        || token_dec.payload.cluster != ctx.state.self_meta.cluster
    {
        tracing::warn!(
            %peer,
            token_tenant = %token_dec.payload.tenant,
            token_cluster = %token_dec.payload.cluster,
            expected_tenant = %ctx.state.self_meta.tenant,
            expected_cluster = %ctx.state.self_meta.cluster,
            "client auth token tenant/cluster mismatch"
        );
        transport.send(&Message::AuthErr).await.ok();
        return Ok(());
    }

    // Verify token with cluster key
    let csk = *ctx.security.csk.read().await;
    if let Err(e) = verify_token(&token_dec, &csk) {
        tracing::error!("Client auth token verification failed for {}: {}", peer, e);
        transport.send(&Message::AuthErr).await.ok();
        return Err(e);
    }

    info!(target: "connection", %peer, "client authenticated successfully");
    transport.send(&Message::AuthOk).await?;

    // Handle client command requests
    loop {
        match transport.recv().await? {
            Some(Message::ClientCommandRequest {
                request_id,
                command,
                args,
                target,
                timeout_secs,
                options,
            }) => {
                info!(
                    target: "connection",
                    %peer,
                    %request_id,
                    %command,
                    "received client command request"
                );

                // Handle the command execution
                handle_client_command(
                    ctx,
                    &mut transport,
                    ClientCommandReq {
                        peer,
                        request_id,
                        command,
                        args,
                        target,
                        timeout_secs,
                        options,
                    },
                )
                .await?;
            }
            Some(Message::ClientCommandCancel { request_id }) => {
                info!(target: "connection", %peer, %request_id, "received client command cancel");
                if let Some(distributor) = &ctx.state.command_distributor {
                    let _ = distributor
                        .cancel_command(&request_id, "client cancelled")
                        .await;
                }
            }
            Some(_) => {
                warn!(target: "connection", %peer, "unexpected message from client");
            }
            None => {
                info!(target: "connection", %peer, "client disconnected");
                if let Some(distributor) = &ctx.state.command_distributor {
                    let _ = distributor
                        .cancel_commands_by_initiator(
                            &format!("client:{}", peer),
                            "client disconnected",
                        )
                        .await;
                }
                break;
            }
        }
    }

    Ok(())
}

/// Handle a single command request from a client
struct ClientCommandReq {
    peer: std::net::SocketAddr,
    request_id: String,
    command: String,
    args: Vec<String>,
    target: String,
    #[allow(dead_code)]
    timeout_secs: u64,
    options: String,
}

async fn handle_client_command(
    ctx: &ManagerContext,
    transport: &mut Box<dyn Transport>,
    req: ClientCommandReq,
) -> Result<(), Report> {
    use volli_commands::{CommandOptions, CommandTarget};

    // Deserialize target and options
    let command_target: CommandTarget = match serde_json::from_str(&req.target) {
        Ok(t) => t,
        Err(e) => {
            let error = format!("Invalid command target: {}", e);
            transport
                .send(&Message::ClientCommandError {
                    request_id: req.request_id,
                    error,
                })
                .await
                .ok();
            return Ok(());
        }
    };

    let command_options: CommandOptions = match serde_json::from_str(&req.options) {
        Ok(o) => o,
        Err(e) => {
            let error = format!("Invalid command options: {}", e);
            transport
                .send(&Message::ClientCommandError {
                    request_id: req.request_id,
                    error,
                })
                .await
                .ok();
            return Ok(());
        }
    };

    info!(
        target: "connection",
        peer=%req.peer,
        request_id=%req.request_id,
        "executing command: {} with args: {:?}",
        req.command,
        req.args
    );

    // Get the command distributor
    let distributor = match &ctx.state.command_distributor {
        Some(d) => d.clone(),
        None => {
            let error = "Command distributor not available".to_string();
            transport
                .send(&Message::ClientCommandError {
                    request_id: req.request_id,
                    error,
                })
                .await
                .ok();
            return Ok(());
        }
    };

    let wire_version = command_options.wire_version;

    // Execute the command using the distributor
    match distributor
        .execute_command(
            req.command.clone(),
            req.args.clone(),
            command_target,
            command_options,
            format!("client:{}", req.peer),
        )
        .await
    {
        Ok(mut result_receiver) => {
            let mut total_workers = 0;
            let mut cancelled = false;

            // Stream results back to client, while also accepting cancel/disconnect.
            loop {
                tokio::select! {
                    update = result_receiver.recv() => {
                        let Some(update) = update else { break };
                        use volli_commands::FleetCommandUpdate;

                        match update {
                            FleetCommandUpdate::WorkerCompleted {
                                worker_id, result, ..
                            } => {
                                let worker_name = lookup_worker_name(&ctx.state.workers, &worker_id).await;
                                let footer_worker_name =
                                    lookup_worker_name(&ctx.state.workers, &worker_id).await;
                                let worker_id_footer = worker_id.clone();
                                let success =
                                    matches!(result.status, volli_commands::CommandStatus::Success);
                                total_workers += 1;
                                let send_result = transport
                                    .send(&Message::ClientCommandResponse {
                                        request_id: req.request_id.clone(),
                                        worker_id: worker_id.clone(),
                                        worker_name,
                                        success,
                                        duration_millis: result.execution_time.as_millis() as u64,
                                        output: String::new(),
                                    })
                                    .await
                                    .map_err(|e| eyre!(e));
                                if send_result.is_err() {
                                    cancelled = true;
                                }
                                if let Some(ref footer) = result.footer {
                                    let payload = volli_commands::wire::encode_payload(footer, wire_version);
                                    let send_result = transport
                                        .send(&Message::ClientCommandFooter {
                                            request_id: req.request_id.clone(),
                                            worker_id: worker_id_footer,
                                            worker_name: footer_worker_name,
                                            payload,
                                            duration_millis: result.execution_time.as_millis() as u64,
                                            success,
                                        })
                                        .await
                                        .map_err(|e| eyre!(e));
                                    if send_result.is_err() {
                                        cancelled = true;
                                    }
                                }
                            }
                            FleetCommandUpdate::WorkerFailed { worker_id, .. } => {
                                let worker_name = lookup_worker_name(&ctx.state.workers, &worker_id).await;
                                total_workers += 1;
                                let result = transport
                                    .send(&Message::ClientCommandResponse {
                                        request_id: req.request_id.clone(),
                                        worker_id,
                                        worker_name,
                                        success: false,
                                        duration_millis: 0,
                                        output: String::new(),
                                    })
                                    .await
                                    .map_err(|e| eyre!(e));
                                if result.is_err() {
                                    cancelled = true;
                                }
                            }
                            FleetCommandUpdate::WorkerStarted {
                                worker_id, header, ..
                            } => {
                                let worker_name = lookup_worker_name(&ctx.state.workers, &worker_id).await;
                                // Stream header to client
                                let payload = volli_commands::wire::encode_payload(&header, wire_version);
                                let result = transport
                                    .send(&Message::ClientCommandHeader {
                                        request_id: req.request_id.clone(),
                                        worker_id,
                                        worker_name,
                                        payload,
                                    })
                                    .await
                                    .map_err(|e| eyre!(e));
                                if result.is_err() {
                                    cancelled = true;
                                }
                            }
                            FleetCommandUpdate::WorkerProgress {
                                worker_id, stream, ..
                            } => {
                                let worker_name = lookup_worker_name(&ctx.state.workers, &worker_id).await;
                                let payload = volli_commands::wire::encode_payload(&stream, wire_version);
                                let result = transport
                                    .send(&Message::ClientCommandStream {
                                        request_id: req.request_id.clone(),
                                        worker_id,
                                        worker_name,
                                        payload,
                                    })
                                    .await
                                    .map_err(|e| eyre!(e));
                                if result.is_err() {
                                    cancelled = true;
                                }
                            }
                            FleetCommandUpdate::Completed { .. } => {
                                let result = transport
                                    .send(&Message::ClientCommandComplete {
                                        request_id: req.request_id.clone(),
                                        total_workers,
                                    })
                                    .await
                                    .map_err(|e| eyre!(e));
                                let _ = result;
                                break;
                            }
                            _ => {
                                // Handle other update types as needed
                            }
                        }

                        if cancelled {
                            let _ = distributor
                                .cancel_command(&req.request_id, "client disconnected")
                                .await;
                            break;
                        }
                    }
                    msg = transport.recv() => {
                        match msg {
                            Ok(Some(Message::ClientCommandCancel { request_id })) => {
                                let _ = distributor
                                    .cancel_command(&request_id, "client cancelled")
                                    .await;
                                break;
                            }
                            Ok(None) => {
                                let _ = distributor
                                    .cancel_command(&req.request_id, "client disconnected")
                                    .await;
                                break;
                            }
                            Ok(Some(_)) => {}
                            Err(_) => {
                                let _ = distributor
                                    .cancel_command(&req.request_id, "client disconnected")
                                    .await;
                                break;
                            }
                        }
                    }
                }
            }
        }
        Err(e) => {
            let error = format!("Command execution failed: {}", e);
            transport
                .send(&Message::ClientCommandError {
                    request_id: req.request_id,
                    error,
                })
                .await
                .ok();
        }
    }

    Ok(())
}

async fn lookup_worker_name(
    table: &crate::workers::WorkerTable,
    worker_id: &str,
) -> Option<String> {
    let map = table.lock().await;
    map.get(worker_id)
        .and_then(|w| w.worker_name.clone())
        .filter(|name| !name.is_empty())
}