#!/usr/bin/env python3
"""
Virtual machine testing utilities for deb-bootc-image-builder.

This module provides utilities for testing built images in virtual machines,
including:
- VM creation and management
- Image boot testing
- Debian-specific VM configurations
"""

import os
import subprocess
import tempfile
import shutil
import json
import time
import logging
from typing import Dict, List, Any, Optional, Tuple
from pathlib import Path

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


class VirtualMachine:
    """Virtual machine for testing Debian images."""

    def __init__(self, name: str, image_path: str, vm_type: str = "qemu"):
        """Initialize virtual machine."""
        self.name = name
        self.image_path = image_path
        self.vm_type = vm_type
        self.vm_process = None
        self.vm_pid = None
        self.network_config = None
        self.memory = "2G"
        self.cpus = 2
        self.disk_size = "10G"
        
        # Debian-specific configurations
        self.debian_release = "trixie"
        self.architecture = "amd64"
        self.boot_timeout = 300  # 5 minutes
        
        logger.info(f"Initialized VM {name} with image {image_path}")

    def configure_network(self, network_type: str = "user", port_forward: Optional[Dict[str, int]] = None):
        """Configure VM network."""
        self.network_config = {
            "type": network_type,
            "port_forward": port_forward or {}
        }
        logger.info(f"Configured network: {network_type}")

    def set_resources(self, memory: str = "2G", cpus: int = 2, disk_size: str = "10G"):
        """Set VM resource limits."""
        self.memory = memory
        self.cpus = cpus
        self.disk_size = disk_size
        logger.info(f"Set resources: {memory} RAM, {cpus} CPUs, {disk_size} disk")

    def start(self) -> bool:
        """Start the virtual machine."""
        try:
            if self.vm_type == "qemu":
                return self._start_qemu()
            else:
                logger.error(f"Unsupported VM type: {self.vm_type}")
                return False
        except Exception as e:
            logger.error(f"Failed to start VM: {e}")
            return False

    def _start_qemu(self) -> bool:
        """Start QEMU virtual machine."""
        cmd = [
            "qemu-system-x86_64",
            "-name", self.name,
            "-m", self.memory,
            "-smp", str(self.cpus),
            "-drive", f"file={self.image_path},if=virtio,format=qcow2",
            "-enable-kvm",
            "-display", "none",
            "-serial", "stdio",
            "-monitor", "none"
        ]
        
        # Add network configuration
        if self.network_config:
            if self.network_config["type"] == "user":
                cmd.extend(["-net", "user"])
                if self.network_config["port_forward"]:
                    for host_port, guest_port in self.network_config["port_forward"].items():
                        cmd.extend(["-net", f"user,hostfwd=tcp::{host_port}-:{guest_port}"])
            elif self.network_config["type"] == "bridge":
                cmd.extend(["-net", "bridge,br=virbr0"])
        
        # Add Debian-specific optimizations
        cmd.extend([
            "-cpu", "host",
            "-machine", "type=q35,accel=kvm",
            "-device", "virtio-net-pci,netdev=net0",
            "-netdev", "user,id=net0"
        ])
        
        logger.info(f"Starting QEMU VM with command: {' '.join(cmd)}")
        
        try:
            self.vm_process = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True
            )
            self.vm_pid = self.vm_process.pid
            logger.info(f"VM started with PID: {self.vm_pid}")
            return True
        except Exception as e:
            logger.error(f"Failed to start QEMU VM: {e}")
            return False

    def stop(self) -> bool:
        """Stop the virtual machine."""
        try:
            if self.vm_process:
                self.vm_process.terminate()
                self.vm_process.wait(timeout=30)
                logger.info("VM stopped gracefully")
                return True
            else:
                logger.warning("No VM process to stop")
                return False
        except subprocess.TimeoutExpired:
            logger.warning("VM did not stop gracefully, forcing termination")
            if self.vm_process:
                self.vm_process.kill()
            return True
        except Exception as e:
            logger.error(f"Failed to stop VM: {e}")
            return False

    def is_running(self) -> bool:
        """Check if VM is running."""
        if self.vm_process:
            return self.vm_process.poll() is None
        return False

    def wait_for_boot(self, timeout: Optional[int] = None) -> bool:
        """Wait for VM to boot up."""
        if timeout is None:
            timeout = self.boot_timeout
        
        logger.info(f"Waiting for VM to boot (timeout: {timeout}s)")
        start_time = time.time()
        
        while time.time() - start_time < timeout:
            if self.is_running():
                # Check if VM is responsive (basic boot check)
                if self._check_boot_status():
                    logger.info("VM booted successfully")
                    return True
            
            time.sleep(5)
        
        logger.error("VM boot timeout exceeded")
        return False

    def _check_boot_status(self) -> bool:
        """Check if VM has booted successfully."""
        # This is a simplified check - in practice, you'd want to:
        # 1. Check for specific boot messages
        # 2. Verify network connectivity
        # 3. Check for system services
        # 4. Verify Debian-specific indicators
        
        try:
            # For now, just check if the process is still running
            return self.is_running()
        except Exception as e:
            logger.debug(f"Boot status check failed: {e}")
            return False

    def execute_command(self, command: str, timeout: int = 30) -> Dict[str, Any]:
        """Execute a command in the VM (requires guest agent or SSH)."""
        # This is a placeholder - in practice, you'd use:
        # 1. QEMU guest agent
        # 2. SSH connection
        # 3. Serial console interaction
        
        logger.info(f"Executing command in VM: {command}")
        
        # For now, return a mock result
        return {
            "success": True,
            "stdout": f"Mock output for: {command}",
            "stderr": "",
            "returncode": 0
        }

    def get_system_info(self) -> Dict[str, Any]:
        """Get system information from the VM."""
        # This would typically use guest agent or SSH
        return {
            "os": "Debian GNU/Linux",
            "release": self.debian_release,
            "architecture": self.architecture,
            "kernel": "Linux",
            "uptime": "0:00:00",
            "memory": "0 MB",
            "disk": "0 MB"
        }

    def check_debian_specific_features(self) -> Dict[str, bool]:
        """Check Debian-specific features in the VM."""
        features = {
            "ostree_integration": False,
            "grub_efi": False,
            "initramfs_tools": False,
            "systemd": False,
            "apt_package_manager": False
        }
        
        # This would check actual VM state
        # For now, return mock results
        features["ostree_integration"] = True
        features["grub_efi"] = True
        features["initramfs_tools"] = True
        features["systemd"] = True
        features["apt_package_manager"] = True
        
        return features


class VMTester:
    """Test runner for virtual machines."""

    def __init__(self, test_config: Dict[str, Any]):
        """Initialize VM tester."""
        self.test_config = test_config
        self.vms: List[VirtualMachine] = []
        self.test_results: List[Dict[str, Any]] = []

    def create_test_vm(self, image_path: str, vm_config: Dict[str, Any]) -> VirtualMachine:
        """Create a test virtual machine."""
        vm_name = vm_config.get("name", f"test-vm-{len(self.vms)}")
        vm = VirtualMachine(vm_name, image_path)
        
        # Apply configuration
        if "network" in vm_config:
            vm.configure_network(**vm_config["network"])
        
        if "resources" in vm_config:
            vm.set_resources(**vm_config["resources"])
        
        if "debian" in vm_config:
            debian_config = vm_config["debian"]
            if "release" in debian_config:
                vm.debian_release = debian_config["release"]
            if "architecture" in debian_config:
                vm.architecture = debian_config["architecture"]
            if "boot_timeout" in debian_config:
                vm.boot_timeout = debian_config["boot_timeout"]
        
        self.vms.append(vm)
        return vm

    def run_basic_boot_test(self, vm: VirtualMachine) -> Dict[str, Any]:
        """Run basic boot test on VM."""
        test_result = {
            "test_name": "basic_boot_test",
            "vm_name": vm.name,
            "start_time": time.time(),
            "success": False,
            "error": None,
            "boot_time": None
        }
        
        try:
            logger.info(f"Starting basic boot test for VM: {vm.name}")
            
            # Start VM
            if not vm.start():
                test_result["error"] = "Failed to start VM"
                return test_result
            
            # Wait for boot
            start_time = time.time()
            if vm.wait_for_boot():
                boot_time = time.time() - start_time
                test_result["boot_time"] = boot_time
                test_result["success"] = True
                logger.info(f"Boot test passed for VM: {vm.name} (boot time: {boot_time:.2f}s)")
            else:
                test_result["error"] = "VM failed to boot within timeout"
            
        except Exception as e:
            test_result["error"] = str(e)
            logger.error(f"Boot test failed for VM {vm.name}: {e}")
        
        finally:
            test_result["end_time"] = time.time()
            if test_result["success"]:
                vm.stop()
        
        return test_result

    def run_debian_feature_test(self, vm: VirtualMachine) -> Dict[str, Any]:
        """Run Debian-specific feature test on VM."""
        test_result = {
            "test_name": "debian_feature_test",
            "vm_name": vm.name,
            "start_time": time.time(),
            "success": False,
            "error": None,
            "features": {}
        }
        
        try:
            logger.info(f"Starting Debian feature test for VM: {vm.name}")
            
            # Start VM
            if not vm.start():
                test_result["error"] = "Failed to start VM"
                return test_result
            
            # Wait for boot
            if not vm.wait_for_boot():
                test_result["error"] = "VM failed to boot"
                return test_result
            
            # Check Debian features
            features = vm.check_debian_specific_features()
            test_result["features"] = features
            
            # Validate required features
            required_features = ["ostree_integration", "grub_efi", "systemd"]
            missing_features = [f for f in required_features if not features.get(f, False)]
            
            if missing_features:
                test_result["error"] = f"Missing required features: {missing_features}"
            else:
                test_result["success"] = True
                logger.info(f"Debian feature test passed for VM: {vm.name}")
            
        except Exception as e:
            test_result["error"] = str(e)
            logger.error(f"Debian feature test failed for VM {vm.name}: {e}")
        
        finally:
            test_result["end_time"] = time.time()
            if vm.is_running():
                vm.stop()
        
        return test_result

    def run_all_tests(self) -> List[Dict[str, Any]]:
        """Run all configured tests."""
        logger.info("Starting VM test suite")
        
        for vm in self.vms:
            # Run basic boot test
            boot_result = self.run_basic_boot_test(vm)
            self.test_results.append(boot_result)
            
            # Run Debian feature test
            feature_result = self.run_debian_feature_test(vm)
            self.test_results.append(feature_result)
        
        # Generate summary
        total_tests = len(self.test_results)
        passed_tests = len([r for r in self.test_results if r["success"]])
        failed_tests = total_tests - passed_tests
        
        logger.info(f"Test suite completed: {passed_tests}/{total_tests} tests passed")
        
        if failed_tests > 0:
            logger.warning(f"{failed_tests} tests failed")
            for result in self.test_results:
                if not result["success"]:
                    logger.error(f"Test {result['test_name']} failed: {result.get('error', 'Unknown error')}")
        
        return self.test_results

    def cleanup(self):
        """Clean up all VMs and resources."""
        logger.info("Cleaning up VM resources")
        
        for vm in self.vms:
            if vm.is_running():
                vm.stop()
        
        self.vms.clear()
        self.test_results.clear()


def create_test_vm_config(image_path: str, 
                         debian_release: str = "trixie",
                         architecture: str = "amd64") -> Dict[str, Any]:
    """Create a test VM configuration."""
    return {
        "name": f"debian-{debian_release}-{architecture}-test",
        "image_path": image_path,
        "network": {
            "type": "user",
            "port_forward": {"2222": 22}  # SSH port forwarding
        },
        "resources": {
            "memory": "2G",
            "cpus": 2,
            "disk_size": "10G"
        },
        "debian": {
            "release": debian_release,
            "architecture": architecture,
            "boot_timeout": 300
        }
    }


def run_vm_test_suite(image_path: str, 
                      debian_release: str = "trixie",
                      architecture: str = "amd64") -> List[Dict[str, Any]]:
    """Run a complete VM test suite."""
    # Create test configuration
    vm_config = create_test_vm_config(image_path, debian_release, architecture)
    
    # Create tester
    tester = VMTester({})
    
    try:
        # Create test VM
        vm = tester.create_test_vm(image_path, vm_config)
        
        # Run tests
        results = tester.run_all_tests()
        
        return results
    
    finally:
        # Cleanup
        tester.cleanup()


if __name__ == "__main__":
    # Example usage
    test_image = "/tmp/test-image.qcow2"
    
    if os.path.exists(test_image):
        print("Running VM test suite...")
        results = run_vm_test_suite(test_image)
        
        for result in results:
            status = "PASS" if result["success"] else "FAIL"
            print(f"{status}: {result['test_name']} - {result['vm_name']}")
            if not result["success"]:
                print(f"  Error: {result.get('error', 'Unknown error')}")
    else:
        print(f"Test image not found: {test_image}")
        print("Create a test image first or update the path")