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

This stage handles Debian package installation using apt/dpkg within an osbuild chroot.
It replaces the DNF 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 sys
from typing import Dict, List, Optional, Any
import logging

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


class AptStage:
    """
    osbuild stage for Debian package management using apt/dpkg.
    
    This stage handles:
    - Repository configuration
    - Package dependency resolution
    - Package installation
    - Cache management
    - OSTree integration considerations
    """
    
    def __init__(self, options: Dict[str, Any]):
        """
        Initialize the APT stage with configuration options.
        
        Args:
            options: Dictionary containing stage configuration
                - packages: List of packages to install
                - repos: List of repository configurations
                - release: Debian release (e.g., 'trixie', 'bookworm')
                - arch: Target architecture (e.g., 'amd64')
                - exclude_packages: List of packages to exclude
                - install_weak_deps: Whether to install weak dependencies
        """
        self.options = options
        self.packages = options.get('packages', [])
        self.repos = options.get('repos', [])
        self.release = options.get('release', 'trixie')
        self.arch = options.get('arch', 'amd64')
        self.exclude_packages = options.get('exclude_packages', [])
        self.install_weak_deps = options.get('install_weak_deps', True)
        
        # Validate required options
        if not self.packages:
            raise ValueError("No packages specified for installation")
        
        logger.info(f"APT Stage initialized for {self.release} ({self.arch})")
        logger.info(f"Packages to install: {self.packages}")
    
    def run(self, context) -> None:
        """
        Execute the APT stage within the osbuild context.
        
        Args:
            context: osbuild context providing chroot access
        """
        logger.info("Starting APT stage execution")
        
        try:
            # Step 1: Set up APT configuration in chroot
            self._setup_apt_config(context)
            
            # Step 2: Configure repositories
            self._configure_repositories(context)
            
            # Step 3: Update package lists
            self._update_package_lists(context)
            
            # Step 4: Install packages with dependency resolution
            self._install_packages(context)
            
            # Step 5: Clean up APT cache
            self._cleanup_cache(context)
            
            logger.info("APT stage completed successfully")
            
        except Exception as e:
            logger.error(f"APT stage failed: {e}")
            raise
    
    def _setup_apt_config(self, context) -> None:
        """
        Set up APT configuration in the chroot environment.
        
        Args:
            context: osbuild context
        """
        logger.info("Setting up APT configuration")
        
        # Create APT configuration directory
        apt_dir = os.path.join(context.root, "etc", "apt")
        os.makedirs(apt_dir, exist_ok=True)
        
        # Create sources.list with default Debian repositories
        sources_list = f"""deb http://deb.debian.org/debian {self.release} main contrib non-free
deb http://deb.debian.org/debian-security {self.release}-security main contrib non-free
deb http://deb.debian.org/debian {self.release}-updates main contrib non-free
deb http://deb.debian.org/debian {self.release}-backports main contrib non-free"""
        
        sources_path = os.path.join(apt_dir, "sources.list")
        with open(sources_path, 'w') as f:
            f.write(sources_list)
        
        # Create apt.conf.d directory
        apt_conf_dir = os.path.join(apt_dir, "apt.conf.d")
        os.makedirs(apt_conf_dir, exist_ok=True)
        
        # Create basic apt configuration
        apt_conf = """APT::Get::AllowUnauthenticated "true";
APT::Get::Assume-Yes "true";
APT::Get::Show-Upgraded "true";
APT::Install-Recommends "false";
APT::Install-Suggests "false";"""
        
        apt_conf_path = os.path.join(apt_conf_dir, "99defaults")
        with open(apt_conf_path, 'w') as f:
            f.write(apt_conf)
        
        logger.info("APT configuration setup completed")
    
    def _configure_repositories(self, context) -> None:
        """
        Configure additional repositories if specified.
        
        Args:
            context: osbuild context
        """
        if not self.repos:
            logger.info("No additional repositories to configure")
            return
        
        logger.info(f"Configuring {len(self.repos)} additional repositories")
        
        # For now, we'll use the default repositories
        # In a full implementation, we would handle custom repository configurations
        logger.info("Using default Debian repositories")
    
    def _update_package_lists(self, context) -> None:
        """
        Update package lists using apt.
        
        Args:
            context: osbuild context
        """
        logger.info("Updating package lists")
        
        # Set environment variables for apt
        env = os.environ.copy()
        env['APT_CONFIG'] = os.path.join(context.root, "etc", "apt", "apt.conf")
        env['APT_STATE_DIR'] = os.path.join(context.root, "var", "lib", "apt")
        env['APT_CACHE_DIR'] = os.path.join(context.root, "var", "cache", "apt")
        
        # Create necessary directories
        apt_state_dir = os.path.join(context.root, "var", "lib", "apt")
        apt_cache_dir = os.path.join(context.root, "var", "cache", "apt")
        os.makedirs(apt_state_dir, exist_ok=True)
        os.makedirs(apt_cache_dir, exist_ok=True)
        
        # Run apt update
        cmd = ["apt", "update"]
        result = context.run(cmd, env=env)
        
        if result.returncode != 0:
            logger.warning(f"apt update failed: {result.stderr}")
            logger.info("Continuing with package installation...")
        else:
            logger.info("Package lists updated successfully")
    
    def _install_packages(self, context) -> None:
        """
        Install packages with dependency resolution.
        
        Args:
            context: osbuild context
        """
        logger.info(f"Installing {len(self.packages)} packages")
        
        # Set environment variables for apt
        env = os.environ.copy()
        env['APT_CONFIG'] = os.path.join(context.root, "etc", "apt", "apt.conf")
        env['APT_STATE_DIR'] = os.path.join(context.root, "var", "lib", "apt")
        env['APT_CACHE_DIR'] = os.path.join(context.root, "var", "cache", "apt")
        
        # Filter out excluded packages
        packages_to_install = [pkg for pkg in self.packages if pkg not in self.exclude_packages]
        
        if not packages_to_install:
            logger.info("No packages to install after filtering")
            return
        
        # Build apt-get install command
        cmd = ["apt-get", "install", "-y"]
        
        if not self.install_weak_deps:
            cmd.append("--no-install-recommends")
        
        cmd.extend(packages_to_install)
        
        logger.info(f"Running command: {' '.join(cmd)}")
        
        # Execute package installation
        result = context.run(cmd, env=env)
        
        if result.returncode != 0:
            logger.error(f"Package installation failed: {result.stderr}")
            logger.error(f"Command executed: {' '.join(cmd)}")
            
            # Try to get more detailed error information
            self._log_apt_errors(context)
            
            raise RuntimeError("Package installation failed")
        
        logger.info("Package installation completed successfully")
    
    def _cleanup_cache(self, context) -> None:
        """
        Clean up APT cache to reduce image size.
        
        Args:
            context: osbuild context
        """
        logger.info("Cleaning up APT cache")
        
        # Set environment variables for apt
        env = os.environ.copy()
        env['APT_CONFIG'] = os.path.join(context.root, "etc", "apt", "apt.conf")
        env['APT_STATE_DIR'] = os.path.join(context.root, "var", "lib", "apt")
        env['APT_CACHE_DIR'] = os.path.join(context.root, "var", "cache", "apt")
        
        # Clean package cache
        cmd = ["apt-get", "clean"]
        result = context.run(cmd, env=env)
        
        if result.returncode != 0:
            logger.warning(f"Cache cleanup failed: {result.stderr}")
        else:
            logger.info("APT cache cleaned successfully")
        
        # Remove package lists
        cmd = ["rm", "-rf", "/var/lib/apt/lists/*"]
        result = context.run(cmd, env=env)
        
        if result.returncode != 0:
            logger.warning(f"Package list cleanup failed: {result.stderr}")
        else:
            logger.info("Package lists removed successfully")
    
    def _log_apt_errors(self, context) -> None:
        """
        Log detailed APT error information for debugging.
        
        Args:
            context: osbuild context
        """
        logger.info("Collecting detailed APT error information")
        
        # Set environment variables for apt
        env = os.environ.copy()
        env['APT_CONFIG'] = os.path.join(context.root, "etc", "apt", "apt.conf")
        env['APT_STATE_DIR'] = os.path.join(context.root, "var", "lib", "apt")
        env['APT_CACHE_DIR'] = os.path.join(context.root, "var", "cache", "apt")
        
        # Check APT status
        cmd = ["apt-get", "check"]
        result = context.run(cmd, env=env)
        logger.info(f"APT check result: {result.stdout}")
        
        # Check for broken packages
        cmd = ["dpkg", "--audit"]
        result = context.run(cmd, env=env)
        if result.stdout:
            logger.error(f"Broken packages detected: {result.stdout}")
        
        # Check package status
        cmd = ["dpkg", "-l"]
        result = context.run(cmd, env=env)
        logger.debug(f"Package status: {result.stdout}")


def main():
    """
    Main entry point for the APT stage.
    
    This function is called by osbuild when executing the stage.
    """
    # Read options from stdin (osbuild passes options as JSON)
    options = json.load(sys.stdin)
    
    # Create and run the stage
    stage = AptStage(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, env=None):
            # 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()
