#!/usr/bin/env python3
"""
Test script for GRUB installation and configuration in Debian GRUB Stage

This script validates the GRUB installation process, including:
- GRUB installation to boot partition
- GRUB configuration file generation
- Bootloader setup validation
- Error handling and edge cases

Author: Debian bootc-image-builder team
"""

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

# Add the parent directory to the path to import the GRUB stage
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from debian_grub_stage import DebianGrubStage

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(levelname)s:%(name)s:%(message)s')
logger = logging.getLogger(__name__)


class GrubInstallationTest:
    """
    Test suite for GRUB installation and configuration.
    """
    
    def __init__(self):
        self.test_results = {}
        self.temp_dirs = []
    
    def run_all_tests(self) -> bool:
        """Run all GRUB installation tests."""
        logger.info("Starting GRUB installation tests")
        
        tests = [
            ("test_grub_directory_setup", self.test_grub_directory_setup),
            ("test_grub_config_generation", self.test_grub_config_generation),
            ("test_grub_installation_process", self.test_grub_installation_process),
            ("test_bootloader_validation", self.test_bootloader_validation),
            ("test_error_handling", self.test_error_handling),
            ("test_grub_customization", self.test_grub_customization),
        ]
        
        all_passed = True
        for test_name, test_func in tests:
            logger.info(f"Running test: {test_name}")
            try:
                result = test_func()
                self.test_results[test_name] = result
                if result:
                    logger.info(f"✓ {test_name} PASSED")
                else:
                    logger.error(f"✗ {test_name} FAILED")
                    all_passed = False
            except Exception as e:
                logger.error(f"✗ {test_name} ERROR: {e}")
                self.test_results[test_name] = False
                all_passed = False
        
        self.print_summary()
        return all_passed
    
    def print_summary(self) -> None:
        """Print test summary."""
        logger.info("=" * 50)
        logger.info("GRUB INSTALLATION TEST SUMMARY")
        logger.info("=" * 50)
        
        passed = sum(1 for result in self.test_results.values() if result)
        total = len(self.test_results)
        
        for test_name, result in self.test_results.items():
            status = "PASSED" if result else "FAILED"
            logger.info(f"{test_name}: {status}")
        
        logger.info(f"Overall: {passed}/{total} tests passed")
        logger.info("=" * 50)
        
        if passed == total:
            logger.info("All GRUB installation tests PASSED")
        else:
            logger.error("Some GRUB installation tests FAILED")
    
    def cleanup(self) -> None:
        """Clean up temporary directories."""
        for temp_dir in self.temp_dirs:
            try:
                shutil.rmtree(temp_dir)
            except Exception as e:
                logger.warning(f"Failed to clean up {temp_dir}: {e}")
    
    def test_grub_directory_setup(self) -> bool:
        """Test GRUB directory setup and structure."""
        test_dir = tempfile.mkdtemp()
        self.temp_dirs.append(test_dir)
        
        try:
            # Create mock context
            context = MockOsbuildContext(test_dir)
            
            # Initialize GRUB stage
            options = {
                'ostree_integration': True,
                'uefi': True,
                'secure_boot': False,
                'timeout': 5,
                'default_entry': 0
            }
            grub_stage = DebianGrubStage(options)
            
            # Test directory setup
            grub_stage._setup_grub_directories(context)
            
            # Verify directories were created
            expected_dirs = [
                f"{test_dir}/boot/grub",
                f"{test_dir}/etc/default",
                f"{test_dir}/etc/grub.d"
            ]
            
            for expected_dir in expected_dirs:
                if not os.path.exists(expected_dir):
                    logger.error(f"Expected directory not created: {expected_dir}")
                    return False
            
            # Verify GRUB configuration file
            grub_default = f"{test_dir}/etc/default/grub"
            if not os.path.exists(grub_default):
                logger.error(f"GRUB default configuration not created: {grub_default}")
                return False
            
            # Check GRUB default configuration content
            with open(grub_default, 'r') as f:
                content = f.read()
                if 'GRUB_TIMEOUT=5' not in content:
                    logger.error("GRUB timeout not set correctly")
                    return False
                if 'GRUB_DEFAULT=0' not in content:
                    logger.error("GRUB default entry not set correctly")
                    return False
            
            logger.info("GRUB directory setup successful")
            return True
            
        except Exception as e:
            logger.error(f"GRUB directory setup test failed: {e}")
            return False
    
    def test_grub_config_generation(self) -> bool:
        """Test GRUB configuration file generation."""
        test_dir = tempfile.mkdtemp()
        self.temp_dirs.append(test_dir)
        
        try:
            # Create mock context
            context = MockOsbuildContext(test_dir)
            
            # Set up basic directory structure
            os.makedirs(f"{test_dir}/boot/grub", exist_ok=True)
            os.makedirs(f"{test_dir}/etc/default", exist_ok=True)
            os.makedirs(f"{test_dir}/etc/grub.d", exist_ok=True)
            
            # Create mock kernel and initramfs files
            kernel_path = f"{test_dir}/boot/vmlinuz-6.1.0-13-amd64"
            initramfs_path = f"{test_dir}/boot/initrd.img-6.1.0-13-amd64"
            os.makedirs(os.path.dirname(kernel_path), exist_ok=True)
            with open(kernel_path, 'w') as f:
                f.write("mock kernel")
            with open(initramfs_path, 'w') as f:
                f.write("mock initramfs")
            
            # Initialize GRUB stage
            options = {
                'ostree_integration': True,
                'uefi': True,
                'secure_boot': False,
                'timeout': 5,
                'default_entry': 0,
                'kernel_path': kernel_path,
                'initramfs_path': initramfs_path
            }
            grub_stage = DebianGrubStage(options)
            
            # Test GRUB configuration generation
            grub_stage._generate_grub_config(context)
            
            # Verify GRUB configuration was created
            grub_cfg = f"{test_dir}/boot/grub/grub.cfg"
            if not os.path.exists(grub_cfg):
                logger.error(f"GRUB configuration not created: {grub_cfg}")
                return False
            
            # Check GRUB configuration content
            with open(grub_cfg, 'r') as f:
                content = f.read()
                if 'set timeout=5' not in content:
                    logger.error("GRUB timeout not set in configuration")
                    return False
                if 'linux' not in content or 'initrd' not in content:
                    logger.error("GRUB configuration missing kernel/initramfs entries")
                    return False
            
            logger.info("GRUB configuration generation successful")
            return True
            
        except Exception as e:
            logger.error(f"GRUB configuration generation test failed: {e}")
            return False
    
    def test_grub_installation_process(self) -> bool:
        """Test GRUB installation to boot partition."""
        test_dir = tempfile.mkdtemp()
        self.temp_dirs.append(test_dir)
        
        try:
            # Create mock context
            context = MockOsbuildContext(test_dir)
            
            # Set up basic directory structure
            os.makedirs(f"{test_dir}/boot/grub", exist_ok=True)
            os.makedirs(f"{test_dir}/etc/default", exist_ok=True)
            os.makedirs(f"{test_dir}/etc/grub.d", exist_ok=True)
            
            # Initialize GRUB stage
            options = {
                'ostree_integration': True,
                'uefi': True,
                'secure_boot': False,
                'timeout': 5,
                'default_entry': 0
            }
            grub_stage = DebianGrubStage(options)
            
            # Test GRUB installation
            grub_stage._install_grub(context)
            
            # Verify GRUB was installed (check for core.img or similar)
            grub_core = f"{test_dir}/boot/grub/i386-pc/core.img"
            grub_efi = f"{test_dir}/boot/efi/EFI/debian/grubx64.efi"
            
            # In a mock environment, we can't actually install GRUB
            # So we verify the installation command was called
            if not hasattr(context, 'last_command') or 'grub-install' not in context.last_command:
                logger.error("GRUB installation command not executed")
                return False
            
            logger.info("GRUB installation process test completed (mock environment)")
            return True
            
        except Exception as e:
            logger.error(f"GRUB installation process test failed: {e}")
            return False
    
    def test_bootloader_validation(self) -> bool:
        """Test bootloader validation and verification."""
        test_dir = tempfile.mkdtemp()
        self.temp_dirs.append(test_dir)
        
        try:
            # Create mock context
            context = MockOsbuildContext(test_dir)
            
            # Set up basic directory structure
            os.makedirs(f"{test_dir}/boot/grub", exist_ok=True)
            os.makedirs(f"{test_dir}/etc/default", exist_ok=True)
            os.makedirs(f"{test_dir}/etc/grub.d", exist_ok=True)
            
            # Create mock GRUB files
            grub_cfg = f"{test_dir}/boot/grub/grub.cfg"
            with open(grub_cfg, 'w') as f:
                f.write("set timeout=5\n")
                f.write("menuentry 'Debian' {\n")
                f.write("    linux /boot/vmlinuz root=/dev/sda1\n")
                f.write("    initrd /boot/initrd.img\n")
                f.write("}\n")
            
            # Initialize GRUB stage
            options = {
                'ostree_integration': True,
                'uefi': True,
                'secure_boot': False,
                'timeout': 5,
                'default_entry': 0
            }
            grub_stage = DebianGrubStage(options)
            
            # Test bootloader validation
            # In a real environment, this would verify GRUB installation
            # For mock testing, we verify the configuration is valid
            
            if not os.path.exists(grub_cfg):
                logger.error("GRUB configuration file not found for validation")
                return False
            
            with open(grub_cfg, 'r') as f:
                content = f.read()
                if 'menuentry' not in content:
                    logger.error("GRUB configuration missing menu entries")
                    return False
                if 'linux' not in content:
                    logger.error("GRUB configuration missing kernel entries")
                    return False
            
            logger.info("Bootloader validation successful")
            return True
            
        except Exception as e:
            logger.error(f"Bootloader validation test failed: {e}")
            return False
    
    def test_error_handling(self) -> bool:
        """Test error handling in GRUB installation."""
        test_dir = tempfile.mkdtemp()
        self.temp_dirs.append(test_dir)
        
        try:
            # Create mock context that simulates failures
            context = MockOsbuildContext(test_dir)
            context.simulate_failure = True
            
            # Initialize GRUB stage
            options = {
                'ostree_integration': True,
                'uefi': True,
                'secure_boot': False,
                'timeout': 5,
                'default_entry': 0
            }
            grub_stage = DebianGrubStage(options)
            
            # Test error handling - the setup should succeed even with simulate_failure
            # because it only affects the run() method, not directory creation
            grub_stage._setup_grub_directories(context)
            
            # Test error handling in a method that uses context.run()
            try:
                grub_stage._install_grub(context)
                logger.error("Expected error was not raised")
                return False
            except Exception as e:
                logger.info(f"Expected error caught: {e}")
                return True
            
        except Exception as e:
            logger.error(f"Error handling test failed: {e}")
            return False
    
    def test_grub_customization(self) -> bool:
        """Test GRUB customization options."""
        test_dir = tempfile.mkdtemp()
        self.temp_dirs.append(test_dir)
        
        try:
            # Create mock context
            context = MockOsbuildContext(test_dir)
            
            # Set up basic directory structure
            os.makedirs(f"{test_dir}/boot/grub", exist_ok=True)
            os.makedirs(f"{test_dir}/etc/default", exist_ok=True)
            os.makedirs(f"{test_dir}/etc/grub.d", exist_ok=True)
            
            # Initialize GRUB stage with custom options
            options = {
                'ostree_integration': True,
                'uefi': True,
                'secure_boot': False,
                'timeout': 10,  # Custom timeout
                'default_entry': 1,  # Custom default entry
                'kernel_params': 'console=ttyS0,115200'  # Custom kernel params
            }
            grub_stage = DebianGrubStage(options)
            
            # Test GRUB customization
            grub_stage._setup_grub_directories(context)
            grub_stage._generate_grub_config(context)
            
            # Verify custom settings were applied
            grub_default = f"{test_dir}/etc/default/grub"
            if os.path.exists(grub_default):
                with open(grub_default, 'r') as f:
                    content = f.read()
                    if 'GRUB_TIMEOUT=10' not in content:
                        logger.error("Custom timeout not applied")
                        return False
                    if 'GRUB_DEFAULT=1' not in content:
                        logger.error("Custom default entry not applied")
                        return False
            
            logger.info("GRUB customization successful")
            return True
            
        except Exception as e:
            logger.error(f"GRUB customization test failed: {e}")
            return False


class MockOsbuildContext:
    """Mock osbuild context for testing."""
    
    def __init__(self, root_path: str):
        self.root_path = root_path
        self.root = root_path
        self.last_command = None
        self.simulate_failure = False
    
    def run(self, cmd: List[str], **kwargs) -> subprocess.CompletedProcess:
        """Mock run method that simulates command execution."""
        self.last_command = ' '.join(cmd)
        
        if self.simulate_failure:
            raise subprocess.CalledProcessError(1, cmd, "Mock failure")
        
        # Handle specific commands
        if 'grub-install' in cmd:
            logger.info(f"Mock GRUB installation: {cmd}")
            # Create mock GRUB files
            if '--target=i386-pc' in ' '.join(cmd):
                os.makedirs(f"{self.root_path}/boot/grub/i386-pc", exist_ok=True)
                with open(f"{self.root_path}/boot/grub/i386-pc/core.img", 'w') as f:
                    f.write("mock grub core")
            elif '--target=x86_64-efi' in ' '.join(cmd):
                os.makedirs(f"{self.root_path}/boot/efi/EFI/debian", exist_ok=True)
                with open(f"{self.root_path}/boot/efi/EFI/debian/grubx64.efi", 'w') as f:
                    f.write("mock grub efi")
        
        elif 'update-grub' in cmd:
            logger.info(f"Mock GRUB update: {cmd}")
            # Create mock grub.cfg
            os.makedirs(f"{self.root_path}/boot/grub", exist_ok=True)
            with open(f"{self.root_path}/boot/grub/grub.cfg", 'w') as f:
                f.write("set timeout=5\n")
                f.write("menuentry 'Debian' {\n")
                f.write("    linux /boot/vmlinuz root=/dev/sda1\n")
                f.write("    initrd /boot/initrd.img\n")
                f.write("}\n")
        
        return subprocess.CompletedProcess(cmd, 0, "", "")


def main():
    """Main test execution function."""
    test_suite = GrubInstallationTest()
    
    try:
        success = test_suite.run_all_tests()
        return 0 if success else 1
    finally:
        test_suite.cleanup()


if __name__ == "__main__":
    sys.exit(main())
