http_storage.py

import json
import numpy as np
from typing import Dict, Any, Optional, Union
import threading
import time
import hashlib
import logging
import uuid
import duckdb
import os
from datetime import datetime
from huggingface_hub import HfApi, HfFileSystem
from config import get_hf_token_cached



class LocalStorage:
    """
    Remote storage implementation using DuckDB and HuggingFace.
    Provides efficient distributed storage and querying capabilities.
    No local filesystem dependencies.
    """
    
    # Singleton instance 
    _instance = None
    _lock = threading.Lock()

    def __new__(cls, db_url: str = "hf://datasets/Fred808/helium/storage.json"):
        with cls._lock:
            if cls._instance is None:
                cls._instance = super().__new__(cls)
                cls._instance._init_singleton(db_url)
            return cls._instance

    def _init_singleton(self, db_url: str):
        if hasattr(self, 'initialized'):
            return
            
        # Setup connection identifier
        self.storage_id = hashlib.md5(db_url.encode()).hexdigest()[:8]
            
        # Setup DuckDB connection
        self.db_url = db_url
        if db_url.startswith('hf://'):
            # Connect directly to HuggingFace dataset
            # Format: hf://datasets/Fred808/helium/storage.json
            _, _, owner, dataset, db_file = db_url.split('/', 4)
            db_path = f"s3://datasets-cached/{owner}/{dataset}/{db_file}"
            print(f"Connecting to database at: {db_path}")
            
            # Get token from environment
            self.hf_token = get_hf_token_cached()
            
            self.conn = duckdb.connect(db_path)
            self.conn.execute("""
                INSTALL httpfs;
                LOAD httpfs;
                SET s3_endpoint='hf.co';
                SET s3_use_ssl=true;
                SET s3_url_style='path';
            """)
            # Configure HuggingFace authentication
            self.conn.execute(f"SET s3_access_key_id='{self.hf_token}';")
            self.conn.execute(f"SET s3_secret_access_key='{self.hf_token}';")
            self.dataset_path = db_path
        else:
            # Remote database
            print(f"Connecting to database at: {db_url}")
            self.conn = duckdb.connect(db_url)
        
        # Basic state management
        self.lock = threading.Lock()
        self._closing = False
        
        # Resource monitoring
        self.resource_monitor = {
            'vram_used': 0,
            'active_tensors': 0,
            'loaded_models': set(),
            'last_updated': time.time()
        }
        
        # Storage statistics
        self.stats = {
            'total_size': 0,
            'available_size': float('inf'),
            'model_count': 0,
            'tensor_count': 0
        }
        
        # Initialize database
        self._init_database()
        
        # Initialize registries and state
        self.model_registry = {}
        self.tensor_registry = {}
        self._connected = True
        self.initialized = True
        
    def _init_database(self):
        """Initialize DuckDB database with required tables"""
        
        # Enable required extensions
        self.conn.execute("""
            INSTALL json;
            LOAD json;
            INSTALL httpfs;
            LOAD httpfs;
        """)
        
        # Create VRAM blocks table
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS vram_blocks (
                block_id VARCHAR PRIMARY KEY,
                size BIGINT,
                data BLOB,  -- Store tensor data directly
                metadata JSON,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                last_accessed TIMESTAMP,
                is_pinned BOOLEAN DEFAULT FALSE,
                device_id VARCHAR
            )
        """)
        
        # Create models table
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS models (
                model_id VARCHAR PRIMARY KEY,
                name VARCHAR,
                version VARCHAR,
                data BLOB,  -- Store model data directly
                metadata JSON,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                last_accessed TIMESTAMP,
                is_loaded BOOLEAN DEFAULT FALSE,
                config JSON
            )
        """)
        
        # Create cache table
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS cache (
                cache_id VARCHAR PRIMARY KEY,
                key VARCHAR UNIQUE,
                data BLOB,  -- Store cached data directly
                metadata JSON,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                accessed_at TIMESTAMP,
                expires_at TIMESTAMP,
                size BIGINT
            )
        """)
        
        # Create states table
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS states (
                state_id VARCHAR PRIMARY KEY,
                name VARCHAR,
                data BLOB,  -- Store state data directly
                metadata JSON,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP,
                parent_id VARCHAR,
                is_checkpoint BOOLEAN DEFAULT FALSE
            )
        """)
        
        # Create tensor operations table
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS tensor_ops (
                op_id VARCHAR PRIMARY KEY,
                core_id VARCHAR,
                operation_type VARCHAR,
                input_tensors JSON,
                output_tensors JSON,
                metadata JSON,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                completed_at TIMESTAMP,
                status VARCHAR,
                execution_time_ns BIGINT
            )
        """)
        
        # Create tensor core states table
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS tensor_core_states (
                core_id VARCHAR PRIMARY KEY,
                array_id VARCHAR,
                current_op VARCHAR,
                register_state JSON,
                shared_memory_state JSON,
                metadata JSON,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP,
                status VARCHAR,
                is_active BOOLEAN
            )
        """)
        
        # Create core communication table
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS core_communication (
                comm_id VARCHAR PRIMARY KEY,
                source_core_id VARCHAR,
                target_core_id VARCHAR,
                data_id VARCHAR,
                metadata JSON,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                completed_at TIMESTAMP,
                status VARCHAR,
                transfer_size_bytes BIGINT
            )
        """)

        # Create indices for faster lookups
        self.conn.execute("CREATE INDEX IF NOT EXISTS idx_vram_blocks_device ON vram_blocks(device_id)")
        self.conn.execute("CREATE INDEX IF NOT EXISTS idx_vram_blocks_accessed ON vram_blocks(last_accessed)")
        self.conn.execute("CREATE INDEX IF NOT EXISTS idx_models_name ON models(name)")
        self.conn.execute("CREATE INDEX IF NOT EXISTS idx_models_loaded ON models(is_loaded)")
        self.conn.execute("CREATE INDEX IF NOT EXISTS idx_cache_accessed ON cache(accessed_at)")
        self.conn.execute("CREATE INDEX IF NOT EXISTS idx_cache_expires ON cache(expires_at)")
        self.conn.execute("CREATE INDEX IF NOT EXISTS idx_states_parent ON states(parent_id)")
        self.conn.execute("CREATE INDEX IF NOT EXISTS idx_states_updated ON states(updated_at)")
        self.conn.execute("CREATE INDEX IF NOT EXISTS idx_tensor_ops_core ON tensor_ops(core_id)")
        self.conn.execute("CREATE INDEX IF NOT EXISTS idx_tensor_ops_status ON tensor_ops(status)")
        self.conn.execute("CREATE INDEX IF NOT EXISTS idx_core_states_array ON tensor_core_states(array_id)")
        self.conn.execute("CREATE INDEX IF NOT EXISTS idx_core_states_active ON tensor_core_states(is_active)")
        self.conn.execute("CREATE INDEX IF NOT EXISTS idx_comm_source ON core_communication(source_core_id)")
        self.conn.execute("CREATE INDEX IF NOT EXISTS idx_comm_target ON core_communication(target_core_id)")
        self.conn.execute("CREATE INDEX IF NOT EXISTS idx_comm_status ON core_communication(status)")
        
        self.conn.commit()
            
    def _sync_to_huggingface(self):
        """Periodically sync the database to HuggingFace if using HF storage"""
        while not self._closing:
            time.sleep(300)  # Sync every 5 minutes
            if hasattr(self, 'hf_api') and not self._closing:
                try:
                    # Close connection temporarily to ensure all changes are written
                    self.conn.close()
                    
                    # Upload the database file to HuggingFace
                    _, _, owner, dataset, db_file = self.db_url.split('/', 4)
                    self.hf_api.upload_file(
                        path_or_fileobj=str(self.db_path),
                        path_in_repo=db_file,
                        repo_id=f"datasets/{owner}/{dataset}",
                        repo_type="dataset"
                    )
                    
                    # Reconnect to the database
                    self.conn = duckdb.connect(str(self.db_path))
                    self.conn.execute("""
                        INSTALL httpfs;
                        LOAD httpfs;
                        SET s3_endpoint='hf.co';
                        SET s3_use_ssl=true;
                        SET s3_url_style='path';
                    """)
                    # Configure HuggingFace authentication
                    self.conn.execute(f"SET s3_access_key_id='{self.hf_token}';")
                    self.conn.execute(f"SET s3_secret_access_key='{self.hf_token}';")
                except Exception as e:
                    logging.error(f"Failed to sync database to HuggingFace: {e}")

    def _store_in_db(self, table: str, data_id: str, data: Union[bytes, np.ndarray], metadata: Dict = None, **kwargs):
        """Store entry in database using DuckDB"""
        metadata_json = json.dumps(metadata) if metadata else None
        
        # Convert numpy arrays to bytes if needed
        if isinstance(data, np.ndarray):
            data = data.tobytes()
            
        # Build dynamic query based on table
        fields = ['data', 'metadata']
        values = [data, metadata_json]
        
        # Add additional fields from kwargs
        for key, value in kwargs.items():
            fields.append(key)
            values.append(value)
            
        # Build query with proper column name for id based on table
        id_column = f"{table[:-1]}_id" if table.endswith('s') else 'id'
        fields_str = ','.join([id_column] + fields)
        placeholders = ','.join(['?' for _ in range(len(fields) + 1)])
        
        # Use REPLACE function of DuckDB
        query = f"""
            DELETE FROM {table} WHERE {id_column} = ?;
            INSERT INTO {table} ({fields_str}) VALUES ({placeholders});
        """
        
        # Execute as transaction
        self.conn.execute("BEGIN TRANSACTION")
        try:
            # Only pass data_id once since we're using it as a single parameter
            self.conn.execute(query, [data_id] + values)
            self.conn.commit()
        except Exception as e:
            self.conn.rollback()
            raise e
            
    def _get_from_db(self, table: str, data_id: str) -> Optional[Dict]:
        """Retrieve entry from database using DuckDB"""
        id_column = f"{table[:-1]}_id" if table.endswith('s') else 'id'
        query = f"SELECT * FROM {table} WHERE {id_column} = ?"
        
        result = self.conn.execute(query, [data_id]).fetchone()
        if result:
            # Convert to dict
            columns = self.conn.execute(f"DESCRIBE {table}").fetchall()
            column_names = [col[0] for col in columns]
            result_dict = dict(zip(column_names, result))
            
            # Parse JSON metadata if present
            if result_dict.get('metadata'):
                result_dict['metadata'] = json.loads(result_dict['metadata'])
            return result_dict
        return None
        
    def is_connected(self) -> bool:
        """Check if storage is connected (always True for local storage)"""
        return self._connected and not self._closing and self.ping()
        
    def close(self):
        """Close storage connection"""
        self._closing = True
        self._connected = False
        
        # Initialize resource monitoring
        self.resource_monitor = {
            'vram_used': 0,
            'active_tensors': 0,
            'loaded_models': set(),
            'last_updated': time.time()
        }
        
        # Initialize model registry and connection state
        self.model_registry = {}
        self._connected = True
        self.model_registry = {}
        self._connected = True
        
    def is_model_loaded(self, model_id: str) -> bool:
        """Check if a model is loaded in storage"""
        if not model_id:
            return False
            
        # Query the models table
        result = self.conn.execute(
            "SELECT is_loaded FROM models WHERE model_id = ?",
            [model_id]
        ).fetchone()
        
        return bool(result[0]) if result else False
        
    def wait_for_connection(self, timeout: float = 30.0) -> bool:
        """Wait for database connection to be ready"""
        end_time = time.time() + timeout
        while time.time() < end_time:
            if self._check_storage_ready():
                return True
            time.sleep(0.5)
        return False
        
    def __init__(self, db_url: str = None):
        """This will actually just return the singleton instance. 
        The actual initialization happens in __new__ and _init_singleton"""
        pass

    def _check_storage_ready(self) -> bool:
        """Check if storage is ready for use"""
        try:
            # Test database connection
            result = self.conn.execute("SELECT 1").fetchone()
            if not result or result[0] != 1:
                return False
                
            # Update storage statistics
            self.stats.update({
                'model_count': self.conn.execute("SELECT COUNT(*) FROM models").fetchone()[0],
                'tensor_count': self.conn.execute("SELECT COUNT(*) FROM vram_blocks").fetchone()[0],
                'total_size': self.conn.execute(
                    "SELECT COALESCE(SUM(size), 0) FROM vram_blocks"
                ).fetchone()[0]
            })
            
            return True
            
        except Exception as e:
            logging.error(f"Storage check failed: {e}")
            return False

    def _check_storage(self) -> Dict[str, Any]:
        """Check storage status and usage"""
        try:
            # Get storage statistics from database
            stats = self.conn.execute("""
                SELECT 
                    COALESCE(SUM(size), 0) as total_size,
                    COUNT(*) as block_count
                FROM vram_blocks
            """).fetchone()
            
            self.storage_monitor.update({
                'total_size': stats[0],
                'last_access': time.time(),
                'block_count': stats[1]
            })
            return {"status": "ok", "monitor": self.storage_monitor}
        except Exception as e:
            logging.error(f"Error checking storage: {e}")
            return {"status": "error", "message": str(e)}

    def store_tensor(self, tensor_id: str, data: np.ndarray, model_size: Optional[int] = None) -> bool:
        """Store tensor data in database"""
        try:
            if data is None:
                raise ValueError("Cannot store None tensor")
                
            # Calculate tensor metadata
            tensor_shape = data.shape
            tensor_dtype = str(data.dtype)
            tensor_size = data.nbytes
            
            # Convert tensor to bytes for storage
            tensor_bytes = data.tobytes()
            
            # Prepare metadata
            metadata = {
                'shape': tensor_shape,
                'dtype': tensor_dtype,
                'size': tensor_size,
                'timestamp': time.time(),
                'model_size': model_size if model_size is not None else -1
            }
            
            # Store in database
            self.conn.execute("""
                INSERT INTO vram_blocks (
                    block_id, data, metadata, size, created_at, last_accessed
                ) VALUES (?, ?, ?, ?, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
                ON CONFLICT (block_id) DO UPDATE SET
                    data = excluded.data,
                    metadata = excluded.metadata,
                    size = excluded.size,
                    last_accessed = CURRENT_TIMESTAMP
            """, [tensor_id, tensor_bytes, json.dumps(metadata), tensor_size])
            
            # Update resource monitor
            with self.lock:
                self.resource_monitor['vram_used'] += tensor_size
                self.resource_monitor['active_tensors'] += 1
            return True
                
        except Exception as e:
            logging.error(f"Error storing tensor {tensor_id}: {str(e)}")
            return False

    def load_tensor(self, tensor_id: str) -> Optional[np.ndarray]:
        """Load tensor data from database"""
        try:
            # Get tensor from database
            result = self.conn.execute("""
                SELECT data, metadata
                FROM vram_blocks 
                WHERE block_id = ?
            """, [tensor_id]).fetchone()
            
            if not result:
                logging.warning(f"Tensor {tensor_id} not found in database")
                return None
                
            tensor_bytes, metadata_str = result
            metadata = json.loads(metadata_str)
            
            # Reconstruct numpy array
            arr = np.frombuffer(tensor_bytes, dtype=metadata['dtype'])
            arr = arr.reshape(metadata['shape'])
            
            # Update access time
            self.conn.execute("""
                UPDATE vram_blocks 
                SET last_accessed = CURRENT_TIMESTAMP 
                WHERE block_id = ?
            """, [tensor_id])
            
            # Update resource monitor
            with self.lock:
                if tensor_id not in self.tensor_registry:
                    self.tensor_registry[tensor_id] = metadata
            
            return arr
                
        except Exception as e:
            logging.error(f"Error loading tensor {tensor_id}: {str(e)}")
            return None

    def store_state(self, component: str, state_id: str, state_data: Dict[str, Any]) -> bool:
        """Store component state in database"""
        try:
            # Store state in database
            metadata = {
                'component': component,
                'timestamp': time.time()
            }
            
            self.conn.execute("""
                INSERT INTO states (
                    state_id, name, data, metadata, created_at, updated_at
                ) VALUES (?, ?, ?, ?, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
                ON CONFLICT (state_id) DO UPDATE SET
                    data = excluded.data,
                    metadata = excluded.metadata,
                    updated_at = CURRENT_TIMESTAMP
            """, [state_id, component, json.dumps(state_data), json.dumps(metadata)])
            
            self.conn.commit()
            
            return True
                
        except Exception as e:
            logging.error(f"Error storing state for {component}/{state_id}: {str(e)}")
            return False

    def load_state(self, component: str, state_id: str) -> Optional[Dict[str, Any]]:
        """Load component state from database"""
        try:
            # Get state from database
            result = self.conn.execute("""
                SELECT data 
                FROM states 
                WHERE state_id = ? AND name = ?
            """, [state_id, component]).fetchone()
            
            if not result:
                logging.warning(f"State not found for {component}/{state_id}")
                return None
            
            # Update access time
            self.conn.execute("""
                UPDATE states 
                SET updated_at = CURRENT_TIMESTAMP 
                WHERE state_id = ?
            """, [state_id])
            
            return json.loads(result[0])
                
        except Exception as e:
            logging.error(f"Error loading state for {component}/{state_id}: {str(e)}")
            return None

    
    def load_model(self, model_name: str, model_data: Optional[Union[bytes, Dict]] = None, model_config: Optional[Dict] = None) -> bool:
        """Load a model into storage"""
        try:
            # Check if model is already loaded
            if self.is_model_loaded(model_name):
                logging.info(f"Model {model_name} already loaded")
                return True

            # Store model in database
            model_id = hashlib.md5(model_name.encode()).hexdigest()
            
            # Convert dict to bytes if needed
            if isinstance(model_data, dict):
                model_data = json.dumps(model_data).encode()
            
            self.conn.execute("""
                INSERT INTO models (
                    model_id, name, version, data, metadata, config, 
                    created_at, last_accessed, is_loaded
                ) VALUES (?, ?, ?, ?, ?, ?, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP, TRUE)
                ON CONFLICT (model_id) DO UPDATE SET
                    data = excluded.data,
                    metadata = excluded.metadata,
                    config = excluded.config,
                    last_accessed = CURRENT_TIMESTAMP,
                    is_loaded = TRUE
            """, [
                model_id, 
                model_name, 
                "1.0",  # Version can be updated if needed
                model_data or b"",
                json.dumps({"source": "direct_load"}),
                json.dumps(model_config) if model_config else "{}"
            ])
            
            # Update model registry
            with self.lock:
                self.model_registry[model_name] = {
                    'id': model_id,
                    'loaded': True,
                    'last_access': time.time()
                }
            
            logging.info(f"Successfully loaded model {model_name}")
            return True
                
        except Exception as e:
            logging.error(f"Error loading model {model_name}: {str(e)}")
            return False


    def ping(self) -> bool:
        """Check if storage is accessible"""
        if self._closing:
            return False
        return self._check_storage_ready()
# Compatibility aliases for existing code