
import threading
import duckdb
from huggingface_hub import HfApi, HfFileSystem

class RAMController:
    DB_URL = "hf://datasets/Fred808/helium/storage.json"
    
    def __init__(self, size_bytes):
        self.size_bytes = size_bytes
        self._connect_with_retries()
        self.db_lock = threading.Lock()
        self._setup_database()
        
    def _connect_with_retries(self, max_retries=3):
        """Establish database connection with retry logic"""
        for attempt in range(max_retries):
            try:
                self.conn = self._init_db_connection()
                return
            except Exception as e:
                if attempt == max_retries - 1:
                    raise RuntimeError(f"Failed to initialize database after {max_retries} attempts: {str(e)}")
                time.sleep(1)
                
    def _init_db_connection(self):
        """Initialize database connection with HuggingFace configuration"""
        # Connect directly to HuggingFace URL
        # First create an in-memory connection to configure settings
        temp_conn = duckdb.connect(":memory:")
        
        # Configure HuggingFace access - must be done before connecting to URL
        temp_conn.execute("INSTALL httpfs;")
        temp_conn.execute("LOAD httpfs;")
        temp_conn.execute("SET s3_endpoint='hf.co';")
        temp_conn.execute("SET s3_use_ssl=true;")
        temp_conn.execute("SET s3_url_style='path';")
        
        # Now create the real connection with the configured settings
        conn = duckdb.connect(self.DB_URL, config={'http_keep_alive': 'true'})
        conn.execute("INSTALL httpfs;")
        conn.execute("LOAD httpfs;")
        conn.execute("SET s3_endpoint='hf.co';")
        conn.execute("SET s3_use_ssl=true;")
        conn.execute("SET s3_url_style='path';")
        
        # Close temporary connection
        temp_conn.close()
        return conn
        
    def _setup_database(self):
        """Initialize database schema"""
        with self.db_lock:
            self.conn.execute("""
                CREATE TABLE IF NOT EXISTS ram_cells (
                    address BIGINT PRIMARY KEY,
                    data BLOB
                )
            """)

    def read(self, address, length):
        if address < 0 or address + length > self.size_bytes:
            raise IndexError("Memory access out of bounds")
        with self.db_lock:
            result = bytearray([0] * length)
            # DuckDB query to get memory contents
            rows = self.conn.execute("""
                SELECT address, data 
                FROM ram_cells 
                WHERE address >= ? AND address < ? 
                ORDER BY address ASC
            """, [address, address + length]).fetchall()
            
            for row in rows:
                addr, data = row
                if address <= addr < address + length:
                    # Convert BLOB data to byte
                    result[addr - address] = data[0] if isinstance(data, (bytes, bytearray)) else data
            return result

    def write(self, address, data):
        if address < 0 or address + len(data) > self.size_bytes:
            raise IndexError("Memory access out of bounds")
        with self.db_lock:
            # Prepare batch insert data
            values = [(address + offset, bytes([value])) 
                     for offset, value in enumerate(data)]
            # DuckDB batch insert/update
            self.conn.execute("""
                INSERT OR REPLACE INTO ram_cells (address, data)
                SELECT * FROM (VALUES ?);
            """, [values])

    def close(self):
        """Close the database connection"""
        with self.db_lock:
            if self.conn:
                self.conn.close()
                self.conn = None
                
    def ensure_connection(self):
        """Ensure database connection is active and valid"""
        try:
            self.conn.execute("SELECT 1")
        except:
            self._connect_with_retries()


