watcheragent/WatcherAgent/src/docker/serverclientcomm.rs

//! Server-client communication utilities for WatcherAgent
//!
//! Handles server commands, Docker image updates, and container management using the Bollard library.
//!
use crate::models::{DockerContainer, ServerMessage};
use crate::docker::container::{get_available_container};

use std::error::Error;
use bollard::Docker;
use bollard::query_parameters::{CreateImageOptions, RestartContainerOptions};
use futures_util::StreamExt;

/// Handles a message from the backend server and dispatches the appropriate action.
///
/// # Arguments
/// * `docker` - Reference to a Bollard Docker client.
/// * `msg` - The server message to handle.
///
/// # Returns
/// * `Result<(), Box<dyn Error + Send + Sync>>` - Ok if handled successfully, error otherwise.
pub async fn handle_server_message(docker: &Docker, msg: ServerMessage) -> Result<(), Box<dyn Error + Send + Sync>> {
    let msg = msg.clone();
    println!("Handling server message: {:?}", msg);

    // Handle different message types
    match msg.message_type.as_str() {
        "update_image" => {
            if let Some(image_name) = msg.data.get("image").and_then(|v| v.as_str()) {
                println!("Received update command for image: {}", image_name);
                // Call your update_docker_image function here
                update_docker_image(docker, image_name).await?;
                Ok(())
            } else {
                Err("Missing image name in update message".into())
            }
        }
        "restart_container" => {
            println!("Received restart container command");
            // Call your restart_container function here
            restart_container(docker).await?;
            Ok(())
        }
        "stop_agent" => {
            println!("Received stop agent command");
            // Implement graceful shutdown
            std::process::exit(0);
        }
        _ => {
            eprintln!("Unknown message type: {}", msg.message_type);
            Err(format!("Unknown message type: {}", msg.message_type).into())
        }
    }
}

/// Pulls a new Docker image and restarts the current container.
///
/// # Arguments
/// * `docker` - Reference to a Bollard Docker client.
/// * `image` - The name of the Docker image to pull.
///
/// # Returns
/// * `Result<(), Box<dyn Error + Send + Sync>>` - Ok if updated successfully, error otherwise.
pub async fn update_docker_image(docker: &Docker, image: &str) -> Result<(), Box<dyn Error + Send + Sync>> {
    println!("Updating to {}", image);

    // 1. Pull new image
    let mut stream = docker.create_image(
        Some(CreateImageOptions {
            from_image: Some(image.to_string()),
            ..Default::default()
        }),
        None,
        None,
    );

    // Use the stream with proper trait bounds
    while let Some(result) = StreamExt::next(&mut stream).await {
        match result {
            Ok(progress) => {
                if let Some(status) = progress.status {
                    println!("Pull status: {}", status);
                }
            }
            Err(e) => {
                eprintln!("Error pulling image: {}", e);
                break;
            }
        }
    }

    // 2. Restart the current container
    let _ = restart_container(docker).await;

    Ok(())
}

/// Finds the Docker container running the agent by image name.
///
/// # Arguments
/// * `docker` - Reference to a Bollard Docker client.
///
/// # Returns
/// * `Result<Option<DockerContainer>, Box<dyn Error + Send + Sync>>` - The agent's container info if found.
pub async fn get_client_container(docker: &Docker) -> Result<Option<DockerContainer>, Box<dyn Error + Send + Sync>> {
    let containers = get_available_container(docker).await;
    let client_image = "watcher-agent";
    
    // Find container with the specific image
    if let Some(container) = containers.iter().find(|c| c.image == client_image) {
        Ok(Some(container.clone()))
    } else {
        Ok(None)
    }
}

/// Restarts the agent's own Docker container.
///
/// # Arguments
/// * `docker` - Reference to a Bollard Docker client.
///
/// # Returns
/// * `Result<(), Box<dyn Error + Send + Sync>>` - Ok if restarted successfully, error otherwise.
pub async fn restart_container(docker: &Docker) -> Result<(), Box<dyn Error + Send + Sync>> {
    if let Ok(Some(container)) = get_client_container(docker).await {
        let container_id = container.clone().ID;
        println!("Restarting container {}", container_id);
        if let Err(e) = docker.restart_container(&container_id.to_string(), Some(RestartContainerOptions { signal: None, t: Some(0) }))
            .await
        {
            eprintln!("Failed to restart container: {}", e);
        }
    } else {
        eprintln!("No container ID found (HOSTNAME not set?)");
    }

    Ok(())
}