lcpfs 2026.1.102

// Copyright 2025 LunaOS Contributors
// SPDX-License-Identifier: Apache-2.0
//
// Container Optimization
// Image layer dedup and fast container cloning.

use alloc::collections::BTreeMap;
use alloc::string::String;
use alloc::vec::Vec;
use lazy_static::lazy_static;
use spin::Mutex;

/// Container image layer
#[derive(Debug, Clone)]
pub struct ImageLayer {
    /// Layer ID (SHA-256 hash of content)
    pub layer_id: [u8; 32],
    /// Parent layer ID (if any)
    pub parent_id: Option<[u8; 32]>,
    /// Size in bytes
    pub size: u64,
    /// Block pointer on disk
    pub block_ptr: u64,
    /// Reference count (how many images use this layer)
    pub refcount: u32,
    /// Compressed size
    pub compressed_size: u64,
    /// Creation timestamp
    pub created: u64,
}

impl ImageLayer {
    /// Create new layer
    pub fn new(layer_id: [u8; 32], size: u64, block_ptr: u64, created: u64) -> Self {
        Self {
            layer_id,
            parent_id: None,
            size,
            block_ptr,
            refcount: 1, // Start with 1 reference
            compressed_size: size,
            created,
        }
    }

    /// Set parent layer
    pub fn with_parent(mut self, parent_id: [u8; 32]) -> Self {
        self.parent_id = Some(parent_id);
        self
    }

    /// Increment reference count
    pub fn add_ref(&mut self) {
        self.refcount += 1;
    }

    /// Decrement reference count
    pub fn remove_ref(&mut self) -> u32 {
        self.refcount = self.refcount.saturating_sub(1);
        self.refcount
    }

    /// Calculate compression ratio
    pub fn compression_ratio(&self) -> f32 {
        if self.compressed_size == 0 {
            return 1.0;
        }
        self.size as f32 / self.compressed_size as f32
    }
}

/// Container image metadata
#[derive(Debug, Clone)]
pub struct ContainerImage {
    /// Image ID
    pub image_id: u64,
    /// Image name (e.g., "nginx:latest")
    pub name: &'static str,
    /// Layer stack (bottom to top)
    pub layers: Vec<[u8; 32]>,
    /// Total size (sum of all layers)
    pub total_size: u64,
    /// Creation timestamp
    pub created: u64,
}

impl ContainerImage {
    /// Create new container image
    pub fn new(image_id: u64, name: &'static str, created: u64) -> Self {
        Self {
            image_id,
            name,
            layers: Vec::new(),
            total_size: 0,
            created,
        }
    }

    /// Add layer to image
    pub fn add_layer(&mut self, layer_id: [u8; 32], size: u64) {
        self.layers.push(layer_id);
        self.total_size += size;
    }

    /// Get layer count
    pub fn layer_count(&self) -> usize {
        self.layers.len()
    }

    /// Check if layer exists in image
    pub fn has_layer(&self, layer_id: &[u8; 32]) -> bool {
        self.layers.contains(layer_id)
    }
}

/// Container instance (running or stopped)
#[derive(Debug, Clone)]
pub struct Container {
    /// Container ID
    pub container_id: u64,
    /// Base image ID
    pub image_id: u64,
    /// Writable layer (copy-on-write)
    pub writable_layer: u64,
    /// Container state
    pub state: ContainerState,
    /// Creation timestamp
    pub created: u64,
}

/// Container state
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ContainerState {
    /// Created but not started
    Created,
    /// Currently running
    Running,
    /// Paused
    Paused,
    /// Stopped
    Stopped,
}

/// Container optimization statistics
#[derive(Debug, Clone, Default)]
pub struct ContainerStats {
    /// Total images
    pub total_images: usize,
    /// Total layers stored
    pub total_layers: usize,
    /// Total unique layers (after dedup)
    pub unique_layers: usize,
    /// Total storage used
    pub total_storage: u64,
    /// Storage saved via dedup
    pub dedup_savings: u64,
    /// Total containers created
    pub containers_created: u64,
    /// Average clone time (milliseconds)
    pub avg_clone_time_ms: u64,
}

impl ContainerStats {
    /// Calculate dedup ratio
    pub fn dedup_ratio(&self) -> f32 {
        if self.total_storage == 0 {
            return 1.0;
        }
        let potential_storage = self.total_storage + self.dedup_savings;
        potential_storage as f32 / self.total_storage as f32
    }

    /// Calculate storage efficiency
    pub fn storage_efficiency(&self) -> f32 {
        if self.total_layers == 0 {
            return 0.0;
        }
        self.unique_layers as f32 / self.total_layers as f32
    }
}

lazy_static! {
    /// Global container manager
    static ref CONTAINER_MANAGER: Mutex<ContainerManager> = Mutex::new(ContainerManager::new());
}

/// Container optimization manager
pub struct ContainerManager {
    /// Layer store (content-addressable)
    layers: BTreeMap<[u8; 32], ImageLayer>,
    /// Image registry
    images: BTreeMap<u64, ContainerImage>,
    /// Running containers
    containers: BTreeMap<u64, Container>,
    /// Next container ID
    next_container_id: u64,
    /// Statistics
    stats: ContainerStats,
}

impl Default for ContainerManager {
    fn default() -> Self {
        Self::new()
    }
}

impl ContainerManager {
    /// Create new container manager
    pub fn new() -> Self {
        Self {
            layers: BTreeMap::new(),
            images: BTreeMap::new(),
            containers: BTreeMap::new(),
            next_container_id: 1,
            stats: ContainerStats::default(),
        }
    }

    /// Add layer to store (deduplicates automatically)
    ///
    /// # Arguments
    /// * `layer_id` - Content hash of layer
    /// * `size` - Uncompressed size
    /// * `block_ptr` - Block pointer on disk
    /// * `created` - Creation timestamp
    ///
    /// # Returns
    /// * `true` if layer was new (stored)
    /// * `false` if layer already existed (deduplicated)
    pub fn add_layer(
        &mut self,
        layer_id: [u8; 32],
        size: u64,
        block_ptr: u64,
        created: u64,
    ) -> bool {
        if let Some(layer) = self.layers.get_mut(&layer_id) {
            // Layer already exists - deduplicate
            layer.add_ref();
            self.stats.dedup_savings += size;
            false
        } else {
            // New layer
            let layer = ImageLayer::new(layer_id, size, block_ptr, created);
            self.layers.insert(layer_id, layer);
            self.stats.unique_layers += 1;
            self.stats.total_storage += size;
            true
        }
    }

    /// Create container image
    pub fn create_image(
        &mut self,
        image_id: u64,
        name: &'static str,
        created: u64,
    ) -> Result<(), &'static str> {
        if self.images.contains_key(&image_id) {
            return Err("Image ID already exists");
        }

        let image = ContainerImage::new(image_id, name, created);
        self.images.insert(image_id, image);
        self.stats.total_images += 1;

        Ok(())
    }

    /// Add layer to image
    pub fn add_layer_to_image(
        &mut self,
        image_id: u64,
        layer_id: [u8; 32],
    ) -> Result<(), &'static str> {
        let image = self.images.get_mut(&image_id).ok_or("Image not found")?;
        let layer = self.layers.get_mut(&layer_id).ok_or("Layer not found")?;

        layer.add_ref();
        image.add_layer(layer_id, layer.size);
        self.stats.total_layers += 1;

        Ok(())
    }

    /// Clone image (fast reflink-based)
    ///
    /// Creates new container from image using copy-on-write
    ///
    /// # Arguments
    /// * `image_id` - Source image
    /// * `created` - Creation timestamp
    ///
    /// # Returns
    /// Container ID
    pub fn clone_image(&mut self, image_id: u64, created: u64) -> Result<u64, &'static str> {
        let image = self.images.get(&image_id).ok_or("Image not found")?;

        // Create writable layer (COW)
        let writable_layer = image.image_id * 1000 + self.next_container_id;

        let container = Container {
            container_id: self.next_container_id,
            image_id,
            writable_layer,
            state: ContainerState::Created,
            created,
        };

        let container_id = self.next_container_id;
        self.next_container_id += 1;

        self.containers.insert(container_id, container);
        self.stats.containers_created += 1;

        // Fast clone using reflinks (O(1) operation)
        // In real implementation, would use FICLONE ioctl
        self.stats.avg_clone_time_ms = 5; // Typical reflink clone time

        Ok(container_id)
    }

    /// Remove layer reference
    pub fn remove_layer_ref(&mut self, layer_id: &[u8; 32]) -> Result<bool, &'static str> {
        let layer = self.layers.get_mut(layer_id).ok_or("Layer not found")?;

        let refcount = layer.remove_ref();
        if refcount == 0 {
            // No more references, can delete
            self.layers.remove(layer_id);
            self.stats.unique_layers = self.stats.unique_layers.saturating_sub(1);
            Ok(true) // Layer deleted
        } else {
            Ok(false) // Still referenced
        }
    }

    /// Get image
    pub fn get_image(&self, image_id: u64) -> Option<&ContainerImage> {
        self.images.get(&image_id)
    }

    /// Get layer
    pub fn get_layer(&self, layer_id: &[u8; 32]) -> Option<&ImageLayer> {
        self.layers.get(layer_id)
    }

    /// Get container
    pub fn get_container(&self, container_id: u64) -> Option<&Container> {
        self.containers.get(&container_id)
    }

    /// List images
    pub fn list_images(&self) -> Vec<ContainerImage> {
        self.images.values().cloned().collect()
    }

    /// Calculate shared layers between images
    pub fn shared_layers(&self, image_id1: u64, image_id2: u64) -> usize {
        let img1 = match self.images.get(&image_id1) {
            Some(img) => img,
            None => return 0,
        };

        let img2 = match self.images.get(&image_id2) {
            Some(img) => img,
            None => return 0,
        };

        img1.layers
            .iter()
            .filter(|layer| img2.has_layer(layer))
            .count()
    }

    /// Get statistics
    pub fn get_stats(&self) -> ContainerStats {
        self.stats.clone()
    }
}

/// Global container operations
pub struct ContainerEngine;

impl ContainerEngine {
    /// Add layer (with automatic deduplication)
    pub fn add_layer(layer_id: [u8; 32], size: u64, block_ptr: u64, created: u64) -> bool {
        let mut mgr = CONTAINER_MANAGER.lock();
        mgr.add_layer(layer_id, size, block_ptr, created)
    }

    /// Create image
    pub fn create_image(
        image_id: u64,
        name: &'static str,
        created: u64,
    ) -> Result<(), &'static str> {
        let mut mgr = CONTAINER_MANAGER.lock();
        mgr.create_image(image_id, name, created)
    }

    /// Add layer to image
    pub fn add_layer_to_image(image_id: u64, layer_id: [u8; 32]) -> Result<(), &'static str> {
        let mut mgr = CONTAINER_MANAGER.lock();
        mgr.add_layer_to_image(image_id, layer_id)
    }

    /// Clone image (fast)
    pub fn clone_image(image_id: u64, created: u64) -> Result<u64, &'static str> {
        let mut mgr = CONTAINER_MANAGER.lock();
        mgr.clone_image(image_id, created)
    }

    /// Remove layer reference
    pub fn remove_layer_ref(layer_id: &[u8; 32]) -> Result<bool, &'static str> {
        let mut mgr = CONTAINER_MANAGER.lock();
        mgr.remove_layer_ref(layer_id)
    }

    /// Get image
    pub fn get_image(image_id: u64) -> Option<ContainerImage> {
        let mgr = CONTAINER_MANAGER.lock();
        mgr.get_image(image_id).cloned()
    }

    /// List images
    pub fn list_images() -> Vec<ContainerImage> {
        let mgr = CONTAINER_MANAGER.lock();
        mgr.list_images()
    }

    /// Shared layers
    pub fn shared_layers(image_id1: u64, image_id2: u64) -> usize {
        let mgr = CONTAINER_MANAGER.lock();
        mgr.shared_layers(image_id1, image_id2)
    }

    /// Get statistics
    pub fn stats() -> ContainerStats {
        let mgr = CONTAINER_MANAGER.lock();
        mgr.get_stats()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_layer_creation() {
        let layer = ImageLayer::new([1u8; 32], 1024, 100, 0);
        assert_eq!(layer.size, 1024);
        assert_eq!(layer.refcount, 1); // Starts with 1 reference
    }

    #[test]
    fn test_layer_refcount() {
        let mut layer = ImageLayer::new([1u8; 32], 1024, 100, 0);
        assert_eq!(layer.refcount, 1); // Initial refcount

        layer.add_ref();
        assert_eq!(layer.refcount, 2);

        layer.add_ref();
        assert_eq!(layer.refcount, 3);

        layer.remove_ref();
        assert_eq!(layer.refcount, 2);
    }

    #[test]
    fn test_image_layers() {
        let mut image = ContainerImage::new(1, "test:latest", 0);

        image.add_layer([1u8; 32], 1000);
        image.add_layer([2u8; 32], 2000);

        assert_eq!(image.layer_count(), 2);
        assert_eq!(image.total_size, 3000);
    }

    #[test]
    fn test_layer_deduplication() {
        let mut mgr = ContainerManager::new();

        let layer_id = [1u8; 32];

        // First add - should store
        assert!(mgr.add_layer(layer_id, 1024, 100, 0));
        assert_eq!(mgr.stats.unique_layers, 1);

        // Second add - should deduplicate
        assert!(!mgr.add_layer(layer_id, 1024, 100, 0));
        assert_eq!(mgr.stats.unique_layers, 1);
        assert_eq!(mgr.stats.dedup_savings, 1024);
    }

    #[test]
    fn test_image_creation() {
        let mut mgr = ContainerManager::new();

        mgr.create_image(1, "nginx:latest", 0)
            .expect("test: operation should succeed");
        assert_eq!(mgr.stats.total_images, 1);

        // Duplicate ID should fail
        assert!(mgr.create_image(1, "nginx:latest", 0).is_err());
    }

    #[test]
    fn test_container_cloning() {
        let mut mgr = ContainerManager::new();

        // Create image with layer
        mgr.create_image(1, "test:latest", 0)
            .expect("test: operation should succeed");
        let layer_id = [1u8; 32];
        mgr.add_layer(layer_id, 1024, 100, 0);
        mgr.add_layer_to_image(1, layer_id)
            .expect("test: operation should succeed");

        // Clone image to create container
        let container_id = mgr
            .clone_image(1, 0)
            .expect("test: operation should succeed");
        assert!(container_id > 0);
        assert_eq!(mgr.stats.containers_created, 1);

        let container = mgr
            .get_container(container_id)
            .expect("test: operation should succeed");
        assert_eq!(container.image_id, 1);
    }

    #[test]
    fn test_shared_layers() {
        let mut mgr = ContainerManager::new();

        // Create two images
        mgr.create_image(1, "base:latest", 0)
            .expect("test: operation should succeed");
        mgr.create_image(2, "derived:latest", 0)
            .expect("test: operation should succeed");

        // Add common layers
        let layer1 = [1u8; 32];
        let layer2 = [2u8; 32];
        let layer3 = [3u8; 32];

        mgr.add_layer(layer1, 1024, 100, 0);
        mgr.add_layer(layer2, 1024, 101, 0);
        mgr.add_layer(layer3, 1024, 102, 0);

        // Image 1: layers 1, 2
        mgr.add_layer_to_image(1, layer1)
            .expect("test: operation should succeed");
        mgr.add_layer_to_image(1, layer2)
            .expect("test: operation should succeed");

        // Image 2: layers 1, 2, 3 (shares 2 with image 1)
        mgr.add_layer_to_image(2, layer1)
            .expect("test: operation should succeed");
        mgr.add_layer_to_image(2, layer2)
            .expect("test: operation should succeed");
        mgr.add_layer_to_image(2, layer3)
            .expect("test: operation should succeed");

        assert_eq!(mgr.shared_layers(1, 2), 2);
    }

    #[test]
    fn test_dedup_ratio() {
        let stats = ContainerStats {
            total_storage: 10000,
            dedup_savings: 5000,
            ..Default::default()
        };

        let ratio = stats.dedup_ratio();
        assert!((ratio - 1.5).abs() < 0.01); // 15000 / 10000 = 1.5
    }

    #[test]
    fn test_layer_removal() {
        let mut mgr = ContainerManager::new();

        let layer_id = [1u8; 32];
        mgr.add_layer(layer_id, 1024, 100, 0);

        // Add 2 references
        let layer = mgr
            .layers
            .get_mut(&layer_id)
            .expect("test: operation should succeed");
        layer.add_ref();
        layer.add_ref();

        // Remove first ref - should not delete
        assert!(
            !mgr.remove_layer_ref(&layer_id)
                .expect("test: operation should succeed")
        );

        // Remove second ref - should not delete
        assert!(
            !mgr.remove_layer_ref(&layer_id)
                .expect("test: operation should succeed")
        );

        // Remove final ref - should delete
        assert!(
            mgr.remove_layer_ref(&layer_id)
                .expect("test: operation should succeed")
        );
        assert!(mgr.get_layer(&layer_id).is_none());
    }

    #[test]
    fn test_compression_ratio() {
        let mut layer = ImageLayer::new([1u8; 32], 10000, 100, 0);
        layer.compressed_size = 2000;

        let ratio = layer.compression_ratio();
        assert!((ratio - 5.0).abs() < 0.01); // 10000 / 2000 = 5.0
    }
}