"""Main OS builder class for creating Debian-based live ISO images.

This module provides the core OSBuilder class that orchestrates the complete
process of building a bootable Debian-based live ISO. The builder uses a
modular architecture where different components (modules) can be registered
to customize the build process.

The build process consists of several stages:
1. Requirement checking (verifies system tools are available)
2. Directory setup (creates build, chroot, and ISO staging directories)
3. Chroot building (uses debootstrap to create base system)
4. ISO creation (packages chroot into squashfs and creates bootable ISO)

The builder supports caching to speed up rebuilds when only runtime scripts
are modified. The chroot is cached and only rebuilt when configuration changes.

Example:
    Basic usage:
        builder = OSBuilder(config_file="config.yaml")
        iso_path = builder.build(output_name="my-iso.iso")
"""

import shutil
import hashlib
import json
import yaml
from pathlib import Path
from typing import Optional, List
from importlib import resources
from pymergetic.os.builder.config import Config
from pymergetic.os.builder.utils import run_command, ensure_dir, require_command, check_command
from pymergetic.os.builder.modules.base import BaseModule
from pymergetic.os.builder.modules.provision import ProvisionModule
from pymergetic.os.builder.console import (
    console, print_header, print_success, print_step, print_info,
    create_info_table
)


class OSBuilder:
    """Main OS builder for creating Debian-based live ISO images.
    
    This class orchestrates the complete build process from configuration
    to final ISO image. It manages the chroot environment, module execution,
    and ISO generation.
    
    Attributes:
        config (Config): Configuration object containing OS and build settings.
        modules (List): List of registered build modules to execute.
    
    Example:
        Create a builder and build an ISO:
            builder = OSBuilder(config_file="pymergetic_os.yaml")
            iso_path = builder.build()
            print(f"ISO created at: {iso_path}")
    """
    
    def __init__(self, config_file: Optional[str] = None) -> None:
        """Initialize the OS builder with configuration.
        
        Loads configuration from a YAML file, environment variable, or uses
        default values. Automatically registers the base module which handles
        the core Debian installation.
        
        Args:
            config_file: Optional path to YAML configuration file. If not
                        provided, checks PYMERGETIC_OS_BUILD_CONFIG environment
                        variable. If neither is set, uses default configuration.
                        Defaults to None.
        
        Raises:
            FileNotFoundError: If the specified config file does not exist.
            ValueError: If configuration cannot be loaded or is invalid.
        
        Example:
            Initialize with config file:
                builder = OSBuilder(config_file="config.yaml")
            
            Initialize with default config:
                builder = OSBuilder()
        """
        if config_file:
            config_path = Path(config_file)
            self.config = Config.from_yaml(config_path)
            self.config_file_path = config_path
        else:
            self.config = Config.from_yaml()
            self.config_file_path = None
        
        # Initialize cache manifest path
        self._cache_manifest_path = self.config.build_dir / 'cache_manifest.json'
        
        self.modules: List = []
        self._register_base_module()
    
    def _register_base_module(self) -> None:
        """Register the base OS module.
        
        The base module is always registered and handles the core Debian
        installation using debootstrap. It installs essential packages,
        configures the system, and sets up the live-boot infrastructure.
        
        Note:
            This method is called automatically during initialization.
            Additional modules can be registered by extending this method
            or adding them manually to the modules list.
        """
        self.modules.append(BaseModule(self.config))
    
    def _register_provision_module(self) -> None:
        """Register the provisioning module.
        
        The provisioning module handles storage detection and mounting
        on boot. It scans for ext4 volumes with fsencrypt and mounts
        .pymergetic storage according to configuration.
        
        Note:
            This module is registered automatically after the base module.
        """
        self.modules.append(ProvisionModule(self.config))
    
    def build(self, output_name: Optional[str] = None, force: bool = False) -> Path:
        """Build the complete OS ISO image.
        
        Orchestrates the entire build process:
        1. Verifies system requirements (debootstrap, mksquashfs, grub-mkrescue)
        2. Sets up build directories
        3. Builds chroot environment by executing all registered modules (cached if possible)
        4. Creates bootable ISO image with GRUB and ISOLINUX support
        
        Args:
            output_name: Optional custom filename for the output ISO. If not
                        provided, defaults to `pymergetic-os-{release}.iso`
                        where {release} is the Debian release from config.
                        Defaults to None.
            force: If True, forces a full rebuild ignoring cache. Defaults to False.
        
        Returns:
            Path: Absolute path to the created ISO image file.
        
        Raises:
            RuntimeError: If required system commands are missing, if chroot
                        verification fails, or if any build stage fails.
            FileNotFoundError: If required boot files (kernel, initrd) are
                              not found in the chroot.
        
        Example:
            Build with default output name:
                iso_path = builder.build()
            
            Build with custom output name:
                iso_path = builder.build(output_name="custom-iso.iso")
            
            Force full rebuild:
                iso_path = builder.build(force=True)
        """
        # Print header with build information
        print_header("PymergeticOS Builder", "Building Debian-based live ISO")
        
        # Calculate ISO output path upfront
        if output_name:
            iso_filename = output_name
        else:
            iso_filename = f"pymergetic-os-{self.config.release}.iso"
        iso_path = self.config.build_dir.parent / iso_filename
        
        # Display build configuration
        config_table = create_info_table("Build Configuration", {
            "Release": self.config.release,
            "Hostname": self.config.hostname,
            "Architecture": self.config.architecture,
            "Packages": str(len(self.config.packages)),
            "ISO Output": str(iso_path.resolve()),
        })
        console.print(config_table)
        console.print()
        
        # Check requirements
        self._check_requirements()
        
        # Setup directories
        self._setup_directories()
        
        # Check cache and build chroot if needed
        if force or not self._is_chroot_cache_valid():
            print_info("Building chroot (cache invalid or --force specified)...")
            self._build_chroot()
            self._save_cache_manifest()
        else:
            print_success("Using cached chroot (configuration unchanged)")
            # Still verify chroot is valid
            self._verify_chroot()
        
        # Create ISO (always rebuild, it's fast)
        iso_path = self._create_iso(output_name)
        
        # Print completion message with ISO location
        console.print()
        size_mb = iso_path.stat().st_size / (1024 * 1024)
        size_gb = size_mb / 1024 if size_mb >= 1024 else None
        size_display = f"{size_gb:.2f} GB" if size_gb else f"{size_mb:.1f} MB"
        
        iso_info = create_info_table("ISO Image Location", {
            "Path": str(iso_path.resolve()),
            "Size": size_display,
            "Filename": iso_path.name
        })
        console.print(iso_info)
        console.print()
        
        return iso_path
    
    def _check_requirements(self) -> None:
        """Verify that all required system tools are available.
        
        Checks for the presence of essential build tools:
        - debootstrap: Creates Debian chroot environments
        - mksquashfs: Creates compressed squashfs filesystem images
        - grub-mkrescue: Builds bootable ISO images with GRUB
        
        Raises:
            RuntimeError: If any required command is not found in PATH.
        
        Note:
            This method is called automatically at the start of the build
            process. Install missing tools using your system package manager
            (e.g., `apt-get install debootstrap squashfs-tools grub-pc-bin`).
        """
        print_step("CHECK", "Verifying requirements...")
        
        required = ['debootstrap', 'mksquashfs', 'grub-mkrescue']
        for cmd in required:
            require_command(cmd)
        
        print_success("All requirements met")
    
    def _setup_directories(self) -> None:
        """Create and prepare build directory structure.
        
        Sets up the following directories:
        - Build directory: Main output directory for all build artifacts
        - Chroot directory: Location for the Debian chroot environment
        - ISO directory: Staging area for ISO contents before final image creation
        
        Note:
            Directory paths are resolved from the configuration and created
            as absolute paths. The chroot is NOT automatically cleaned here
            to support caching - it's only cleaned if cache is invalid.
        """
        print_step("SETUP", "Creating build directories...")
        
        ensure_dir(self.config.build_dir)
        ensure_dir(self.config.chroot_dir)
        ensure_dir(self.config.iso_dir)
    
    def _is_chroot_cache_valid(self) -> bool:
        """Check if the chroot cache is still valid.
        
        The cache is valid if:
        1. Chroot directory exists and has required files
        2. Cache manifest exists
        3. Configuration hash matches cached hash
        4. Chroot is not corrupted (has required directories)
        
        Returns:
            bool: True if cache is valid, False otherwise.
        """
        # Check if chroot exists
        if not self.config.chroot_dir.exists():
            return False
        
        # Check if chroot has required directories
        required_dirs = ['etc', 'usr', 'bin', 'sbin', 'lib']
        for dir_name in required_dirs:
            if not (self.config.chroot_dir / dir_name).exists():
                return False
        
        # Check if cache manifest exists
        if not self._cache_manifest_path.exists():
            return False
        
        try:
            # Load cache manifest
            with open(self._cache_manifest_path, 'r') as f:
                cache_data = json.load(f)
            
            # Calculate current config hash
            current_hash = self._calculate_config_hash()
            
            # Compare with cached hash
            if cache_data.get('config_hash') != current_hash:
                return False
            
            # Check if chroot is newer than cache (sanity check)
            chroot_mtime = self.config.chroot_dir.stat().st_mtime
            cache_mtime = self._cache_manifest_path.stat().st_mtime
            if chroot_mtime < cache_mtime:
                # Chroot is older than cache, something's wrong
                return False
            
            return True
        except (json.JSONDecodeError, KeyError, OSError):
            return False
    
    def _calculate_config_hash(self) -> str:
        """Calculate a hash of the current configuration.
        
        The hash includes:
        - Config file content (if file-based)
        - Config values (release, hostname, architecture, packages, desktop)
        
        Returns:
            str: SHA256 hash of the configuration.
        """
        import yaml
        
        # Include config file if it exists
        hash_input = []
        
        if self.config_file_path and self.config_file_path.exists():
            with open(self.config_file_path, 'rb') as f:
                hash_input.append(f.read())
        else:
            # Serialize config to YAML for hashing
            config_dict = {
                'os': {
                    'release': self.config.release,
                    'hostname': self.config.hostname,
                    'architecture': self.config.architecture,
                    'packages': sorted(self.config.packages),
                    'desktop': self.config.desktop,
                },
                'build': {
                    'build_dir': self.config.build.build_dir,
                    'chroot_dir': self.config.build.chroot_dir,
                    'iso_dir': self.config.build.iso_dir,
                },
            }
            hash_input.append(yaml.dump(config_dict, sort_keys=True).encode())
        
        # Create hash
        hash_obj = hashlib.sha256()
        for data in hash_input:
            hash_obj.update(data)
        
        return hash_obj.hexdigest()
    
    def _save_cache_manifest(self) -> None:
        """Save cache manifest with current build information.
        
        Stores configuration hash and build metadata for cache validation.
        """
        manifest = {
            'config_hash': self._calculate_config_hash(),
            'release': self.config.release,
            'architecture': self.config.architecture,
            'chroot_dir': str(self.config.chroot_dir),
        }
        
        self._cache_manifest_path.parent.mkdir(parents=True, exist_ok=True)
        with open(self._cache_manifest_path, 'w') as f:
            json.dump(manifest, f, indent=2)
    
    def _build_chroot(self) -> None:
        """Build the complete chroot environment.
        
        Executes all registered modules in sequence to build the chroot.
        Each module performs specific installation and configuration tasks.
        The base module is always executed first and handles the core
        Debian installation.
        
        After all modules complete, verifies that the chroot contains
        all required directories and files.
        
        Raises:
            RuntimeError: If any module fails during installation or if
                         chroot verification fails.
        
        Note:
            Modules are executed in registration order. The base module
            should always be registered first.
        """
        # Clean chroot if it exists (we're doing a fresh build)
        if self.config.chroot_dir.exists():
            print_info("Cleaning existing chroot for fresh build...")
            run_command(['rm', '-rf', str(self.config.chroot_dir)], sudo=True)
        
        print()
        print_step("CHROOT", "Building chroot environment...")
        
        # Run modules
        for module in self.modules:
            print_info(f"Running module: [bold]{module.name}[/bold] - {module.description}")
            module.install(self.config.chroot_dir)
            
            # Copy assets if module has them
            assets = module.get_assets()
            if assets and assets.exists():
                print_info(f"Copying assets from [bold]{module.name}[/bold]...")
                # Copy assets to chroot
                # Implementation depends on module structure
        
        # Verify chroot
        self._verify_chroot()
    
    def _verify_chroot(self) -> None:
        """Verify that the chroot environment is properly configured.
        
        Checks for the presence of essential directories that should exist
        in any valid Debian system:
        - /etc: System configuration files
        - /usr: User programs and data
        - /bin: Essential binaries
        - /sbin: System administration binaries
        - /lib: Shared libraries
        
        Raises:
            RuntimeError: If any required directory is missing, indicating
                         the chroot build failed or is incomplete.
        
        Note:
            This verification is performed after all modules complete their
            installation. Missing directories indicate a build failure.
        """
        print_step("VERIFY", "Verifying chroot...")
        
        required_dirs = ['etc', 'usr', 'bin', 'sbin', 'lib']
        for dir_name in required_dirs:
            if not (self.config.chroot_dir / dir_name).exists():
                raise RuntimeError(f"Chroot missing required directory: {dir_name}")
        
        print_success("Chroot verified")
    
    def _create_iso(self, output_name: Optional[str] = None) -> Path:
        """Create the final bootable ISO image.
        
        Orchestrates the ISO creation process:
        1. Prepares ISO directory structure (live/, boot/grub/, isolinux/)
        2. Copies kernel and initrd from chroot to ISO
        3. Creates compressed squashfs filesystem from chroot
        4. Generates boot configuration (GRUB and ISOLINUX)
        5. Builds final ISO image using grub-mkrescue
        
        Args:
            output_name: Optional custom filename for the ISO. If not provided,
                        defaults to `pymergetic-os-{release}.iso`. Defaults to None.
        
        Returns:
            Path: Absolute path to the created ISO image file.
        
        Raises:
            RuntimeError: If kernel or initrd files are not found, if squashfs
                         creation fails, or if ISO generation fails.
        
        Note:
            The ISO supports both BIOS (ISOLINUX) and UEFI (GRUB) boot modes.
            The squashfs is compressed with xz for optimal size.
        """
        console.print()
        print_step("ISO", "Creating ISO image...")
        
        # Prepare ISO structure
        self._prepare_iso_structure()
        
        # Copy kernel and initrd
        self._copy_boot_files()
        
        # Create squashfs
        self._create_squashfs()
        
        # Create boot configuration
        self._create_boot_config()
        
        # Build ISO
        iso_path = self._build_iso_image(output_name)
        
        return iso_path
    
    def _prepare_iso_structure(self) -> None:
        """Create the ISO directory structure and metadata.
        
        Sets up the required directory layout for a bootable ISO:
        - live/: Contains the live filesystem (squashfs, kernel, initrd)
        - boot/grub/: GRUB boot configuration and modules
        - isolinux/: ISOLINUX boot configuration (for BIOS boot)
        - .disk/: ISO metadata and identification
        
        Also creates a .disk/info file identifying the ISO as PymergeticOS.
        
        Note:
            This structure follows Debian live-boot conventions and is
            compatible with both GRUB (UEFI) and ISOLINUX (BIOS) bootloaders.
        """
        print_info("Preparing ISO structure...")
        
        # Create directories
        ensure_dir(self.config.iso_dir / 'live')
        ensure_dir(self.config.iso_dir / 'boot' / 'grub')
        ensure_dir(self.config.iso_dir / 'isolinux')
        ensure_dir(self.config.iso_dir / '.disk')
        
        # Copy .disk/info from resources
        disk_info = resources.files('pymergetic.os.builder.resources').joinpath('.disk', 'info')
        (self.config.iso_dir / '.disk' / 'info').write_text(disk_info.read_text())
    
    def _copy_boot_files(self) -> None:
        """Copy kernel and initrd images from chroot to ISO.
        
        Locates the kernel (vmlinuz-*) and matching initrd (initrd.img-*)
        files in the chroot's /boot directory and copies them to the ISO's
        live/ directory with standardized names (vmlinuz, initrd.gz).
        
        The kernel version is extracted from the kernel filename to find
        the matching initrd. If no version-specific initrd is found,
        falls back to any initrd file.
        
        Raises:
            RuntimeError: If no kernel files are found in /boot/vmlinuz-*
                         or if no initrd files are found.
        
        Note:
            Requires sudo privileges to copy files from the chroot.
            The files are renamed to standard names expected by the bootloader.
        """
        print_info("Copying boot files...")
        
        chroot_boot = self.config.chroot_dir / 'boot'
        
        # Find kernel
        kernel_files = list(chroot_boot.glob('vmlinuz-*'))
        if not kernel_files:
            raise RuntimeError("No kernel found in chroot")
        
        kernel_file = kernel_files[0]
        run_command([
            'cp', str(kernel_file),
            str(self.config.iso_dir / 'live' / 'vmlinuz'),
        ], sudo=True)
        
        # Find initrd
        kernel_version = kernel_file.name.replace('vmlinuz-', '')
        initrd_files = list(chroot_boot.glob(f'initrd*{kernel_version}*'))
        if not initrd_files:
            initrd_files = list(chroot_boot.glob('initrd*'))
        
        if not initrd_files:
            raise RuntimeError("No initrd found in chroot")
        
        initrd_file = initrd_files[0]
        run_command([
            'cp', str(initrd_file),
            str(self.config.iso_dir / 'live' / 'initrd.gz'),
        ], sudo=True)
    
    def _create_squashfs(self) -> None:
        """Create a compressed squashfs filesystem from the chroot.
        
        Packages the entire chroot directory into a compressed squashfs
        image using xz compression. Excludes temporary and virtual
        filesystem directories (boot, dev, proc, sys, tmp, run, mnt, media).
        
        The resulting filesystem.squashfs file is placed in the ISO's
        live/ directory and contains the complete root filesystem that
        will be mounted by live-boot.
        
        Raises:
            RuntimeError: If mksquashfs command fails or if the chroot
                         directory is not accessible.
        
        Note:
            Requires sudo privileges. The compression uses xz for optimal
            size. The resulting file size is printed in MB for verification.
        """
        print_info("Creating squashfs filesystem...")
        
        squashfs_path = self.config.iso_dir / 'live' / 'filesystem.squashfs'
        
        run_command([
            'mksquashfs',
            str(self.config.chroot_dir),
            str(squashfs_path),
            '-comp', 'xz',
            '-e', 'boot',
            '-e', 'dev/*',
            '-e', 'proc/*',
            '-e', 'sys/*',
            '-e', 'tmp/*',
            '-e', 'run/*',
            '-e', 'mnt',
            '-e', 'media',
            '-e', 'lost+found',
        ], sudo=True)
        
        size_mb = squashfs_path.stat().st_size / (1024 * 1024)
        print_success(f"Squashfs created: [bold]{size_mb:.1f} MB[/bold]")
    
    def _create_boot_config(self) -> None:
        """Generate boot configuration files for GRUB and ISOLINUX.
        
        Creates bootloader configuration files:
        - boot/grub/grub.cfg: GRUB configuration for UEFI boot
        - isolinux/isolinux.cfg: ISOLINUX configuration for BIOS boot
        
        Both configurations set up a boot menu with a single entry for
        "PymergeticOS Live" that boots the kernel with live-boot parameters.
        
        If isolinux.bin is available on the system, it is also copied to
        the isolinux/ directory for BIOS boot support.
        
        Note:
            The boot configurations use standard Debian live-boot parameters:
            - boot=live: Enables live mode
            - quiet: Reduces kernel message verbosity
            - splash: Shows boot splash screen (if available)
        """
        print_info("Creating boot configuration...")
        
        # GRUB config - load from resources
        grub_cfg_resource = resources.files('pymergetic.os.builder.resources').joinpath('boot', 'grub', 'grub.cfg')
        grub_cfg = self.config.iso_dir / 'boot' / 'grub' / 'grub.cfg'
        grub_cfg.write_text(grub_cfg_resource.read_text())
        
        # ISOLINUX config - load from resources
        isolinux_cfg_resource = resources.files('pymergetic.os.builder.resources').joinpath('isolinux', 'isolinux.cfg')
        isolinux_cfg = self.config.iso_dir / 'isolinux' / 'isolinux.cfg'
        isolinux_cfg.write_text(isolinux_cfg_resource.read_text())
        
        # Copy ISOLINUX files if available
        if check_command('isolinux.bin'):
            isolinux_bin = shutil.which('isolinux.bin')
            if isolinux_bin:
                run_command([
                    'cp', isolinux_bin,
                    str(self.config.iso_dir / 'isolinux' / 'isolinux.bin'),
                ], sudo=True)
    
    def _build_iso_image(self, output_name: Optional[str] = None) -> Path:
        """Build the final bootable ISO image using grub-mkrescue.
        
        Uses grub-mkrescue to create a hybrid ISO image that supports both
        BIOS and UEFI boot modes. The ISO includes all files from the ISO
        staging directory, including the squashfs, boot files, and bootloader
        configurations.
        
        Args:
            output_name: Optional custom filename. If not provided, defaults
                         to `pymergetic-os-{release}.iso`. Defaults to None.
        
        Returns:
            Path: Absolute path to the created ISO image file.
        
        Raises:
            RuntimeError: If grub-mkrescue fails or if the ISO staging
                         directory is incomplete.
        
        Note:
            Requires sudo privileges. The ISO is created in the build directory's
            parent directory. The final file size is printed in MB.
        """
        print_info("Building ISO image...")
        
        if not output_name:
            output_name = f"pymergetic-os-{self.config.release}.iso"
        
        iso_path = self.config.build_dir.parent / output_name
        
        # Use grub-mkrescue
        run_command([
            'grub-mkrescue',
            '-o', str(iso_path),
            str(self.config.iso_dir),
        ], sudo=True)
        
        return iso_path

