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

This script validates UEFI boot setup, including:
- UEFI boot entry creation
- EFI partition configuration
- Secure Boot integration
- UEFI firmware interaction

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 UEFIBootTest:
    """
    Test suite for UEFI boot configuration.
    """
    
    def __init__(self):
        self.test_results = {}
        self.temp_dirs = []
    
    def run_all_tests(self) -> bool:
        """Run all UEFI boot tests."""
        logger.info("Starting UEFI boot configuration tests")
        
        tests = [
            ("test_uefi_boot_entry_creation", self.test_uefi_boot_entry_creation),
            ("test_efi_partition_setup", self.test_efi_partition_setup),
            ("test_grub_efi_installation", self.test_grub_efi_installation),
            ("test_uefi_boot_order", self.test_uefi_boot_order),
            ("test_secure_boot_integration", self.test_secure_boot_integration),
            ("test_uefi_fallback_boot", self.test_uefi_fallback_boot),
        ]
        
        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("UEFI BOOT CONFIGURATION 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 UEFI boot configuration tests PASSED")
        else:
            logger.error("Some UEFI boot configuration 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_uefi_boot_entry_creation(self) -> bool:
        """Test UEFI boot entry creation."""
        test_dir = tempfile.mkdtemp()
        self.temp_dirs.append(test_dir)
        
        try:
            # Create mock context
            context = MockOsbuildContext(test_dir)
            
            # Set up EFI directory structure
            os.makedirs(f"{test_dir}/boot/efi/EFI/debian", exist_ok=True)
            os.makedirs(f"{test_dir}/boot/grub", exist_ok=True)
            os.makedirs(f"{test_dir}/etc/default", exist_ok=True)
            
            # Create mock EFI files
            grub_efi = f"{test_dir}/boot/efi/EFI/debian/grubx64.efi"
            with open(grub_efi, 'w') as f:
                f.write("mock grub efi")
            
            # Initialize GRUB stage with UEFI enabled
            options = {
                'ostree_integration': True,
                'uefi': True,
                'secure_boot': False,
                'timeout': 5,
                'default_entry': 0
            }
            grub_stage = DebianGrubStage(options)
            
            # Test UEFI boot entry creation
            grub_stage._configure_uefi_boot(context)
            
            # Verify UEFI boot entry was created
            if not hasattr(context, 'last_command') or 'efibootmgr' not in context.last_command:
                logger.error("UEFI boot entry creation command not executed")
                return False
            
            # Verify EFI files exist
            if not os.path.exists(grub_efi):
                logger.error(f"GRUB EFI file not created: {grub_efi}")
                return False
            
            logger.info("UEFI boot entry creation successful")
            return True
            
        except Exception as e:
            logger.error(f"UEFI boot entry creation test failed: {e}")
            return False
    
    def test_efi_partition_setup(self) -> bool:
        """Test EFI partition setup and configuration."""
        test_dir = tempfile.mkdtemp()
        self.temp_dirs.append(test_dir)
        
        try:
            # Create mock context
            context = MockOsbuildContext(test_dir)
            
            # Set up EFI directory structure
            os.makedirs(f"{test_dir}/boot/efi", exist_ok=True)
            os.makedirs(f"{test_dir}/boot/efi/EFI", exist_ok=True)
            os.makedirs(f"{test_dir}/boot/efi/EFI/debian", exist_ok=True)
            
            # Create mock EFI partition files
            efi_partition = f"{test_dir}/boot/efi"
            
            # Initialize GRUB stage
            options = {
                'ostree_integration': True,
                'uefi': True,
                'secure_boot': False,
                'timeout': 5,
                'default_entry': 0
            }
            grub_stage = DebianGrubStage(options)
            
            # Test EFI partition setup
            # In a real environment, this would verify EFI partition mounting
            # For mock testing, we verify the directory structure
            
            if not os.path.exists(efi_partition):
                logger.error(f"EFI partition directory not created: {efi_partition}")
                return False
            
            # Verify EFI directory structure
            expected_dirs = [
                f"{test_dir}/boot/efi/EFI",
                f"{test_dir}/boot/efi/EFI/debian"
            ]
            
            for expected_dir in expected_dirs:
                if not os.path.exists(expected_dir):
                    logger.error(f"Expected EFI directory not created: {expected_dir}")
                    return False
            
            logger.info("EFI partition setup successful")
            return True
            
        except Exception as e:
            logger.error(f"EFI partition setup test failed: {e}")
            return False
    
    def test_grub_efi_installation(self) -> bool:
        """Test GRUB EFI installation."""
        test_dir = tempfile.mkdtemp()
        self.temp_dirs.append(test_dir)
        
        try:
            # Create mock context
            context = MockOsbuildContext(test_dir)
            
            # Set up directory structure
            os.makedirs(f"{test_dir}/boot/efi/EFI/debian", exist_ok=True)
            os.makedirs(f"{test_dir}/boot/grub", exist_ok=True)
            os.makedirs(f"{test_dir}/etc/default", 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 EFI installation
            grub_stage._install_grub(context)
            
            # Verify GRUB EFI installation command was called
            if not hasattr(context, 'last_command') or 'grub-install' not in context.last_command:
                logger.error("GRUB EFI installation command not executed")
                return False
            
            # Verify EFI target was specified
            if '--target=x86_64-efi' not in context.last_command:
                logger.error("GRUB EFI installation not using correct target")
                return False
            
            # Verify EFI file was created
            grub_efi = f"{test_dir}/boot/efi/EFI/debian/grubx64.efi"
            if not os.path.exists(grub_efi):
                logger.error(f"GRUB EFI file not created: {grub_efi}")
                return False
            
            logger.info("GRUB EFI installation successful")
            return True
            
        except Exception as e:
            logger.error(f"GRUB EFI installation test failed: {e}")
            return False
    
    def test_uefi_boot_order(self) -> bool:
        """Test UEFI boot order configuration."""
        test_dir = tempfile.mkdtemp()
        self.temp_dirs.append(test_dir)
        
        try:
            # Create mock context
            context = MockOsbuildContext(test_dir)
            
            # Set up directory structure
            os.makedirs(f"{test_dir}/boot/efi/EFI/debian", exist_ok=True)
            os.makedirs(f"{test_dir}/boot/grub", exist_ok=True)
            os.makedirs(f"{test_dir}/etc/default", 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 UEFI boot order configuration
            grub_stage._configure_uefi_boot(context)
            
            # Verify boot order command was executed
            if not hasattr(context, 'last_command') or 'efibootmgr' not in context.last_command:
                logger.error("UEFI boot order configuration command not executed")
                return False
            
            # Verify boot order was set correctly
            if '--bootorder' not in context.last_command:
                logger.error("UEFI boot order not configured")
                return False
            
            logger.info("UEFI boot order configuration successful")
            return True
            
        except Exception as e:
            logger.error(f"UEFI boot order test failed: {e}")
            return False
    
    def test_secure_boot_integration(self) -> bool:
        """Test Secure Boot integration."""
        test_dir = tempfile.mkdtemp()
        self.temp_dirs.append(test_dir)
        
        try:
            # Create mock context
            context = MockOsbuildContext(test_dir)
            
            # Set up directory structure
            os.makedirs(f"{test_dir}/boot/efi/EFI/debian", exist_ok=True)
            os.makedirs(f"{test_dir}/boot/grub", exist_ok=True)
            os.makedirs(f"{test_dir}/etc/default", exist_ok=True)
            
            # Create mock GRUB EFI file
            grub_efi = f"{test_dir}/boot/efi/EFI/debian/grubx64.efi"
            with open(grub_efi, 'w') as f:
                f.write("mock grub efi")
            
            # Initialize GRUB stage with Secure Boot enabled
            options = {
                'ostree_integration': True,
                'uefi': True,
                'secure_boot': True,
                'timeout': 5,
                'default_entry': 0
            }
            grub_stage = DebianGrubStage(options)
            
            # Test Secure Boot configuration
            grub_stage._configure_secure_boot(context)
            
            # Verify Secure Boot configuration was applied
            if not hasattr(context, 'last_command'):
                logger.error("Secure Boot configuration command not executed")
                return False
            
            # Verify signing commands were called
            if 'sbsign' not in context.last_command and 'pesign' not in context.last_command:
                logger.error("Secure Boot signing commands not executed")
                return False
            
            logger.info("Secure Boot integration successful")
            return True
            
        except Exception as e:
            logger.error(f"Secure Boot integration test failed: {e}")
            return False
    
    def test_uefi_fallback_boot(self) -> bool:
        """Test UEFI fallback boot configuration."""
        test_dir = tempfile.mkdtemp()
        self.temp_dirs.append(test_dir)
        
        try:
            # Create mock context
            context = MockOsbuildContext(test_dir)
            
            # Set up directory structure
            os.makedirs(f"{test_dir}/boot/efi/EFI/debian", exist_ok=True)
            os.makedirs(f"{test_dir}/boot/efi/EFI/BOOT", exist_ok=True)
            os.makedirs(f"{test_dir}/boot/grub", exist_ok=True)
            os.makedirs(f"{test_dir}/etc/default", exist_ok=True)
            
            # Create fallback boot files
            fallback_efi = f"{test_dir}/boot/efi/EFI/BOOT/BOOTX64.EFI"
            with open(fallback_efi, 'w') as f:
                f.write("mock fallback efi")
            
            # Initialize GRUB stage
            options = {
                'ostree_integration': True,
                'uefi': True,
                'secure_boot': False,
                'timeout': 5,
                'default_entry': 0
            }
            grub_stage = DebianGrubStage(options)
            
            # Test UEFI fallback boot configuration
            grub_stage._configure_uefi_boot(context)
            
            # Verify fallback boot file exists
            if not os.path.exists(fallback_efi):
                logger.error(f"UEFI fallback boot file not created: {fallback_efi}")
                return False
            
            # Verify fallback boot was configured
            if not hasattr(context, 'last_command') or 'efibootmgr' not in context.last_command:
                logger.error("UEFI fallback boot configuration command not executed")
                return False
            
            logger.info("UEFI fallback boot configuration successful")
            return True
            
        except Exception as e:
            logger.error(f"UEFI fallback boot 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 EFI installation: {cmd}")
            # Create mock GRUB EFI files
            if '--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")
                # Create fallback boot file
                os.makedirs(f"{self.root_path}/boot/efi/EFI/BOOT", exist_ok=True)
                with open(f"{self.root_path}/boot/efi/EFI/BOOT/BOOTX64.EFI", 'w') as f:
                    f.write("mock fallback efi")
        
        elif 'efibootmgr' in cmd:
            logger.info(f"Mock UEFI boot manager: {cmd}")
            # Simulate UEFI boot entry creation
            if '--create' in ' '.join(cmd):
                logger.info("Mock UEFI boot entry created")
            elif '--bootorder' in ' '.join(cmd):
                logger.info("Mock UEFI boot order set")
        
        elif 'sbsign' in cmd or 'pesign' in cmd:
            logger.info(f"Mock Secure Boot signing: {cmd}")
            # Simulate Secure Boot signing
            if 'grubx64.efi' in ' '.join(cmd):
                logger.info("Mock GRUB EFI signed for Secure Boot")
        
        return subprocess.CompletedProcess(cmd, 0, "", "")


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


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