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

import os
import json
import logging
import sqlite3
import threading
import time
from typing import Dict, Any, Optional, List
from pathlib import Path
import hashlib

class QEMUStorage:
    def __init__(self, db_path: str = "db/qemu/storage.db"):
        """Initialize QEMU storage with SQLite backend"""
        self.db_path = db_path
        os.makedirs(os.path.dirname(db_path), exist_ok=True)
        self.lock = threading.Lock()
        self._init_storage()

    def _init_storage(self):
        """Initialize storage tables in SQLite database"""
        self.conn = sqlite3.connect(self.db_path, check_same_thread=False)
        
        with self.lock:
            # Virtual disk storage (chunked)
            self.conn.execute("""
                CREATE TABLE IF NOT EXISTS virtual_disk (
                    chunk_id TEXT PRIMARY KEY,
                    offset INTEGER NOT NULL,
                    data BLOB NOT NULL,
                    checksum TEXT NOT NULL,
                    last_modified REAL NOT NULL
                )
            """)

            # CPU state storage
            self.conn.execute("""
                CREATE TABLE IF NOT EXISTS cpu_states (
                    cpu_id TEXT PRIMARY KEY,
                    state BLOB NOT NULL,
                    timestamp REAL NOT NULL,
                    metadata TEXT
                )
            """)

            # VM state storage
            self.conn.execute("""
                CREATE TABLE IF NOT EXISTS vm_state (
                    key TEXT PRIMARY KEY,
                    value TEXT NOT NULL,
                    last_updated REAL NOT NULL
                )
            """)

            # OS image storage
            self.conn.execute("""
                CREATE TABLE IF NOT EXISTS os_images (
                    id TEXT PRIMARY KEY,
                    name TEXT NOT NULL,
                    size INTEGER NOT NULL,
                    checksum TEXT NOT NULL,
                    created_at REAL NOT NULL
                )
            """)

            # OS image chunks for large file storage
            self.conn.execute("""
                CREATE TABLE IF NOT EXISTS os_image_chunks (
                    chunk_id TEXT PRIMARY KEY,
                    image_id TEXT NOT NULL,
                    chunk_index INTEGER NOT NULL,
                    data BLOB NOT NULL,
                    chunk_size INTEGER NOT NULL,
                    FOREIGN KEY (image_id) REFERENCES os_images(id)
                )
            """)
            self.conn.execute("""
                CREATE TABLE IF NOT EXISTS vm_state (
                    key TEXT PRIMARY KEY,
                    value TEXT NOT NULL,
                    last_updated REAL NOT NULL
                )
            """)

            # VM snapshots
            self.conn.execute("""
                CREATE TABLE IF NOT EXISTS vm_snapshots (
                    snapshot_id TEXT PRIMARY KEY,
                    name TEXT NOT NULL,
                    data BLOB NOT NULL,
                    metadata TEXT,
                    created_at REAL NOT NULL
                )
            """)

            # Initialize VM state if empty
            self.conn.execute("""
                INSERT OR IGNORE INTO vm_state (key, value, last_updated)
                VALUES (?, ?, ?)
            """, ('status', json.dumps({
                'status': 'initialized',
                'installation': {
                    'completed': False,
                    'iso_url': None,
                    'checksum': None
                }
            }), time.time()))

            self.conn.commit()

    def write_disk_chunk(self, chunk_id: str, offset: int, data: bytes):
        """Write a chunk of data to virtual disk"""
        checksum = hashlib.sha256(data).hexdigest()
        
        with self.lock:
            self.conn.execute("""
                INSERT OR REPLACE INTO virtual_disk 
                (chunk_id, offset, data, checksum, last_modified)
                VALUES (?, ?, ?, ?, ?)
            """, (chunk_id, offset, data, checksum, time.time()))
            self.conn.commit()

    def read_disk_chunk(self, chunk_id: str) -> Optional[Dict[str, Any]]:
        """Read a chunk of data from virtual disk"""
        with self.lock:
            result = self.conn.execute("""
                SELECT offset, data, checksum, last_modified
                FROM virtual_disk
                WHERE chunk_id = ?
            """, (chunk_id,)).fetchone()
            
        if result:
            return {
                'chunk_id': chunk_id,
                'offset': result[0],
                'data': result[1],
                'checksum': result[2],
                'last_modified': result[3]
            }
        return None

    def save_cpu_state(self, cpu_id: str, state: bytes, metadata: Optional[Dict] = None):
        """Save CPU state"""
        with self.lock:
            self.conn.execute("""
                INSERT OR REPLACE INTO cpu_states (cpu_id, state, timestamp, metadata)
                VALUES (?, ?, ?, ?)
            """, (cpu_id, state, time.time(), 
                 json.dumps(metadata) if metadata else None))
            self.conn.commit()

    def load_cpu_state(self, cpu_id: str) -> Optional[Dict[str, Any]]:
        """Load CPU state"""
        with self.lock:
            result = self.conn.execute("""
                SELECT state, timestamp, metadata
                FROM cpu_states
                WHERE cpu_id = ?
            """, (cpu_id,)).fetchone()
            
        if result:
            return {
                'cpu_id': cpu_id,
                'state': result[0],
                'timestamp': result[1],
                'metadata': json.loads(result[2]) if result[2] else None
            }
        return None

    def update_vm_state(self, key: str, value: Any):
        """Update VM state"""
        with self.lock:
            self.conn.execute("""
                INSERT OR REPLACE INTO vm_state (key, value, last_updated)
                VALUES (?, ?, ?)
            """, (key, json.dumps(value), time.time()))
            self.conn.commit()

    def get_vm_state(self, key: str) -> Optional[Any]:
        """Get VM state"""
        with self.lock:
            result = self.conn.execute("""
                SELECT value
                FROM vm_state
                WHERE key = ?
            """, (key,)).fetchone()
            
        if result:
            return json.loads(result[0])
        return None

    def create_snapshot(self, name: str, data: bytes, metadata: Optional[Dict] = None):
        """Create VM snapshot"""
        snapshot_id = hashlib.sha256(f"{name}_{time.time()}".encode()).hexdigest()
        
        with self.lock:
            self.conn.execute("""
                INSERT INTO vm_snapshots 
                (snapshot_id, name, data, metadata, created_at)
                VALUES (?, ?, ?, ?, ?)
            """, (snapshot_id, name, data, 
                 json.dumps(metadata) if metadata else None,
                 time.time()))
            self.conn.commit()
            
        return snapshot_id

    def load_snapshot(self, snapshot_id: str) -> Optional[Dict[str, Any]]:
        """Load VM snapshot"""
        with self.lock:
            result = self.conn.execute("""
                SELECT name, data, metadata, created_at
                FROM vm_snapshots
                WHERE snapshot_id = ?
            """, (snapshot_id,)).fetchone()
            
        if result:
            return {
                'snapshot_id': snapshot_id,
                'name': result[0],
                'data': result[1],
                'metadata': json.loads(result[2]) if result[2] else None,
                'created_at': result[3]
            }
        return None

    def list_snapshots(self) -> List[Dict[str, Any]]:
        """List all VM snapshots"""
        with self.lock:
            results = self.conn.execute("""
                SELECT snapshot_id, name, metadata, created_at
                FROM vm_snapshots
                ORDER BY created_at DESC
            """).fetchall()
            
        return [{
            'snapshot_id': r[0],
            'name': r[1],
            'metadata': json.loads(r[2]) if r[2] else None,
            'created_at': r[3]
        } for r in results]

    def delete_snapshot(self, snapshot_id: str):
        """Delete VM snapshot"""
        with self.lock:
            self.conn.execute("""
                DELETE FROM vm_snapshots
                WHERE snapshot_id = ?
            """, (snapshot_id,))
            self.conn.commit()

    def update_vm_state(self, key: str, value: Any):
        """Update VM state"""
        with self.lock:
            self.conn.execute("""
                INSERT OR REPLACE INTO vm_state (key, value, last_updated)
                VALUES (?, ?, ?)
            """, (key, json.dumps(value), time.time()))
            self.conn.commit()

    def get_vm_state(self, key: str) -> Optional[Any]:
        """Get VM state by key"""
        with self.lock:
            result = self.conn.execute("""
                SELECT value FROM vm_state WHERE key = ?
            """, (key,)).fetchone()
            
        if result:
            try:
                return json.loads(result[0])
            except json.JSONDecodeError:
                return None
        return None

    def store_os_image_chunk(self, image_id: str, chunk_index: int, data: bytes):
        """Store a chunk of OS image"""
        chunk_id = hashlib.sha256(f"{image_id}_{chunk_index}".encode()).hexdigest()
        
        with self.lock:
            # Ensure image entry exists
            exists = self.conn.execute("""
                SELECT 1 FROM os_images WHERE id = ?
            """, (image_id,)).fetchone()
            
            if not exists:
                self.conn.execute("""
                    INSERT INTO os_images (id, name, size, checksum, created_at)
                    VALUES (?, ?, ?, ?, ?)
                """, (image_id, "os_image", 0, "", time.time()))
            
            self.conn.execute("""
                INSERT OR REPLACE INTO os_image_chunks (chunk_id, image_id, chunk_index, data, chunk_size)
                VALUES (?, ?, ?, ?, ?)
            """, (chunk_id, image_id, chunk_index, data, len(data)))
            self.conn.commit()

    def get_os_image(self, image_id: str) -> Optional[bytes]:
        """Retrieve full OS image"""
        with self.lock:
            # Get all chunks in order
            chunks = self.conn.execute("""
                SELECT data FROM os_image_chunks 
                WHERE image_id = ? 
                ORDER BY chunk_index
            """, (image_id,)).fetchall()
            
            if not chunks:
                return None
            
            # Combine chunks
            return b''.join(chunk[0] for chunk in chunks)

    def close(self):
        """Close database connection"""
        if hasattr(self, 'conn'):
            self.conn.close()

    def __del__(self):
        self.close()
