"""
QEMU Virtual Machine Manager with Local Storage Integration
"""

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

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

class QEMUManager:
    def __init__(self, storage=None, config_path: str = "config.json"):
        self.config = self.load_config(str(Path(__file__).parent / config_path))
        self.vm_process: Optional[subprocess.Popen] = None
        self.vnc_port = 5900
        self.storage = None  # Initialize storage as None
        
        # Initialize VM state
        self.vm_state = {
            "status": "initialized",
            "installation": {
                "completed": False,
                "iso_url": None,
                "checksum": None,
                "image_id": None
            },
            "hardware_config": self.config
        }
        
        # Save initial state
        if self.storage:
            for key, value in self.vm_state.items():
                self.storage.update_vm_state(key, value)

    def load_config(self, config_path: str) -> Dict:
        """Load QEMU configuration"""
        try:
            with open(config_path, 'r') as f:
                return json.load(f)
        except FileNotFoundError:
            logger.warning("Config file not found, using defaults")
            return {
                'vm': {'memory': '4G', 'cpus': 4},
                'display': {'type': 'vnc', 'port': 5900},
                'storage': {'type': 'qcow2', 'size': '20G'}
            }

    async def install_os(self, iso_url: str):
        """Install OS from ISO with virtual hardware support"""
        # Generate unique ID for this OS image
        image_id = hashlib.sha256(f"os_image_{time.time()}".encode()).hexdigest()
        
        # Update VM state for installation
        self.vm_state["installation"].update({
            "iso_url": iso_url,
            "image_id": image_id,
            "completed": False
        })
        self.vm_state["status"] = "installing"
        
        if self.storage:
            for key, value in self.vm_state.items():
                self.storage.update_vm_state(key, value)
        
        logger.info(f"Starting OS installation from {iso_url}")
        
        # Create temporary directory for downloaded ISO
        iso_path = os.path.join(os.path.dirname(__file__), 'temp', f'os_{image_id}.iso')
        os.makedirs(os.path.dirname(iso_path), exist_ok=True)
        
        try:
            # Download ISO to temporary file 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()
                    
                    # Download and store chunks
                    with open(iso_path, 'wb') as f:
                        while True:
                            chunk = await response.content.read(chunk_size)
                            if not chunk:
                                break
                                
                            # Update checksum and write chunk
                            checksum.update(chunk)
                            f.write(chunk)
                            
                            # Store chunk in database if storage is available
                            if self.storage:
                                self.storage.store_os_image_chunk(image_id, downloaded // chunk_size, chunk)
                            
                            # Update progress
                            downloaded += len(chunk)
                            progress = (downloaded / total_size) * 100
                            logger.info(f"Download progress: {progress:.1f}%")
                    
                    # Update installation state
                    self.vm_state["installation"].update({
                        "checksum": checksum.hexdigest(),
                        "completed": True
                    })
                    
                    if self.storage:
                        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)
            if self.storage:
                for key, value in self.vm_state.items():
                    self.storage.update_vm_state(key, value)
            raise RuntimeError(f"Failed to download ISO: {e}")
        
        # 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)
        
        # Create QCOW2 disk image
        qemu_img_cmd = f"qemu-img create -f qcow2 {disk_path} {self.config['storage']['size']}"
        subprocess.run(qemu_img_cmd, shell=True, check=True)
        
        # Start QEMU with installation media
        cmd = self.get_qemu_command(disk_path, iso_path)
        logger.info("Starting OS installation...")
        
        self.vm_process = subprocess.Popen(cmd, shell=True)
        return self.vnc_port

    def get_qemu_command(self, disk_path: str, iso_path: Optional[str] = None) -> str:
        """Generate QEMU command with hardware configuration"""
        cmd = [
            "qemu-system-x86_64",
            "-machine pc,accel=whpx,v=on",  # Use Windows Hypervisor Platform
            f"-smp {self.config['vm']['cpus']}",
            f"-m {self.config['vm']['memory']}",
            
            # Drive configuration
            "-drive",
            f"file={disk_path},if=virtio,format=qcow2",
            
            # Network configuration
            "-netdev user,id=net0",
            "-device virtio-net-pci,netdev=net0",
            
            # Display configuration
            "-display none",
            f"-vnc :{self.vnc_port - 5900}",
            
            # Other options
            "-rtc base=localtime",
            "-usb",
            "-device usb-tablet"
        ]
        
        # Add ISO if installing
        if iso_path:
            cmd.extend([
                "-cdrom", iso_path,
                "-boot", "d"
            ])
            
        return " ".join(cmd)

    async def boot_os(self):
        """Boot existing OS"""
        disk_path = os.path.join(os.path.dirname(__file__), 'disk', 'os.qcow2')
        if not os.path.exists(disk_path):
            raise FileNotFoundError("No OS installation found.")
            
        # Start QEMU
        cmd = self.get_qemu_command(disk_path)
        logger.info("Booting OS...")
        
        self.vm_process = subprocess.Popen(cmd, shell=True)
        return self.vnc_port

    async def shutdown(self):
        """Shutdown running VM"""
        if self.vm_process:
            self.vm_process.terminate()
            await asyncio.sleep(1)  # Give it a moment to shutdown gracefully
            self.vm_process.kill()  # Force kill if still running
            await self.vm_process.wait()
            self.vm_process = None
            logger.info("VM shut down.")