watcheragent/WatcherAgent/src/docker/mod.rs

/// # Docker Module
///
/// This module provides Docker integration for WatcherAgent, including container enumeration, statistics, and lifecycle management.
///
/// ## Responsibilities
/// - **Container Management:** Lists, inspects, and manages Docker containers relevant to the agent.
/// - **Statistics Aggregation:** Collects network and CPU statistics for all managed containers.
/// - **Lifecycle Operations:** Supports container restart and ID lookup for agent self-management.
///
pub mod container;
pub mod serverclientcomm;
pub mod stats;

use crate::models::{
    DockerCollectMetricDto, DockerContainer, DockerContainerCpuDto, DockerContainerInfo,
    DockerContainerNetworkDto, DockerContainerRamDto, DockerMetricDto, DockerRegistrationDto,
};
use bollard::Docker;
use std::error::Error;

/// Main Docker manager that holds the Docker client and provides all operations
#[derive(Debug, Clone)]
pub struct DockerManager {
    pub docker: Docker,
}

impl Default for DockerManager {
    fn default() -> Self {
        Self {
            docker: Docker::connect_with_local_defaults()
                .unwrap_or_else(|e| panic!("Failed to create default Docker connection: {}", e)),
        }
    }
}

impl DockerManager {
    /// Creates a new DockerManager instance
    pub fn new() -> Result<Self, Box<dyn Error + Send + Sync>> {
        let docker = Docker::connect_with_local_defaults()
            .map_err(|e| format!("Failed to connect to Docker: {}", e))?;

        Ok(Self { docker })
    }

    /// Creates a DockerManager instance with optional Docker connection
    pub fn new_optional() -> Option<Self> {
        Docker::connect_with_local_defaults()
            .map(|docker| Self { docker })
            .ok()
    }

    /// Finds the Docker container running the agent by image name
    pub async fn get_client_container(
        &self,
    ) -> Result<Option<DockerContainer>, Box<dyn Error + Send + Sync>> {
        let containers = container::get_available_containers(&self.docker).await;
        let client_image = "watcher-agent";

        Ok(containers
            .into_iter()
            .find(|c| c.clone().image.unwrap().contains(client_image))
            .map(|container| DockerContainer {
                id: container.id,
                image: container.image,
                name: container.name,
            }))
    }

    /// Gets the current client version (image name) if running in Docker
    pub async fn get_client_version(&self) -> String {
        match self.get_client_container().await {
            Ok(Some(container)) => container
                .image
                .clone()
                .unwrap()
                .split(':')
                .next()
                .unwrap_or("unknown")
                .to_string(),
            Ok(None) => {
                println!("Warning: No WatcherAgent container found");
                "unknown".to_string()
            }
            Err(e) => {
                println!("Warning: Could not get current image version: {}", e);
                "unknown".to_string()
            }
        }
    }

    /// Checks if Docker is available and the agent is running in a container
    pub async fn is_dockerized(&self) -> bool {
        self.get_client_container()
            .await
            .map(|c| c.is_some())
            .unwrap_or(false)
    }

    /// Gets all available containers as DTOs for registration
    pub async fn get_containers(
        &self,
    ) -> Result<Vec<DockerContainer>, Box<dyn Error + Send + Sync>> {
        let containers = container::get_available_containers(&self.docker).await;

        Ok(containers
            .into_iter()
            .map(|container| DockerContainer {
                id: container.id,
                image: container.image,
                name: container.name,
            })
            .collect())
    }

    /// Gets the number of running containers
    pub async fn get_container_count(&self) -> Result<usize, Box<dyn Error + Send + Sync>> {
        let containers = container::get_available_containers(&self.docker).await;
        Ok(containers.len())
    }

    /// Restarts a specific container by ID
    pub async fn restart_container(
        &self,
        container_id: &str,
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        container::restart_container(&self.docker, container_id).await
    }

    /// Collects Docker metrics for all containers
    pub async fn collect_metrics(&self) -> Result<DockerMetricDto, Box<dyn Error + Send + Sync>> {
        let containers = self.get_containers().await?;
        if let Some(first_container) = containers.first() {
            println!("Debug: Testing stats for container {}", first_container.id);
            let _ = self.debug_container_stats(&first_container.id).await;
        }


        // Get stats with proper error handling
        let stats_result = stats::get_container_stats(&self.docker).await;
        let (cpu_stats, net_stats, mem_stats) = match stats_result {
            Ok(stats) => stats,
            Err(e) => {
                eprintln!("Warning: Failed to get container stats: {}", e);
                // Return empty stats instead of failing completely
                (Vec::new(), Vec::new(), Vec::new())
            }
        };

        println!("Debug: Found {} containers, {} CPU stats, {} network stats, {} memory stats", 
            containers.len(), cpu_stats.len(), net_stats.len(), mem_stats.len());

        let container_infos_total: Vec<_> = containers
            .into_iter()
            .map(|container| {
                // Use short ID for matching (first 12 chars)
                let container_short_id = if container.id.len() > 12 { 
                    &container.id[..12] 
                } else { 
                    &container.id 
                };

                let cpu = cpu_stats
                    .iter()
                    .find(|c| {
                        c.container_id.as_ref()
                            .map(|id| id.starts_with(container_short_id))
                            .unwrap_or(false)
                    })
                    .cloned();
                
                let network = net_stats
                    .iter()
                    .find(|n| {
                        n.container_id.as_ref()
                            .map(|id| id.starts_with(container_short_id))
                            .unwrap_or(false)
                    })
                    .cloned();
                
                let ram = mem_stats
                    .iter()
                    .find(|m| {
                        m.container_id.as_ref()
                            .map(|id| id.starts_with(container_short_id))
                            .unwrap_or(false)
                    })
                    .cloned();

                // Debug output for this container
                if cpu.is_none() || network.is_none() || ram.is_none() {
                    println!("Debug: Container {} - CPU: {:?}, Network: {:?}, RAM: {:?}", 
                        container_short_id, cpu.is_some(), network.is_some(), ram.is_some());
                }

                DockerContainerInfo {
                    container: Some(container),
                    status: None,
                    cpu,
                    network,
                    ram,
                }
            })
            .collect();

        let container_infos: Vec<DockerCollectMetricDto> = container_infos_total
            .into_iter()
            .filter_map(|info| {
                let container = match info.container {
                    Some(c) => c,
                    None => {
                        eprintln!("Warning: Container info missing container data, skipping");
                        return None;
                    }
                };

                // Safely handle CPU data with defaults
                let cpu_dto = if let Some(cpu) = info.cpu {
                    DockerContainerCpuDto {
                        cpu_load: cpu.cpu_usage_percent,
                    }
                } else {
                    DockerContainerCpuDto { cpu_load: None }
                };

                // Safely handle RAM data with defaults
                let ram_dto = if let Some(ram) = info.ram {
                    DockerContainerRamDto {
                        ram_load: ram.memory_usage_percent,
                    }
                } else {
                    DockerContainerRamDto { ram_load: None }
                };

                // Safely handle network data with defaults
                let network_dto = if let Some(net) = info.network {
                    DockerContainerNetworkDto {
                        net_in: net.rx_bytes.map(|bytes| bytes as f64),
                        net_out: net.tx_bytes.map(|bytes| bytes as f64),
                    }
                } else {
                    DockerContainerNetworkDto {
                        net_in: None,
                        net_out: None,
                    }
                };

                Some(DockerCollectMetricDto {
                    id: container.id,
                    cpu: cpu_dto,
                    ram: ram_dto,
                    network: network_dto,
                })
            })
            .collect();

        let dto = DockerMetricDto {
            server_id: 0, // This should be set by the caller
            containers: serde_json::to_string(&container_infos)?,
        };

        Ok(dto)
    }

    pub async fn create_registration_dto(
        &self,
    ) -> Result<DockerRegistrationDto, Box<dyn Error + Send + Sync>> {
        let containers = self.get_containers().await?;
        let dto = DockerRegistrationDto {
            server_id: 0, // This will be set by the caller
            containers: serde_json::to_string(&containers)
                .unwrap_or_else(|_| "[]".to_string()), // Fallback to empty array
        };

        Ok(dto)
    }

    /// Debug function to check stats collection for a specific container
    pub async fn debug_container_stats(
        &self, 
        container_id: &str
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        println!("=== DEBUG STATS FOR CONTAINER {} ===", container_id);
        
        let (cpu_info, net_info, mem_info) = stats::get_single_container_stats(&self.docker, container_id).await?;
        
        println!("CPU Info: {:?}", cpu_info);
        println!("Network Info: {:?}", net_info);
        println!("Memory Info: {:?}", mem_info);
        
        // Also try the individual stats functions
        println!("--- Individual CPU Stats ---");
        match stats::cpu::get_single_container_cpu_stats(&self.docker, container_id).await {
            Ok(cpu) => println!("CPU: {:?}", cpu),
            Err(e) => println!("CPU Error: {}", e),
        }
        
        println!("--- Individual Network Stats ---");
        match stats::network::get_single_container_network_stats(&self.docker, container_id).await {
            Ok(net) => println!("Network: {:?}", net),
            Err(e) => println!("Network Error: {}", e),
        }
        
        println!("--- Individual Memory Stats ---");
        match stats::ram::get_single_container_memory_stats(&self.docker, container_id).await {
            Ok(mem) => println!("Memory: {:?}", mem),
            Err(e) => println!("Memory Error: {}", e),
        }
        
        Ok(())
    }
}

// Keep these as utility functions if needed, but they should use DockerManager internally
impl DockerContainer {
    /// Returns the container ID
    pub fn id(&self) -> &str {
        &self.id
    }

    /// Returns the image name
    pub fn image(&self) -> &str {
        &self.image.as_deref().unwrap_or("unknown")
    }

    /// Returns the container name
    pub fn name(&self) -> &str {
        &self.name.as_deref().unwrap_or("unknown")
    }
}