import sqlite3
from typing import Dict, Any
import json
import os

class QEMUMemoryDB:
    def __init__(self, db_path: str = "db/qemu_memory.db"):
        """Initialize memory database for QEMU VM"""
        self.db_path = db_path
        os.makedirs(os.path.dirname(db_path), exist_ok=True)
        
        self.conn = sqlite3.connect(db_path)
        self.cursor = self.conn.cursor()
        
        # Create tables if they don't exist
        self._init_db()
    
    def _init_db(self):
        """Initialize database schema"""
        self.cursor.execute('''
        CREATE TABLE IF NOT EXISTS cpu_memory (
            key TEXT PRIMARY KEY,
            data BLOB,
            metadata TEXT
        )
        ''')
        self.conn.commit()
    
    def write(self, key: str, data: bytes, metadata: Dict[str, Any] = None):
        """Write data to memory"""
        self.cursor.execute(
            'INSERT OR REPLACE INTO cpu_memory (key, data, metadata) VALUES (?, ?, ?)',
            (key, data, json.dumps(metadata or {}))
        )
        self.conn.commit()
    
    def read(self, key: str) -> tuple[bytes, Dict[str, Any]]:
        """Read data from memory"""
        self.cursor.execute('SELECT data, metadata FROM cpu_memory WHERE key = ?', (key,))
        row = self.cursor.fetchone()
        if row:
            return row[0], json.loads(row[1])
        return None, {}
    
    def delete(self, key: str):
        """Delete data from memory"""
        self.cursor.execute('DELETE FROM cpu_memory WHERE key = ?', (key,))
        self.conn.commit()