apt-ostree/src/commands/compose/composer.rs

//! Tree composer for apt-ostree compose

use std::path::PathBuf;
use apt_ostree::lib::error::{AptOstreeError, AptOstreeResult};
use super::treefile::Treefile;
use super::package_manager::PackageManager;
use super::ostree_integration::OstreeIntegration;
use super::container::ContainerGenerator;

/// Main tree composer that orchestrates the composition process
pub struct TreeComposer {
    workdir: PathBuf,
    package_manager: PackageManager,
    ostree_integration: OstreeIntegration,
    container_generator: ContainerGenerator,
}

impl TreeComposer {
    /// Create a new tree composer instance
    pub fn new(options: &crate::commands::compose::ComposeOptions) -> AptOstreeResult<Self> {
        let workdir = options.workdir.clone().unwrap_or_else(|| {
            std::env::temp_dir().join("apt-ostree-compose")
        });
        
        let package_manager = PackageManager::new(options)?;
        let ostree_integration = OstreeIntegration::new(options.repo.as_deref(), &workdir)?;
        let container_generator = ContainerGenerator::new(&workdir, &workdir);
        
        Ok(Self {
            workdir,
            package_manager,
            ostree_integration,
            container_generator,
        })
    }
    
    /// Compose a complete tree from a treefile
    pub async fn compose_tree(&self, treefile: &Treefile) -> AptOstreeResult<String> {
        println!("Starting tree composition for: {}", treefile.metadata.ref_name);
        
        // Step 1: Set up build environment
        self.setup_build_environment(treefile).await?;
        
        // Step 2: Initialize base system
        if let Some(base_image) = &treefile.base_image {
            self.package_manager.initialize_base_system(base_image).await?;
        }
        
        // Step 3: Configure package sources
        if !treefile.repositories.is_empty() {
            self.package_manager.setup_package_sources(&treefile.repositories).await?;
        }
        
        // Step 4: Update package cache
        self.package_manager.update_cache().await?;
        
        // Step 5: Install base packages
        if let Some(packages) = &treefile.packages.base {
            self.install_packages(packages, "base").await?;
        }
        
        // Step 6: Install additional packages
        if let Some(packages) = &treefile.packages.additional {
            self.install_packages(packages, "additional").await?;
        }
        
        // Step 7: Apply customizations
        if let Some(customizations) = &treefile.customizations {
            self.apply_customizations(customizations).await?;
        }
        
        // Step 8: Run post-installation scripts
        self.package_manager.run_post_install_scripts().await?;
        
        // Step 9: Update package database
        self.package_manager.update_package_database().await?;
        
        // Step 10: Initialize OSTree repository
        self.ostree_integration.init_repository().await?;
        
        // Step 11: Create OSTree commit
        let parent_ref = self.get_parent_reference(treefile).await?;
        let commit_hash = self.ostree_integration.create_commit(&treefile.metadata, parent_ref.as_deref()).await?;
        
        // Step 12: Update reference
        self.ostree_integration.update_reference(&treefile.metadata.ref_name, &commit_hash).await?;
        
        // Step 13: Create repository summary
        self.ostree_integration.create_summary().await?;
        
        // Step 14: Generate container image if requested
        if let Some(output_config) = &treefile.output {
            if output_config.generate_container {
                self.container_generator.generate_image(&treefile.metadata.ref_name, output_config).await?;
            }
        }
        
        // Step 15: Clean up build artifacts
        self.cleanup_build_artifacts().await?;
        
        println!("✅ Tree composition completed successfully");
        println!("Commit hash: {}", commit_hash);
        println!("Reference: {}", treefile.metadata.ref_name);
        
        Ok(commit_hash)
    }
    
    /// Set up the build environment
    async fn setup_build_environment(&self, treefile: &Treefile) -> AptOstreeResult<()> {
        println!("Setting up build environment...");
        
        // Create working directory
        std::fs::create_dir_all(&self.workdir)
            .map_err(|e| AptOstreeError::System(format!("Failed to create work directory: {}", e)))?;
        
        // Create build root directory
        let build_root = self.workdir.join("build-root");
        if build_root.exists() {
            std::fs::remove_dir_all(&build_root)
                .map_err(|e| AptOstreeError::System(format!("Failed to clean build root: {}", e)))?;
        }
        std::fs::create_dir_all(&build_root)
            .map_err(|e| AptOstreeError::System(format!("Failed to create build root: {}", e)))?;
        
        // Create necessary subdirectories
        let dirs = ["etc", "var", "usr", "tmp"];
        for dir in &dirs {
            let path = build_root.join(dir);
            std::fs::create_dir_all(&path)
                .map_err(|e| AptOstreeError::System(format!("Failed to create directory {}: {}", dir, e)))?;
        }
        
        println!("✅ Build environment set up successfully");
        Ok(())
    }
    
    /// Install packages
    async fn install_packages(&self, packages: &[String], category: &str) -> AptOstreeResult<()> {
        println!("Installing {} packages: {:?}", category, packages);
        
        // Resolve dependencies first
        let all_packages = self.package_manager.resolve_dependencies(packages).await?;
        println!("Resolved {} packages (including dependencies)", all_packages.len());
        
        // Install packages
        for (i, package) in all_packages.iter().enumerate() {
            println!("[{}/{}] Installing {}", i + 1, all_packages.len(), package);
            self.package_manager.install_package(package).await?;
        }
        
        println!("✅ {} packages installed successfully", category);
        Ok(())
    }
    
    /// Apply customizations
    async fn apply_customizations(&self, customizations: &super::treefile::Customizations) -> AptOstreeResult<()> {
        println!("Applying customizations...");
        
        let build_root = self.workdir.join("build-root");
        
        // Apply file customizations
        if let Some(files) = &customizations.files {
            for file_custom in files {
                let file_path = build_root.join(&file_custom.path.trim_start_matches('/'));
                
                // Create parent directory if it doesn't exist
                if let Some(parent) = file_path.parent() {
                    std::fs::create_dir_all(parent)
                        .map_err(|e| AptOstreeError::System(format!("Failed to create directory for {}: {}", file_custom.path, e)))?;
                }
                
                // Write file content if provided
                if let Some(content) = &file_custom.content {
                    std::fs::write(&file_path, content)
                        .map_err(|e| AptOstreeError::System(format!("Failed to write file {}: {}", file_custom.path, e)))?;
                    
                    println!("Created file: {}", file_custom.path);
                }
            }
        }
        
        // Apply system customizations
        if let Some(system_mods) = &customizations.system {
            for system_mod in system_mods {
                println!("Applying system modification: {:?}", system_mod);
                // TODO: Implement system modifications
            }
        }
        
        // Apply script customizations
        if let Some(scripts) = &customizations.scripts {
            for script in scripts {
                println!("Running script: {}", script.name);
                // TODO: Implement script execution
            }
        }
        
        println!("✅ Customizations applied successfully");
        Ok(())
    }
    
    /// Get parent reference
    async fn get_parent_reference(&self, treefile: &Treefile) -> AptOstreeResult<Option<String>> {
        // Check if parent reference is specified in treefile metadata
        if let Some(parent) = &treefile.metadata.parent {
            // Verify parent reference exists
            if self.ostree_integration.reference_exists(parent).await? {
                println!("Using parent reference: {}", parent);
                return Ok(Some(parent.clone()));
            } else {
                println!("Warning: Parent reference {} not found, creating without parent", parent);
            }
        }
        
        // Check if we can find a previous commit for the same reference
        if let Ok(Some(commit_hash)) = self.ostree_integration.get_commit_hash(&treefile.metadata.ref_name).await {
            println!("Using previous commit as parent: {}", commit_hash);
            return Ok(Some(commit_hash));
        }
        
        println!("No parent reference found, creating initial commit");
        Ok(None)
    }
    
    /// Clean up build artifacts
    async fn cleanup_build_artifacts(&self) -> AptOstreeResult<()> {
        println!("Cleaning up build artifacts...");
        
        // Clean up package manager state
        self.package_manager.cleanup().await?;
        
        // Remove temporary files
        let temp_dirs = ["tmp", "var/tmp"];
        let build_root = self.workdir.join("build-root");
        
        for temp_dir in &temp_dirs {
            let path = build_root.join(temp_dir);
            if path.exists() {
                std::fs::remove_dir_all(&path)
                    .map_err(|e| AptOstreeError::System(format!("Failed to remove temp directory {}: {}", temp_dir, e)))?;
            }
        }
        
        println!("✅ Build artifacts cleaned up successfully");
        Ok(())
    }
}