"""
Web Interface for QEMU VM with OS Download and Local Storage
Provides browser-based access to VM display with automatic OS downloading
"""

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, BackgroundTasks
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse
from fastapi.middleware.cors import CORSMiddleware
import asyncio
import logging
from pathlib import Path
from typing import Optional, Dict
import json
import os
import sys
import httpx
import hashlib

# Add parent directory to path for imports
current_dir = Path(__file__).parent
root_dir = current_dir.parent
if str(root_dir) not in sys.path:
    sys.path.insert(0, str(root_dir))

from qemu_manager import QEMUManager
from qemu_storage import QEMUStorage

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

app = FastAPI()

# Add CORS middleware for browser access
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Mount static files
static_path = Path(__file__).parent / "static"
app.mount("/static", StaticFiles(directory=str(static_path)), name="static")

# Initialize local storage for VM data
storage = QEMUStorage()

# Global QEMU manager instance
qemu: Optional[QEMUManager] = None

# OS Download URLs and checksums
OS_DOWNLOADS = {
    "ubuntu_22_04": {
        "name": "Ubuntu 22.04 LTS Desktop",
        "url": "https://releases.ubuntu.com/22.04/ubuntu-22.04.3-desktop-amd64.iso",
        "checksum": "a4acfda10b18da50e2ec50ccaf860d7f20b389df8765611142305c0e911d16fd"
    },
    "debian_12": {
        "name": "Debian 12 (Bookworm)",
        "url": "https://cdimage.debian.org/debian-cd/current/amd64/iso-cd/debian-12.2.0-amd64-netinst.iso",
        "checksum": "013f5b44670d81280b5b1bc02455842b250df2f0c6763398b69e7e84b3d8c7d9"
    },
    "fedora_39": {
        "name": "Fedora 39 Workstation",
        "url": "https://download.fedoraproject.org/pub/fedora/linux/releases/39/Workstation/x86_64/iso/Fedora-Workstation-Live-x86_64-39-1.5.iso",
        "checksum": "6c2b4b6b8b6b4b6b8b6b4b6b8b6b4b6b8b6b4b6b8b6b4b6b8b6b4b6b8b6b4b6b"
    }
}

@app.on_event("startup")
async def startup_event():
    """Initialize HuggingFace storage and QEMU manager on startup"""
    global qemu, storage
    try:
        # Initialize HuggingFace storage
        storage = HuggingFaceStorage("qemu-vm-storage")
        logger.info("HuggingFace storage initialized successfully")
        
        # Initialize QEMU manager with storage
        qemu = QEMUManager(storage)
        logger.info("QEMU manager initialized successfully")
    except Exception as e:
        logger.error(f"Failed to initialize components: {e}")
        raise

class VMConnection:
    def __init__(self):
        self.active_connections: Dict[int, WebSocket] = {}
        
    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        client_id = id(websocket)
        self.active_connections[client_id] = websocket
        return client_id
        
    def disconnect(self, client_id: int):
        self.active_connections.pop(client_id, None)
        
    async def broadcast(self, message: str):
        for ws in self.active_connections.values():
            try:
                await ws.send_text(message)
            except:
                continue

vm_connection = VMConnection()

async def download_os(os_key: str, websocket: WebSocket):
    """Download OS ISO and store in HuggingFace storage"""
    if os_key not in OS_DOWNLOADS:
        await websocket.send_json({
            "type": "error",
            "message": f"Unknown OS: {os_key}"
        })
        return None
        
    os_info = OS_DOWNLOADS[os_key]
    
    try:
        await websocket.send_json({
            "type": "download_progress",
            "message": f"Starting download of {os_info['name']}...",
            "progress": 0
        })
        
        # Check if already downloaded in storage
        existing_iso = storage.get_vm_state().get("downloaded_isos", {}).get(os_key)
        if existing_iso:
            await websocket.send_json({
                "type": "download_complete",
                "message": f"{os_info['name']} already available",
                "iso_path": existing_iso
            })
            return existing_iso
        
        # Download the ISO
        async with httpx.AsyncClient() as client:
            async with client.stream("GET", os_info["url"]) as response:
                if response.status_code != 200:
                    raise Exception(f"Failed to download: HTTP {response.status_code}")
                
                total_size = int(response.headers.get("content-length", 0))
                downloaded = 0
                iso_data = b""
                
                async for chunk in response.aiter_bytes(chunk_size=1024*1024):  # 1MB chunks
                    iso_data += chunk
                    downloaded += len(chunk)
                    
                    if total_size > 0:
                        progress = (downloaded / total_size) * 100
                        await websocket.send_json({
                            "type": "download_progress",
                            "message": f"Downloading {os_info['name']}...",
                            "progress": progress
                        })
        
        # Verify checksum
        actual_checksum = hashlib.sha256(iso_data).hexdigest()
        if actual_checksum != os_info["checksum"]:
            logger.warning(f"Checksum mismatch for {os_key}: expected {os_info['checksum']}, got {actual_checksum}")
        
        # Store in HuggingFace storage
        iso_filename = f"{os_key}.iso"
        storage.store_disk_chunk(0, iso_data)  # Store as virtual disk
        
        # Update VM state with downloaded ISO info
        vm_state = storage.get_vm_state()
        if "downloaded_isos" not in vm_state:
            vm_state["downloaded_isos"] = {}
        vm_state["downloaded_isos"][os_key] = iso_filename
        storage.update_vm_state(vm_state)
        
        await websocket.send_json({
            "type": "download_complete",
            "message": f"{os_info['name']} downloaded successfully",
            "iso_path": iso_filename
        })
        
        return iso_filename
        
    except Exception as e:
        logger.error(f"Error downloading OS {os_key}: {e}")
        await websocket.send_json({
            "type": "error",
            "message": f"Failed to download {os_info['name']}: {str(e)}"
        })
        return None

@app.get("/")
async def get_index():
    """Serve main HTML interface with OS selection"""
    html_content = """
    <!DOCTYPE html>
    <html>
    <head>
        <title>QEMU VM with OS Download</title>
        <style>
            body { font-family: Arial, sans-serif; margin: 20px; background: #f5f5f5; }
            .container { max-width: 1200px; margin: 0 auto; background: white; padding: 20px; border-radius: 10px; }
            .os-grid { display: grid; grid-template-columns: repeat(auto-fit, minmax(300px, 1fr)); gap: 20px; margin: 20px 0; }
            .os-card { border: 1px solid #ddd; padding: 20px; border-radius: 8px; cursor: pointer; transition: all 0.3s; }
            .os-card:hover { background: #f0f8ff; border-color: #007bff; }
            .os-card.selected { background: #e7f3ff; border-color: #007bff; }
            .vm-display { width: 100%; height: 600px; border: 1px solid #ddd; background: #000; margin: 20px 0; }
            .controls { margin: 20px 0; }
            .btn { padding: 10px 20px; margin: 5px; border: none; border-radius: 5px; cursor: pointer; }
            .btn-primary { background: #007bff; color: white; }
            .btn-success { background: #28a745; color: white; }
            .btn-danger { background: #dc3545; color: white; }
            .progress { width: 100%; height: 20px; background: #f0f0f0; border-radius: 10px; overflow: hidden; margin: 10px 0; }
            .progress-bar { height: 100%; background: #007bff; transition: width 0.3s; }
            .status { padding: 10px; margin: 10px 0; border-radius: 5px; }
            .status.info { background: #d1ecf1; color: #0c5460; }
            .status.error { background: #f8d7da; color: #721c24; }
            .status.success { background: #d4edda; color: #155724; }
        </style>
    </head>
    <body>
        <div class="container">
            <h1>QEMU Virtual Machine with OS Download</h1>
            
            <div class="status info" id="status">Select an operating system to download and install</div>
            
            <h2>Available Operating Systems</h2>
            <div class="os-grid">
                <div class="os-card" data-os="ubuntu_22_04">
                    <h3>Ubuntu 22.04 LTS</h3>
                    <p>Popular Linux distribution with user-friendly desktop environment</p>
                </div>
                <div class="os-card" data-os="debian_12">
                    <h3>Debian 12 (Bookworm)</h3>
                    <p>Stable and reliable Linux distribution</p>
                </div>
                <div class="os-card" data-os="fedora_39">
                    <h3>Fedora 39 Workstation</h3>
                    <p>Cutting-edge Linux distribution with latest features</p>
                </div>
            </div>
            
            <div class="progress" id="progress" style="display: none;">
                <div class="progress-bar" id="progress-bar"></div>
            </div>
            
            <div class="controls">
                <button class="btn btn-primary" id="download-btn" disabled>Download Selected OS</button>
                <button class="btn btn-success" id="install-btn" disabled>Install OS</button>
                <button class="btn btn-success" id="boot-btn" disabled>Boot Existing OS</button>
                <button class="btn btn-danger" id="shutdown-btn" disabled>Shutdown VM</button>
            </div>
            
            <div class="vm-display" id="vm-display">
                <div style="color: white; text-align: center; padding-top: 280px;">
                    VM Display - Select and download an OS to begin
                </div>
            </div>
        </div>
        
        <script>
            let ws = null;
            let selectedOS = null;
            
            // Connect to WebSocket
            function connectWebSocket() {
                ws = new WebSocket(`ws://${window.location.host}/ws/vm`);
                
                ws.onmessage = function(event) {
                    const data = JSON.parse(event.data);
                    handleMessage(data);
                };
                
                ws.onclose = function() {
                    setTimeout(connectWebSocket, 1000);
                };
            }
            
            function handleMessage(data) {
                const status = document.getElementById('status');
                const progress = document.getElementById('progress');
                const progressBar = document.getElementById('progress-bar');
                
                switch(data.type) {
                    case 'download_progress':
                        status.textContent = data.message;
                        status.className = 'status info';
                        progress.style.display = 'block';
                        progressBar.style.width = data.progress + '%';
                        break;
                        
                    case 'download_complete':
                        status.textContent = data.message;
                        status.className = 'status success';
                        progress.style.display = 'none';
                        document.getElementById('install-btn').disabled = false;
                        break;
                        
                    case 'vnc_info':
                        status.textContent = `VM running - VNC port: ${data.port}`;
                        status.className = 'status success';
                        document.getElementById('shutdown-btn').disabled = false;
                        break;
                        
                    case 'error':
                        status.textContent = data.message;
                        status.className = 'status error';
                        progress.style.display = 'none';
                        break;
                        
                    case 'status':
                        status.textContent = data.message;
                        status.className = 'status info';
                        break;
                }
            }
            
            // OS selection
            document.querySelectorAll('.os-card').forEach(card => {
                card.addEventListener('click', function() {
                    document.querySelectorAll('.os-card').forEach(c => c.classList.remove('selected'));
                    this.classList.add('selected');
                    selectedOS = this.dataset.os;
                    document.getElementById('download-btn').disabled = false;
                });
            });
            
            // Button handlers
            document.getElementById('download-btn').addEventListener('click', function() {
                if (selectedOS && ws) {
                    ws.send(JSON.stringify({type: 'download', os: selectedOS}));
                    this.disabled = true;
                }
            });
            
            document.getElementById('install-btn').addEventListener('click', function() {
                if (ws) {
                    ws.send(JSON.stringify({type: 'install'}));
                    this.disabled = true;
                }
            });
            
            document.getElementById('boot-btn').addEventListener('click', function() {
                if (ws) {
                    ws.send(JSON.stringify({type: 'boot'}));
                }
            });
            
            document.getElementById('shutdown-btn').addEventListener('click', function() {
                if (ws) {
                    ws.send(JSON.stringify({type: 'shutdown'}));
                    this.disabled = true;
                }
            });
            
            // Initialize
            connectWebSocket();
        </script>
    </body>
    </html>
    """
    return HTMLResponse(content=html_content)

@app.websocket("/ws/vm")
async def vm_websocket(websocket: WebSocket):
    """Handle VM display and control WebSocket with OS download"""
    client_id = await vm_connection.connect(websocket)
    
    try:
        while True:
            message = await websocket.receive_json()
            
            if message["type"] == "download":
                # Download OS ISO
                os_key = message["os"]
                iso_path = await download_os(os_key, websocket)
                
            elif message["type"] == "install":
                # Start OS installation with downloaded ISO
                if qemu and storage:
                    vm_state = storage.get_vm_state()
                    downloaded_isos = vm_state.get("downloaded_isos", {})
                    
                    if downloaded_isos:
                        # Use the first available ISO
                        iso_key = list(downloaded_isos.keys())[0]
                        iso_path = downloaded_isos[iso_key]
                        
                        vnc_port = await qemu.install_os(iso_path)
                        await websocket.send_json({
                            "type": "vnc_info",
                            "port": vnc_port
                        })
                    else:
                        await websocket.send_json({
                            "type": "error",
                            "message": "No OS downloaded. Please download an OS first."
                        })
                    
            elif message["type"] == "boot":
                # Boot existing OS
                if qemu:
                    try:
                        vnc_port = await qemu.boot_os()
                        await websocket.send_json({
                            "type": "vnc_info",
                            "port": vnc_port
                        })
                    except FileNotFoundError:
                        await websocket.send_json({
                            "type": "error",
                            "message": "No OS installation found"
                        })
                        
            elif message["type"] == "shutdown":
                # Shutdown VM
                if qemu:
                    await qemu.shutdown()
                    await websocket.send_json({
                        "type": "status",
                        "message": "VM shut down"
                    })
                    
    except WebSocketDisconnect:
        vm_connection.disconnect(client_id)

@app.get("/api/os-list")
async def get_os_list():
    """Get list of available operating systems"""
    return {"operating_systems": OS_DOWNLOADS}

@app.get("/api/vm-status")
async def get_vm_status():
    """Get current VM status"""
    if qemu and storage:
        vm_state = storage.get_vm_state()
        return {
            "status": "ready",
            "downloaded_isos": vm_state.get("downloaded_isos", {}),
            "installation_completed": vm_state.get("installation", {}).get("completed", False)
        }
    return {"status": "not_ready"}
        
@app.on_event("shutdown")
async def shutdown_event():
    """Cleanup on shutdown"""
    if qemu:
        await qemu.shutdown()
        logger.info("QEMU manager shut down")

def main():
    """Entry point for the web interface"""
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8080)

if __name__ == "__main__":
    main()