#!/usr/bin/env python3
"""
Debian Kernel Stage for osbuild

This stage handles Debian kernel installation, initramfs generation using initramfs-tools,
and OSTree integration. It replaces the dracut stage used in Fedora/RHEL systems.

Author: Debian bootc-image-builder team
License: Same as original bootc-image-builder
"""

import os
import subprocess
import tempfile
import json
import glob
import shutil
from typing import Dict, List, Optional, Any
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class DebianKernelStage:
    """
    osbuild stage for Debian kernel handling and initramfs generation.
    
    This stage handles:
    - Kernel detection in Debian paths
    - initramfs-tools integration with OSTree
    - Kernel module path management
    - OSTree-specific kernel configuration
    """
    
    def __init__(self, options: Dict[str, Any]):
        """
        Initialize the Debian Kernel stage with configuration options.
        
        Args:
            options: Dictionary containing stage configuration
                - kernel_package: Kernel package to install (e.g., 'linux-image-amd64')
                - initramfs_tools: Whether to use initramfs-tools
                - ostree_integration: Whether to enable OSTree integration
                - modules_autoload: Whether to auto-load kernel modules
                - kernel_version: Specific kernel version (optional)
        """
        self.options = options
        self.kernel_package = options.get('kernel_package', 'linux-image-amd64')
        self.initramfs_tools = options.get('initramfs_tools', True)
        self.ostree_integration = options.get('ostree_integration', True)
        self.modules_autoload = options.get('modules_autoload', True)
        self.kernel_version = options.get('kernel_version', None)
        
        logger.info(f"Debian Kernel Stage initialized")
        logger.info(f"Kernel package: {self.kernel_package}")
        logger.info(f"OSTree integration: {self.ostree_integration}")
    
    def run(self, context) -> None:
        """
        Execute the Debian Kernel stage within the osbuild context.
        
        Args:
            context: osbuild context providing chroot access
        """
        logger.info("Starting Debian Kernel stage execution")
        
        try:
            # Step 1: Detect kernel after package installation
            kernel_info = self._detect_kernel(context)
            
            # Step 2: Set up kernel module paths
            self._setup_kernel_modules(context, kernel_info)
            
            # Step 3: Configure initramfs-tools for OSTree
            if self.initramfs_tools:
                self._configure_initramfs_tools(context, kernel_info)
            
            # Step 4: Generate initramfs
            if self.initramfs_tools:
                self._generate_initramfs(context, kernel_info)
            
            # Step 5: Set up OSTree integration
            if self.ostree_integration:
                self._setup_ostree_integration(context, kernel_info)
            
            logger.info("Debian Kernel stage completed successfully")
            
        except Exception as e:
            logger.error(f"Debian Kernel stage failed: {e}")
            raise
    
    def _detect_kernel(self, context) -> Dict[str, str]:
        """
        Detect kernel information in Debian paths.
        
        Args:
            context: osbuild context
            
        Returns:
            Dictionary containing kernel information
        """
        logger.info("Detecting kernel in Debian paths")
        
        kernel_info = {}
        
        # Look for kernel in /boot/vmlinuz-*
        boot_path = os.path.join(context.root, "boot")
        if os.path.exists(boot_path):
            kernel_files = glob.glob(os.path.join(boot_path, "vmlinuz-*"))
            if kernel_files:
                # Sort by modification time to get the latest
                kernel_files.sort(key=lambda x: os.path.getmtime(x), reverse=True)
                kernel_path = kernel_files[0]
                kernel_info['kernel_path'] = kernel_path
                kernel_info['kernel_version'] = os.path.basename(kernel_path).replace('vmlinuz-', '')
                logger.info(f"Found kernel: {kernel_path}")
        
        # Look for kernel modules in /usr/lib/modules/
        modules_path = os.path.join(context.root, "usr", "lib", "modules")
        if os.path.exists(modules_path):
            module_dirs = [d for d in os.listdir(modules_path) 
                          if os.path.isdir(os.path.join(modules_path, d))]
            if module_dirs:
                # Sort by version to get the latest (handle Debian version format)
                def version_key(version):
                    try:
                        # Split version like "6.1.0-13-amd64" into parts
                        parts = version.split('-')[0].split('.')  # Get "6.1.0"
                        return [int(y) for y in parts]
                    except (ValueError, IndexError):
                        # If parsing fails, use string comparison
                        return [0]
                
                module_dirs.sort(key=version_key, reverse=True)
                kernel_info['modules_path'] = os.path.join(modules_path, module_dirs[0])
                kernel_info['kernel_version'] = module_dirs[0]
                logger.info(f"Found kernel modules: {kernel_info['modules_path']}")
        
        if not kernel_info:
            raise RuntimeError("No kernel found in expected Debian paths")
        
        return kernel_info
    
    def _setup_kernel_modules(self, context, kernel_info: Dict[str, str]) -> None:
        """
        Set up kernel module paths and dependencies.
        
        Args:
            context: osbuild context
            kernel_info: Kernel information dictionary
        """
        logger.info("Setting up kernel modules")
        
        modules_path = kernel_info.get('modules_path')
        if not modules_path:
            logger.warning("No kernel modules path found")
            return
        
        # Create modules.dep file if it doesn't exist
        modules_dep = os.path.join(modules_path, "modules.dep")
        if not os.path.exists(modules_dep):
            logger.info("Generating modules.dep file")
            cmd = ["depmod", "-b", context.root, kernel_info['kernel_version']]
            result = context.run(cmd)
            if result.returncode != 0:
                logger.warning(f"Failed to generate modules.dep: {result.stderr}")
        
        # Set up module autoload if enabled
        if self.modules_autoload:
            self._setup_module_autoload(context, kernel_info)
    
    def _setup_module_autoload(self, context, kernel_info: Dict[str, str]) -> None:
        """
        Set up kernel module autoloading.
        
        Args:
            context: osbuild context
            kernel_info: Kernel information dictionary
        """
        logger.info("Setting up kernel module autoloading")
        
        # Create /etc/modules-load.d/ directory
        modules_load_dir = os.path.join(context.root, "etc", "modules-load.d")
        os.makedirs(modules_load_dir, exist_ok=True)
        
        # Common modules that should be autoloaded
        common_modules = [
            "loop",
            "overlay",
            "fuse",
            "dm_mod",
            "dm_crypt",
            "ext4",
            "vfat",
            "nls_utf8",
            "nls_cp437",
            "nls_iso8859-1"
        ]
        
        # Write modules to autoload
        modules_file = os.path.join(modules_load_dir, "osbuild.conf")
        with open(modules_file, 'w') as f:
            for module in common_modules:
                f.write(f"{module}\n")
        
        os.chmod(modules_file, 0o644)
        logger.info(f"Created module autoload file: {modules_file}")
    
    def _configure_initramfs_tools(self, context, kernel_info: Dict[str, str]) -> None:
        """
        Configure initramfs-tools for OSTree integration.
        
        Args:
            context: osbuild context
            kernel_info: Kernel information dictionary
        """
        logger.info("Configuring initramfs-tools for OSTree")
        
        # Create initramfs-tools configuration directory
        initramfs_conf_dir = os.path.join(context.root, "etc", "initramfs-tools")
        os.makedirs(initramfs_conf_dir, exist_ok=True)
        
        # Configure initramfs-tools
        config_file = os.path.join(initramfs_conf_dir, "initramfs.conf")
        config_content = [
            "# Debian initramfs-tools configuration for OSTree",
            "# Generated by osbuild debian-kernel-stage",
            "",
            "# Kernel modules to include",
            "MODULES=most",
            "",
            "# Busybox configuration",
            "BUSYBOX=y",
            "",
            "# Include additional tools",
            "KEYMAP=y",
            "COMPCACHE_SIZE=\"\"",
            "COMPRESS=gzip",
            "",
            "# OSTree integration",
            "OSTREE=y",
            "",
            "# Additional hooks",
            "HOOKS=\"ostree\"",
            ""
        ]
        
        with open(config_file, 'w') as f:
            f.write('\n'.join(config_content))
        
        os.chmod(config_file, 0o644)
        logger.info(f"Created initramfs configuration: {config_file}")
        
        # Create OSTree hook for initramfs-tools
        self._create_ostree_hook(context)
    
    def _create_ostree_hook(self, context) -> None:
        """
        Create OSTree hook for initramfs-tools.
        
        Args:
            context: osbuild context
        """
        logger.info("Creating OSTree hook for initramfs-tools")
        
        # Create hooks directory
        hooks_dir = os.path.join(context.root, "etc", "initramfs-tools", "hooks")
        os.makedirs(hooks_dir, exist_ok=True)
        
        # Create OSTree hook script
        hook_script = os.path.join(hooks_dir, "ostree")
        hook_content = [
            "#!/bin/sh",
            "# OSTree hook for initramfs-tools",
            "# This hook ensures OSTree support in the initramfs",
            "",
            "PREREQ=\"\"",
            "",
            "prereqs() {",
            "    echo \"$PREREQ\"",
            "}",
            "",
            "case \"$1\" in",
            "    prereqs)",
            "        prereqs",
            "        exit 0",
            "        ;;",
            "esac",
            "",
            "# Add OSTree binaries to initramfs",
            "if [ -x /usr/bin/ostree ]; then",
            "    copy_exec /usr/bin/ostree /usr/bin/",
            "fi",
            "",
            "# Add OSTree libraries",
            "if [ -d /usr/lib/x86_64-linux-gnu ]; then",
            "    for lib in /usr/lib/x86_64-linux-gnu/libostree*.so*; do",
            "        if [ -f \"$lib\" ]; then",
            "            copy_exec \"$lib\" /usr/lib/x86_64-linux-gnu/",
            "        fi",
            "    done",
            "fi",
            "",
            "# Add OSTree configuration",
            "if [ -d /etc/ostree ]; then",
            "    copy_file /etc/ostree /etc/",
            "fi",
            "",
            "# Add OSTree boot script",
            "cat > \"$DESTDIR/scripts/ostree-boot\" << 'EOF'",
            "#!/bin/sh",
            "# OSTree boot script for initramfs",
            "",
            "ostree_boot() {",
            "    # Mount OSTree repository",
            "    if [ -d /ostree ]; then",
            "        mount -t tmpfs tmpfs /ostree",
            "    fi",
            "",
            "    # Find and mount OSTree deployment",
            "    if [ -f /etc/ostree/remotes.d/ostree.conf ]; then",
            "        ostree admin deploy --os=debian-atomic",
            "    fi",
            "}",
            "",
            "case \"$1\" in",
            "    ostree)",
            "        ostree_boot",
            "        ;;",
            "esac",
            "EOF",
            "",
            "chmod +x \"$DESTDIR/scripts/ostree-boot\"",
            ""
        ]
        
        with open(hook_script, 'w') as f:
            f.write('\n'.join(hook_content))
        
        os.chmod(hook_script, 0o755)
        logger.info(f"Created OSTree hook: {hook_script}")
    
    def _generate_initramfs(self, context, kernel_info: Dict[str, str]) -> None:
        """
        Generate initramfs using update-initramfs.
        
        Args:
            context: osbuild context
            kernel_info: Kernel information dictionary
        """
        logger.info("Generating initramfs")
        
        kernel_version = kernel_info['kernel_version']
        
        # Run update-initramfs to generate initramfs
        cmd = ["update-initramfs", "-c", "-k", kernel_version]
        result = context.run(cmd)
        
        if result.returncode != 0:
            raise RuntimeError(f"Failed to generate initramfs: {result.stderr}")
        
        # Verify initramfs was created
        initramfs_path = os.path.join(context.root, "boot", f"initrd.img-{kernel_version}")
        if not os.path.exists(initramfs_path):
            raise RuntimeError(f"Initramfs not found at expected path: {initramfs_path}")
        
        logger.info(f"Generated initramfs: {initramfs_path}")
    
    def _setup_ostree_integration(self, context, kernel_info: Dict[str, str]) -> None:
        """
        Set up OSTree integration for the kernel.
        
        Args:
            context: osbuild context
            kernel_info: Kernel information dictionary
        """
        logger.info("Setting up OSTree integration")
        
        # Create OSTree configuration directory
        ostree_conf_dir = os.path.join(context.root, "etc", "ostree")
        os.makedirs(ostree_conf_dir, exist_ok=True)
        
        # Create OSTree configuration
        ostree_conf = os.path.join(ostree_conf_dir, "ostree.conf")
        ostree_content = [
            "[core]",
            "repo_mode=bare-user",
            "",
            "[sysroot]",
            "readonly=true",
            "bootloader=grub2",
            ""
        ]
        
        with open(ostree_conf, 'w') as f:
            f.write('\n'.join(ostree_content))
        
        os.chmod(ostree_conf, 0o644)
        logger.info(f"Created OSTree configuration: {ostree_conf}")
        
        # Set up kernel for OSTree boot
        self._setup_ostree_kernel_config(context, kernel_info)
    
    def _setup_ostree_kernel_config(self, context, kernel_info: Dict[str, str]) -> None:
        """
        Set up kernel configuration for OSTree boot.
        
        Args:
            context: osbuild context
            kernel_info: Kernel information dictionary
        """
        logger.info("Setting up kernel configuration for OSTree")
        
        # Create /usr/lib/ostree-boot directory structure
        ostree_boot_dir = os.path.join(context.root, "usr", "lib", "ostree-boot")
        os.makedirs(ostree_boot_dir, exist_ok=True)
        
        kernel_version = kernel_info['kernel_version']
        
        # Copy kernel files to OSTree boot directory
        kernel_path = kernel_info['kernel_path']
        if os.path.exists(kernel_path):
            ostree_kernel = os.path.join(ostree_boot_dir, "vmlinuz")
            shutil.copy2(kernel_path, ostree_kernel)
            logger.info(f"Copied kernel to OSTree boot: {ostree_kernel}")
        
        # Copy initramfs to OSTree boot directory
        initramfs_path = os.path.join(context.root, "boot", f"initrd.img-{kernel_version}")
        if os.path.exists(initramfs_path):
            ostree_initramfs = os.path.join(ostree_boot_dir, "initramfs.img")
            shutil.copy2(initramfs_path, ostree_initramfs)
            logger.info(f"Copied initramfs to OSTree boot: {ostree_initramfs}")
        
        # Copy kernel modules to OSTree boot directory
        modules_path = kernel_info.get('modules_path')
        if modules_path and os.path.exists(modules_path):
            ostree_modules = os.path.join(ostree_boot_dir, "modules")
            if os.path.exists(ostree_modules):
                shutil.rmtree(ostree_modules)
            shutil.copytree(modules_path, ostree_modules)
            logger.info(f"Copied kernel modules to OSTree boot: {ostree_modules}")


def main():
    """
    Main entry point for the Debian Kernel stage.
    
    This function is called by osbuild when executing the stage.
    """
    import sys
    
    # Read options from stdin (osbuild passes options as JSON)
    options = json.load(sys.stdin)
    
    # Create and run the stage
    stage = DebianKernelStage(options)
    
    # Note: In a real osbuild stage, the context would be provided by osbuild
    # For now, this is a placeholder for testing
    class MockContext:
        def __init__(self, root):
            self.root = root
        
        def run(self, cmd):
            # Mock implementation for testing
            logger.info(f"Would run: {' '.join(cmd)}")
            return type('Result', (), {'returncode': 0, 'stdout': '', 'stderr': ''})()
    
    # For testing purposes
    if len(sys.argv) > 1 and sys.argv[1] == '--test':
        context = MockContext('/tmp/test-chroot')
        stage.run(context)
    else:
        # In real osbuild environment, context would be provided
        raise NotImplementedError("This stage must be run within osbuild")


if __name__ == "__main__":
    main()
