"""
QEMU Virtual Machine Manager with Local Storage I        # Set up base paths
        storage_base = os.path.join(os.path.dirname(__file__), 'storage')
        os.makedirs(storage_base, exist_ok=True)
        
        # Initialize disk-based storage for virtual RAM
        self.memory_storage = DiskStorage(
            base_path=os.path.join(storage_base, 'virtual_ram'),
            page_size=4096  # 4KB pages
        )
        
        # Initialize CPU storage
        from cpu.enhanced_cpu import storage, CPUGroupType
        if storage is None:
            storage = LocalStorageManager(os.path.join(storage_base, 'cpu.db'))
            storage.init_db()
            
        # Initialize CPU with proper group type
        self.cpu = EnhancedCPU(cpu_id=0, group_type=CPUGroupType.SYSTEM_TASKS)
        
        # Initialize HAL with storage
        hal_storage = LocalStorage(os.path.join(storage_base, 'hal.db'))
        self.hal = HardwareAbstractionLayer(storage_manager=hal_storage)anages VM state, virtual GPU, VRAM, and CPU states through SQLite storage
"""

import os
import sys
import asyncio
import subprocess
import time
from pathlib import Path
from typing import Dict, Optional, List
import json
import logging
from datetime import datetime
import hashlib
import aiohttp
import aiofiles

# Add parent directory to path for imports
parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(parent_dir)

# Import virtual hardware components
from cpu.enhanced_cpu import EnhancedCPU
from virtual_gpu_driver.src.hal.hal import HardwareAbstractionLayer, HardwareType
from disk_storage import DiskStorageManager  # For disk-based memory
from http_storage import LocalStorage  # For HAL initialization
from local_storage_manager import LocalStorageManager  # For CPU storage
from .qemu_storage import QEMUStorage

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

class QEMUManager:
    def __init__(self, storage_path: str = None):
        """Initialize QEMU Manager with virtual hardware"""
        self.vm_process = None
        self.vnc_port = 5900
        self.vm_state = {
            "status": "stopped",
            "installation": {},
            "hardware": {
                "cpu": {},
                "gpu": {},
                "memory": {}
            }
        }
        
        # Initialize storage
        self.storage_path = storage_path or os.path.join(os.path.dirname(__file__), 'vm_storage.db')
        self.storage = QEMUStorage(self.storage_path)
        
        # Initialize storage and virtual hardware
        self.storage_manager = LocalStorage()
        
        # Initialize CPU storage and system
        import cpu.enhanced_cpu as cpu_module
        if cpu_module.storage is None:
            storage_path = os.path.join(os.path.dirname(__file__), 'storage', 'cpu.db')
            os.makedirs(os.path.dirname(storage_path), exist_ok=True)
            cpu_module.storage = LocalStorageManager(storage_path)
            
        # Initialize virtual CPU system first
        # Set class attributes
        cpu_module.VirtualCPU.CORES_PER_CPU = 8
        cpu_module.VirtualCPU.THREADS_PER_CORE = 2
        
        # Create system config table
        cpu_module.storage.cursor.execute("""
        CREATE TABLE IF NOT EXISTS cpu_system_config (
            id INTEGER PRIMARY KEY,
            total_cpus INTEGER,
            cores_per_cpu INTEGER,
            threads_per_core INTEGER,
            batch_size INTEGER,
            initialized INTEGER DEFAULT 0,
            last_updated INTEGER DEFAULT 0
        )""")
        
        # Insert system configuration
        cpu_module.storage.cursor.execute("""
        INSERT OR REPLACE INTO cpu_system_config 
        (id, total_cpus, cores_per_cpu, threads_per_core, batch_size, initialized, last_updated)
        VALUES (1, ?, ?, ?, ?, 1, ?)
        """, (1, 8, 2, 1000, int(time.time())))
        cpu_module.storage.conn.commit()
        
        # Initialize the system
        cpu_module.VirtualCPU.initialize_system(
            total_cpus=1,  # Single CPU for QEMU
            cores_per_cpu=8,  # 8 cores
            threads_per_core=2,  # 2 threads per core
            batch_size=1000
        )
        
        # Initialize CPU with proper storage
        self.cpu = EnhancedCPU(cpu_id=0, group_type=cpu_module.CPUGroupType.SYSTEM_TASKS)
        
        # Initialize HAL with LocalStorage for GPU state
        hal_storage_path = os.path.join(os.path.dirname(__file__), 'storage', 'hal.db')
        hal_storage = LocalStorage(hal_storage_path)
        self.hal = HardwareAbstractionLayer(storage_manager=hal_storage)
        
        # Initialize memory manager with HAL
        self.memory_manager = MemoryManager(hal=self.hal)
        # Configure virtual hardware
        self._configure_virtual_hardware()
        
        # Load saved state if exists
        self._load_state()
    
    def _configure_virtual_hardware(self):
        """Configure virtual hardware components"""
        # CPU is already configured via module-level storage
        
        # Configure GPU via HAL
        gpu_config = {
            "compute_units": 8,
            "tensor_cores": 4,
            "shader_units": 16,
            "vram_size": 8 * 1024 * 1024 * 1024  # 8GB
        }
        self.hal.configure_hardware(HardwareType.COMPUTE_UNIT, gpu_config)
        
        # Configure memory
        self.memory_manager.create_memory_pool(8 * 1024 * 1024 * 1024)  # 8GB RAM
        
        # Update state
        self.vm_state["hardware"].update({
            "cpu": self.cpu.get_config(),
            "gpu": gpu_config,
            "memory": {"total": 8 * 1024 * 1024 * 1024}
        })
    
    def _load_state(self):
        """Load saved VM state"""
        try:
            for key in ["status", "installation", "hardware"]:
                value = self.storage.get_vm_state(key)
                if value:
                    self.vm_state[key] = value
        except Exception as e:
            logging.warning(f"Failed to load VM state: {e}")
    def get_qemu_command(self, disk_path: str, iso_path: Optional[str] = None) -> str:
        """Generate QEMU command with virtual hardware configuration using disk-based memory"""
        # Create memory backing file
        memory_path = os.path.join(os.path.dirname(__file__), 'storage', 'virtual_ram', 'memory.backing')
        
        cmd = [
            'qemu-system-x86_64',
            '-enable-kvm',
            # Use minimal host RAM, most memory will be on disk
            '-m 256',  # Only 256MB in actual RAM
            f'-cpu host',
            f'-smp {self.cpu.get_config()["cores"]}',
            # Main disk
            f'-drive file={disk_path},format=qcow2',
            # Memory backing file (uses disk as RAM)
            f'-object memory-backend-file,id=mem1,size=8G,mem-path={memory_path},share=on',
            f'-numa node,memdev=mem1',
            # Graphics
            f'-vga virtio',
            f'-vnc :{self.vnc_port - 5900}',
            # Enable memory ballooning for dynamic allocation
            '-device virtio-balloon-pci,id=balloon0',
            # Add memory prealloc to reduce runtime allocation overhead
            '-mem-prealloc'
        ]
        
        # Add installation ISO if provided
        if iso_path:
            cmd.append(f'-cdrom {iso_path}')
        
        # Add virtual hardware configurations
        gpu_config = self.vm_state["hardware"]["gpu"]
        cmd.extend([
            f'-device vfio-pci,sysfsdev=/sys/class/mdev/vgpu-{gpu_config["compute_units"]}'
        ])
        
        return ' '.join(cmd)

    async def capture_frame(self) -> bytes:
        """Capture current frame from QEMU VGA output with reduced size"""
        if not self.vm_process:
            return None
            
        try:
            # Get frame from HAL's frame buffer
            raw_frame = self.hal.get_frame_buffer()
            if not raw_frame:
                return None
                
            # Convert raw frame to compressed JPEG for browser
            return self._compress_frame(raw_frame)

        except Exception as e:
            logger.error(f"Frame capture error: {e}")
            return None
    
    def _compress_frame(self, frame_data: bytes) -> bytes:
        """Compress frame data for browser rendering"""
        try:
            # Get frame dimensions from GPU config
            width = 1024  # Default width
            height = 768  # Default height
            
            # Use HAL to get actual dimensions if available
            gpu_info = self.hal.get_hardware_info(HardwareType.COMPUTE_UNIT)
            if gpu_info and 'resolution' in gpu_info:
                width, height = gpu_info['resolution']
            
            # Compress frame to JPEG with reduced quality for browser
            import cv2
            import numpy as np
            
            # Convert raw frame to numpy array
            frame = np.frombuffer(frame_data, dtype=np.uint8)
            frame = frame.reshape((height, width, 4))  # RGBA format
            
            # Reduce resolution if too large
            max_width = 800
            if width > max_width:
                scale = max_width / width
                new_width = int(width * scale)
                new_height = int(height * scale)
                frame = cv2.resize(frame, (new_width, new_height))
            
            # Convert to JPEG with moderate compression
            encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), 70]
            _, buffer = cv2.imencode('.jpg', frame, encode_param)
            
            return buffer.tobytes()
            
        except Exception as e:
            logger.error(f"Frame compression error: {e}")
            return None
    
    async def start_frame_streaming(self, websocket):
        """Stream frames to browser with controlled FPS"""
        fps = 30  # Target FPS
        frame_delay = 1.0 / fps
        
        while self.vm_process and self.vm_process.poll() is None:
            try:
                frame_start = time.time()
                
                # Capture and send frame
                frame = await self.capture_frame()
                if frame:
                    await websocket.send_bytes(frame)
                
                # Control frame rate
                elapsed = time.time() - frame_start
                if elapsed < frame_delay:
                    await asyncio.sleep(frame_delay - elapsed)
                    
            except Exception as e:
                logger.error(f"Streaming error: {e}")
                break
        self.vnc_port = 5900
        self.memory_service = None
        self.storage_path = str(Path(__file__).parent / "data" / "qemu_state.json")
        os.makedirs(os.path.dirname(self.storage_path), exist_ok=True)
        
        # Use provided storage or initialize local storage
        self.storage = storage or QEMUStorage()
        
        # Load or initialize VM state
        status = self.storage.get_vm_state('status')
        if not status:
            self.vm_state = {
                "status": "initialized",
                "installation": {
                    "completed": False,
                    "iso_url": None,
                    "checksum": None
                },
                "hardware_config": self.config
            }
            for key, value in self.vm_state.items():
                self.storage.update_vm_state(key, value)
        else:
            self.vm_state = {
                "status": status,
                "installation": self.storage.get_vm_state('installation') or {},
                "hardware_config": self.storage.get_vm_state('hardware_config') or self.config
            }
        
        # Memory service socket path
        self.memory_socket = "/tmp/memory-backend.sock"

        # Initialize simplified virtual hardware
        self.gpus = [VirtualGPU() for _ in range(8)]
        self.cpus = [VirtualCPU(i) for i in range(4)]

        # MMIO regions for device communication
        self.mmio_regions = {
            'gpu_cmd': 0xF0000000,
            'gpu_fb': 0xF1000000,
            'gpu_status': 0xF3000000
        }
        
        # Initialize data storage
        self.data = {'devices': {}, 'cpu_states': {}, 'gpu_state': {}, 'vram_mapping': {}}
        
        # Initialize all virtual hardware components
        self.init_virtual_hardware()
        
    def _load_data(self) -> Dict:
        """Load data from storage file"""
        if os.path.exists(self.storage_path):
            try:
                with open(self.storage_path, 'r') as f:
                    data = json.load(f)
                    if 'devices' not in data:
                        data['devices'] = {}
                    return data
            except json.JSONDecodeError:
                pass
        return {'devices': {}}
        
    def _save_data(self, data: Dict):
        """Save data to storage file"""
        with open(self.storage_path, 'w') as f:
            json.dump(data, f, indent=2)
        
    def _init_database_schema(self):
        """Initialize the database schema if it doesn't exist"""
        self.data = self._load_data()
        if not os.path.exists(self.storage_path):
            self._save_data(self.data)
        
        # Initialize GPU devices if they don't exist
        self.con.execute("""
            INSERT OR IGNORE INTO devices (id, type, state)
            SELECT 'gpu_' || CAST(id AS VARCHAR), 'gpu_device', '{}'
            FROM range(0, 8) t(id);
        """)
        
        # Create views for easier access
        self.con.execute("""
            CREATE VIEW IF NOT EXISTS gpu_state AS
            SELECT 
                id as device_id,
                COALESCE(json_extract(state, '$.commands'), '[]') as command_buffer,
                COALESCE(json_extract(state, '$.frame_data'), NULL) as framebuffer,
                state as status,
                updated_at as timestamp
            FROM devices 
            WHERE type = 'gpu_device';
        """)
        
        # Create view for VRAM mapping
        self.con.execute("""
            CREATE VIEW IF NOT EXISTS vram_mapping AS
            SELECT 
                memory_address as address,
                block_size as size,
                memory_data as content,
                access_flags as flags,
                COALESCE(last_accessed, CURRENT_TIMESTAMP) as last_access
            FROM raw_data 
            WHERE type = 'vram_block';
        """)
        
        # Create view for general hardware state
        self.con.execute("""
            CREATE VIEW IF NOT EXISTS virtual_hardware AS
            SELECT * FROM raw_data
            WHERE type = 'hardware_config';
        """)
        
        return self.con
        
    def init_virtual_hardware(self):
        """Initialize virtual hardware components"""
        self._init_cpu_state()
        self._init_virtual_gpu()
        self._init_virtual_vram()
        
    def _init_cpu_state(self):
        """Initialize CPU state tracking with EnhancedCPU and GPU integration"""
        self.cpu_groups = {
            'ui_display': (0, 499),
            'computation': (500, 999),
            'io_storage': (1000, 1499),
            'system_tasks': (1500, 1999)
        }
        
        # Initialize virtual CPUs
        self.virtual_cpus = {}
        for group_type, (start, end) in self.cpu_groups.items():
            for cpu_id in range(start, end + 1):
                # Create CPU with basic virtualization
                cpu = VirtualCPU(cpu_id=cpu_id)
                self.virtual_cpus[cpu_id] = cpu
                
                # Initialize state tracking in local storage
                if 'cpu_states' not in self.data:
                    self.data['cpu_states'] = {}
                if 'gpu_state' not in self.data:
                    self.data['gpu_state'] = {}
                if 'vram_mapping' not in self.data:
                    self.data['vram_mapping'] = {}
                self._save_data(self.data)
                
    def _init_virtual_gpu(self):
        """Initialize virtual GPU state"""
        if 'gpu_state' not in self.data:
            self.data['gpu_state'] = {
                'devices': {str(i): {'status': 'ready'} for i in range(8)},
                'framebuffers': {str(i): bytearray(1024 * 768 * 4) for i in range(8)}
            }
            self._save_data(self.data)
        
    def _init_virtual_vram(self):
        """Initialize virtual VRAM mapping"""
        if 'vram_mapping' not in self.data:
            self.data['vram_mapping'] = {
                'addresses': {},
                'contents': {},
                'flags': {},
                'access_times': {}
            }
            self._save_data(self.data)
            
    def load_config(self, config_path: str) -> Dict:
        """Load VM configuration"""
        with open(config_path) as f:
            config = json.load(f)
            
        # Hardware-based configuration limits for 2025 server hardware
        max_cpus = 8  # Maximum physical CPU sockets (high-end server)
        max_cores = 128  # Maximum cores per CPU (e.g., future EPYC)
        max_threads = 222  # Standard hyperthreading (2 threads per core)
        max_memory = '8192G'  # 8TB RAM maximum for 2025 server
        max_disk = '16384G'  # 16TB maximum storage
        max_vram = '192G'  # Maximum VRAM (future datacenter GPU)
        max_fps = 144  # Maximum refresh rate supported
        
        # Validate and set defaults with hardware limits
        cpu_count = min(int(config.get('cpus', max_cpus)), max_cpus)
        core_count = min(int(config.get('cores_per_cpu', max_cores)), max_cores)
        thread_count = min(int(config.get('threads_per_core', max_threads)), max_threads)
        fps = min(int(config.get('fps', 60)), max_fps)  # Default to 60 FPS, cap at max_fps
        
        # Convert memory sizes to GB for validation
        def parse_size(size_str):
            if isinstance(size_str, (int, float)):
                return size_str
            unit = size_str[-1].upper()
            value = float(size_str[:-1])
            if unit == 'T': value *= 1024
            elif unit == 'G': value = value
            elif unit == 'M': value /= 1024
            return value
            
        mem_size = min(parse_size(config.get('memory', max_memory)), parse_size(max_memory))
        disk_size = min(parse_size(config.get('disk_size', max_disk)), parse_size(max_disk))
        vram_size = min(parse_size(config.get('vgpu_memory', max_vram)), parse_size(max_vram))
        
        # Set validated configuration
        config.update({
            'cpus': cpu_count,  # Number of CPU sockets
            'cores_per_cpu': core_count,  # Cores per CPU
            'threads_per_core': thread_count,  # Threads per core
            'memory': f'{int(mem_size)}G',  # System memory
            'disk_size': f'{int(disk_size)}G',  # Virtual disk size
            'vgpu_memory': f'{int(vram_size)}G',  # VRAM size
            'fps': fps,  # Display refresh rate
        })
        
        # Log actual configuration
        total_vcpus = cpu_count * core_count * thread_count
        logger.info(f"VM Configuration:")
        logger.info(f" - Total vCPUs: {total_vcpus} ({cpu_count} sockets × {core_count} cores × {thread_count} threads)")
        logger.info(f" - Memory: {config['memory']}")
        logger.info(f" - Disk Size: {config['disk_size']}")
        logger.info(f" - VGPU Memory: {config['vgpu_memory']}")
        
        return config
        
    def _start_memory_service(self):
        """Start the memory management service that handles virtual memory operations"""
        try:
            # Start memory service in a separate process
            import multiprocessing
            from virtual_gpu_driver import memory_service
            
            # Create service instance
            self.memory_service = memory_service.MemoryService(
                socket_path=self.memory_socket,
                storage=self.storage,
                memory_manager=self.memory_manager
            )
            
            # Start in separate process
            self.memory_service_process = multiprocessing.Process(
                target=self.memory_service.start,
                daemon=True  # Ensure process exits when main program exits
            )
            self.memory_service_process.start()
            
            # Wait for service to be ready
            import time
            start_time = time.time()
            while not os.path.exists(self.memory_socket):
                if time.time() - start_time > 10:  # 10 second timeout
                    raise TimeoutError("Memory service failed to start")
                time.sleep(0.1)
                
            logger.info("Memory service started successfully")
            
        except Exception as e:
            logger.error(f"Failed to start memory service: {e}")
            raise
            
    def create_virtual_disk(self) -> str:
        """Create virtual disk using HuggingFace dataset as backend storage"""
        logger.info("Initializing HuggingFace dataset-backed virtual disk")
        
        from vram.nand_memory import NANDMemory, DynamicNANDMemory
        from vram.dram_cache import DRAMCache, DynamicDRAMCache
        from vram.ftl import FTL
        from vram.dynamic_ftl import DynamicFTL
        from vram.nvme import NVMeController
        from huggingface_hub import HfApi, Repository
        
        # Initialize HuggingFace dataset storage
        repo = Repository(
            local_dir=self.storage_path,
            clone_from=self.hf_dataset_path,
            repo_type="dataset",
            use_auth_token=True
        )
        
        # Use HuggingFace-backed storage
        storage = self.storage
        
        # Create dynamic NAND memory configuration
        self.nand_memory = DynamicNANDMemory(
            storage=storage,
            initial_planes=16,       # Start with 16 planes
            initial_blocks=512,      # 512 blocks per plane
            initial_pages=512,       # 512 pages per block
            bytes_per_page=32768,    # 32KB pages
            auto_scale=True,         # Enable automatic scaling
            max_scale_factor=float('inf')  # No upper limit
        )
        
        # Initialize dynamic cache and FTL
        self.dram_cache = DynamicDRAMCache(
            initial_size_mb=4096,    # Start with 4GB cache
            max_size_mb=None,        # No size limit
            auto_scale=True          # Automatically scale with usage
        )
        
        self.ftl = DynamicFTL(
            nand_memory=self.nand_memory,
            dynamic_mapping=True,     # Enable dynamic mapping table
            unlimited_blocks=True     # Allow unlimited block allocation
        )
        
        # Create NVMe controller with unlimited queue depth
        self.nvme_controller = NVMeController(
            nand_memory=self.nand_memory,
            dram_cache=self.dram_cache,
            ftl=self.ftl,
            max_queues=1024,          # Large number of I/O queues
            queue_depth=65535,        # Maximum queue depth
            dynamic_queues=True       # Enable dynamic queue creation
        )
        
        # Return special NVMe device path for unlimited storage
        return f"nvme://{self.storage_manager.dataset_path}?type=virtual_disk&dynamic=true"
        
    def get_qemu_command(self, disk_path: str, iso_path: Optional[str] = None) -> str:
        """Generate QEMU command with all hardware configuration"""
        # Calculate maximum CPUs to match our EnhancedCPU groups
        total_cpus = 2000  # Total CPUs from CPU groups (0-1999)
        max_cpus = 4       # 4 physical CPU sockets (500 CPUs per group * 4 groups)
        cores_per_cpu = 250  # 250 cores per CPU to get 1000 cores total
        threads_per_core = 2  # Each core has 2 threads to reach 2000 total vCPUs
        
        # This gives us exactly 2000 vCPUs (4 sockets × 250 cores × 2 threads)
        # Matching our EnhancedCPU groups of 500 CPUs each:
        # - UI_DISPLAY (0-499)
        # - COMPUTATION (500-999)
        # - IO_STORAGE (1000-1499)
        # - SYSTEM_TASKS (1500-1999)
        
        # Update config with CPU values matching EnhancedCPU setup
        self.config['max_cpus'] = max_cpus
        self.config['max_cores'] = cores_per_cpu
        self.config['max_threads'] = threads_per_core
        
        # Start our virtual memory management service
        if not hasattr(self, 'memory_service_process'):
            self._start_memory_service()
        
        cmd = [
            "qemu-system-x86_64",
            
            # Machine configuration - using our virtual hardware
            "-machine microvm",  # Lightweight VM without emulated BIOS/hardware
            "-nodefaults",  # Don't create default devices
            "-no-acpi",    # Using our own ACPI implementation
            
            # Enhanced Virtual CPU configuration
            "-cpu custom,vendor=Virtual,family=2,model=1,stepping=1",
            f"cores={self.config['cores_per_cpu']}",
            f"threads={self.config['threads_per_core']}",
            "virtual-insns=on,virtual-mmu=on,virtual-timer=on,",  # Virtual CPU features
            "enhanced-grid=on,cpu-groups=4,",  # Enhanced CPU features
            f"group-sizes={','.join(str(end-start+1) for start, end in self.cpu_groups.values())},",
            "-no-hpet,-no-tsc",  # Using pure virtual timing
            
            # CPU topology using our virtual cores
            f"-smp {self.config['cpus']}",
            f"cores={self.config['cores_per_cpu']}",
            f"threads={self.config['threads_per_core']}",
            f"sockets=1,maxcpus={total_cpus}",
            
            # System memory using RemoteStorageManager backend
            "-object memory-backend-file,",
            f"size={self.config['memory']},id=ram0,",
            f"mem-path={self.storage_path}?type=system,",
            "share=on",
            "-numa node,memdev=ram0,nodeid=0",
            
            # RNG device configuration
            "-object rng-random,filename=/dev/urandom,id=rng0",
            "-device virtio-rng-pci,rng=rng0,bus=pcie.0",
            
            # Pure virtual GPU configuration with integrated VRAM
            "-device virtio-gpu-pci,",
            "id=gpu0,max_outputs=1,",
            "virtual-gpu=on,",
            f"bus=pcie.0,addr=0x2,",
            "virtual-render-api=custom,",
            f"virtual-gpu-path={self.hal.get_gpu_path()},",
            f"vram-path={self.storage_path}?type=vram,",
            f"vram-size={self.config['vgpu_memory']},",
            "shared-vram=on",
            
            # MMIO region mapping
            "-device ivshmem-doorbell,",
            f"vectors={len(self.mmio_regions)},",
            "id=shmem0",
            
            # Storage configuration using unlimited virtual NAND
            "-device nvme,",
            "id=nvme0,",
            "serial=unlimited-nvme0,",  # Unique serial for unlimited drive
            "num_queues=1024,",         # Large number of queues
            "queue_size=65535,",        # Maximum queue size
            "max_transfer=1048576,",    # 1MB max transfer
            "dynamic_queues=on,",       # Enable dynamic queue scaling
            f"path={self.nvme_controller.get_storage_path()},",
            "unlimited=on,",            # Enable unlimited storage mode
            "auto_scale=on",           # Enable automatic scaling
            
            # Virtual disk configuration with unlimited storage
            "-drive",
            f"file=nvme://{self.storage_path}?type=virtual_disk&dynamic=true,",
            "if=none,",
            "id=nvme0,",
            "format=raw,",
            "cache=directsync,",        # Direct sync with remote storage
            "discard=unmap,",          # Enable TRIM support
            "detect-zeroes=unmap,",    # Optimize zero writes
            "aio=native,",             # Use native async I/O
            "readonly=off"             # Allow writes
            
            # Network configuration
            "-netdev user,id=net0",
            "-device virtio-net-pci,netdev=net0",
            
            # Display and monitoring
            "-display none",  # Headless mode
            f"-vnc :{self.vnc_port - 5900}",  # VNC for remote access
            "-monitor unix:qemu-monitor-socket,server,nowait",
            
            # Debug and performance options
            "-no-user-config",  # Don't load user config
            "-nodefaults",  # Don't add default devices
            "-global kvm-pit.lost_tick_policy=delay",  # Better timing
            "-rtc base=localtime,clock=host",  # Use host clock
            "-boot strict=off"  # Flexible boot options
        ]
        
        # Add ISO if installing
        if iso_path:
            cmd.extend([
                "-cdrom", iso_path,
                "-boot", "d"
            ])
            
        return " ".join(cmd)
        
    async def install_os(self, iso_url: str):
        """Install OS from ISO with virtual hardware support using local storage"""
        # Update VM state for installation
        self.vm_state["installation"]["iso_url"] = iso_url
        self.vm_state["status"] = "installing"
        for key, value in self.vm_state.items():
            self.storage.update_vm_state(key, value)
        
        logger.info(f"Starting OS installation from {iso_url}")
        try:
            # Download ISO to memory first
            async with aiohttp.ClientSession() as session:
                async with session.get(iso_url) as response:
                    if response.status != 200:
                        raise RuntimeError(f"Failed to download ISO: {response.status}")
                    
                    # Stream download with progress
                    total_size = int(response.headers.get('content-length', 0))
                    chunk_size = 1024 * 1024  # 1MB chunks
                    downloaded = 0
                    checksum = hashlib.sha256()
                    
                    # Store ISO in chunks directly to HuggingFace dataset
                    while True:
                        chunk = await response.content.read(chunk_size)
                        if not chunk:
                            break
                            
                        # Update checksum
                        checksum.update(chunk)
                        
                        # Store chunk in database
                        self.storage.store_os_image_chunk(iso_path, downloaded // chunk_size, chunk)
                        
                        # Update progress
                        downloaded += len(chunk)
                        progress = (downloaded / total_size) * 100
                        logger.info(f"Download progress: {progress:.1f}%")
                    
                    # Create OS image entry
                    image_id = hashlib.sha256(f"os_image_{time.time()}".encode()).hexdigest()
                    
                    # Update installation state
                    self.vm_state["installation"].update({
                        "checksum": checksum.hexdigest(),
                        "image_id": image_id,
                        "completed": True
                    })
                    
                    for key, value in self.vm_state.items():
                        self.storage.update_vm_state(key, value)
                    
                    logger.info(f"ISO downloaded and stored. SHA256: {checksum.hexdigest()}")
                    
        except Exception as e:
            self.vm_state["status"] = "error"
            self.vm_state["installation"]["error"] = str(e)
            for key, value in self.vm_state.items():
                self.storage.update_vm_state(key, value)
            raise RuntimeError(f"Failed to download ISO: {e}")
                
        # Get temporary path for the OS image
        image_path = os.path.join(os.path.dirname(__file__), 'temp', f'os_{image_id}.iso')
        os.makedirs(os.path.dirname(image_path), exist_ok=True)
            
        # Write OS image to temporary file for QEMU
        image_data = self.storage.get_os_image(image_id)
        if not image_data:
            raise RuntimeError("Failed to retrieve OS image from storage")
            
        with open(image_path, 'wb') as f:
            f.write(image_data)
            
        # Create virtual disk for installation
        disk_path = os.path.join(os.path.dirname(__file__), 'disk', 'os.qcow2')
        os.makedirs(os.path.dirname(disk_path), exist_ok=True)
        
        logger.info("Preparing virtual hardware for OS installation...")
        
        # Start hardware monitoring
        monitor_task = asyncio.create_task(self._monitor_hardware())
        
        # Start QEMU with ISO and virtual hardware
        cmd = self.get_qemu_command(disk_url, iso_path)
        logger.info("Starting OS installation with virtual hardware...")
        
        # Set environment for virtual devices during install
        env = os.environ.copy()
        env.update({
            'VGPU_DB_URL': self.db_url,
            'VGPU_MMIO_CMD': hex(self.mmio_regions['gpu_cmd']),
            'VGPU_MMIO_FB': hex(self.mmio_regions['gpu_fb']),
            'VGPU_MMIO_STATUS': hex(self.mmio_regions['gpu_status'])
        })
        
        self.vm_process = subprocess.Popen(cmd, shell=True, env=env)
        
        # Wait for installation to complete
        await self.vm_process.wait()
        
        # Stop monitoring
        monitor_task.cancel()
        
        logger.info("OS installation finished.")
        return self.vnc_port
        
    async def boot_os(self):
        """Boot existing OS with virtual hardware"""
        # Get the OS image from storage
        installation = self.storage.get_vm_state('installation')
        if not installation or not installation.get('completed'):
            raise FileNotFoundError("No OS installation found.")
            
        image_id = installation.get('image_id')
        if not image_id:
            raise FileNotFoundError("OS image ID not found in installation state.")
            
        # Get the image data
        image_data = self.storage.get_os_image(image_id)
        if not image_data:
            raise FileNotFoundError("OS image not found in storage.")
            
        # Write to temporary disk file
        disk_path = os.path.join(os.path.dirname(__file__), 'disk', 'os.qcow2')
        os.makedirs(os.path.dirname(disk_path), exist_ok=True)
        with open(disk_path, 'wb') as f:
            f.write(image_data)
            
        # Virtual hardware is already initialized in __init__
        self._map_virtual_hardware()
        
        # Configure GPU for boot
        await self.hal.set_gpu_mode_async('boot')
        
        # Start hardware monitoring
        monitor_task = asyncio.create_task(self._monitor_hardware())
        
        # Start QEMU
        cmd = self.get_qemu_command(disk_path)
        logger.info("Booting OS with virtual hardware...")
        
        env = os.environ.copy()
        env.update({
            'VGPU_DB_URL': self.db_url,
            'VGPU_MMIO_CMD': hex(self.mmio_regions['gpu_cmd']),
            'VGPU_MMIO_FB': hex(self.mmio_regions['gpu_fb']),
            'VGPU_MMIO_STATUS': hex(self.mmio_regions['gpu_status'])
        })
        
        self.vm_process = subprocess.Popen(cmd, shell=True, env=env)
        
        return self.vnc_port
        
    async def shutdown(self):
        """Shutdown running VM"""
        if self.vm_process:
            self.vm_process.terminate()
            await self.vm_process.wait()
            self.vm_process = None
            logger.info("VM shut down.")


