"""
QEMU Storage Backend using Local SQLite Storage
Provides virtual disk and state storage using local SQLite database
"""

import os
import json
import logging
import sqlite3
from typing import Dict, Any, Optional
from pathlib import Path
import threading
import numpy as np

class QEMULocalStorage:
    """Local storage implementation for QEMU state and disk data"""
    
    def __init__(self, db_path: Optional[str] = None):
        self.db_path = db_path or os.path.join('data', 'qemu_storage.db')
        os.makedirs(os.path.dirname(self.db_path), exist_ok=True)
        
        self.lock = threading.Lock()
        self.conn = sqlite3.connect(self.db_path, check_same_thread=False)
        self.cursor = self.conn.cursor()
        self._init_db()
        
    def _init_db(self):
        """Initialize database tables"""
        with self.lock:
            # Table for VM disk states
            self.cursor.execute("""
                CREATE TABLE IF NOT EXISTS disk_states (
                    disk_id TEXT PRIMARY KEY,
                    data BLOB,
                    metadata TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)
            
            # Table for VM memory states
            self.cursor.execute("""
                CREATE TABLE IF NOT EXISTS memory_states (
                    memory_id TEXT PRIMARY KEY,
                    data BLOB,
                    metadata TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)
            
            # Table for VM device states
            self.cursor.execute("""
                CREATE TABLE IF NOT EXISTS device_states (
                    device_id TEXT,
                    state_id TEXT,
                    data TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    PRIMARY KEY (device_id, state_id)
                )
            """)
            
            self.conn.commit()
    
    def store_disk_state(self, disk_id: str, data: bytes, metadata: Optional[Dict] = None) -> bool:
        """Store disk state in database"""
        try:
            with self.lock:
                self.cursor.execute("""
                    INSERT OR REPLACE INTO disk_states (disk_id, data, metadata, updated_at)
                    VALUES (?, ?, ?, CURRENT_TIMESTAMP)
                """, (
                    disk_id,
                    data,
                    json.dumps(metadata) if metadata else None
                ))
                self.conn.commit()
                return True
        except Exception as e:
            logging.error(f"Error storing disk state {disk_id}: {e}")
            return False
    
    def load_disk_state(self, disk_id: str) -> Optional[bytes]:
        """Load disk state from database"""
        try:
            with self.lock:
                self.cursor.execute("SELECT data FROM disk_states WHERE disk_id = ?", (disk_id,))
                row = self.cursor.fetchone()
                return row[0] if row else None
        except Exception as e:
            logging.error(f"Error loading disk state {disk_id}: {e}")
            return None
    
    def store_memory_state(self, memory_id: str, data: bytes, metadata: Optional[Dict] = None) -> bool:
        """Store memory state in database"""
        try:
            with self.lock:
                self.cursor.execute("""
                    INSERT OR REPLACE INTO memory_states (memory_id, data, metadata, updated_at)
                    VALUES (?, ?, ?, CURRENT_TIMESTAMP)
                """, (
                    memory_id,
                    data,
                    json.dumps(metadata) if metadata else None
                ))
                self.conn.commit()
                return True
        except Exception as e:
            logging.error(f"Error storing memory state {memory_id}: {e}")
            return False
    
    def load_memory_state(self, memory_id: str) -> Optional[bytes]:
        """Load memory state from database"""
        try:
            with self.lock:
                self.cursor.execute("SELECT data FROM memory_states WHERE memory_id = ?", (memory_id,))
                row = self.cursor.fetchone()
                return row[0] if row else None
        except Exception as e:
            logging.error(f"Error loading memory state {memory_id}: {e}")
            return None
            
    def store_device_state(self, device_id: str, state_id: str, state: Dict) -> bool:
        """Store device state in database"""
        try:
            with self.lock:
                self.cursor.execute("""
                    INSERT OR REPLACE INTO device_states (device_id, state_id, data)
                    VALUES (?, ?, ?)
                """, (
                    device_id,
                    state_id,
                    json.dumps(state)
                ))
                self.conn.commit()
                return True
        except Exception as e:
            logging.error(f"Error storing device state {device_id}/{state_id}: {e}")
            return False
    
    def load_device_state(self, device_id: str, state_id: str) -> Optional[Dict]:
        """Load device state from database"""
        try:
            with self.lock:
                self.cursor.execute("""
                    SELECT data FROM device_states 
                    WHERE device_id = ? AND state_id = ?
                """, (device_id, state_id))
                row = self.cursor.fetchone()
                return json.loads(row[0]) if row else None
        except Exception as e:
            logging.error(f"Error loading device state {device_id}/{state_id}: {e}")
            return None
    
    def cleanup(self):
        """Clean up resources"""
        with self.lock:
            if hasattr(self, 'conn'):
                self.conn.close()
