"""
JSON-based storage manager with compression for minimal disk and RAM usage.
Stores all data directly to disk in compressed JSON format with streaming operations.
"""

import json
import os
import zlib
import lzma
from typing import Dict, Any, Optional, Union
from pathlib import Path
import threading
import time
from datetime import datetime
import mmap
import shutil

class DiskStorageManager:
    """
    Disk-based storage manager that minimizes RAM and disk usage by:
    1. Using memory mapping for large files
    2. Streaming JSON operations with compression
    3. Immediate disk syncing
    4. No in-memory caching
    5. Automatic cleanup of old blocks
    6. Multiple compression algorithms
    """
    
    def __init__(self, storage_path: str = "storage", compression: str = "lzma", max_size_gb: float = 10.0):
        self.base_path = Path(storage_path)
        self.base_path.mkdir(exist_ok=True)
        self.lock = threading.Lock()
        self.compression = compression
        self.max_size_gb = max_size_gb
        self.cleanup_thread = threading.Thread(target=self._cleanup_loop, daemon=True)
        self.cleanup_thread.start()
        
    def _get_block_path(self, block_id: str) -> Path:
        """Get path for a specific block"""
        ext = ".xz" if self.compression == "lzma" else ".gz"
        return self.base_path / f"{block_id}{ext}"
        
    def _compress_data(self, data: bytes) -> bytes:
        """Compress data using selected algorithm"""
        if self.compression == "lzma":
            return lzma.compress(data, preset=9)  # Maximum compression
        else:
            return zlib.compress(data, level=9)  # Maximum compression
            
    def _decompress_data(self, data: bytes) -> bytes:
        """Decompress data using selected algorithm"""
        if self.compression == "lzma":
            return lzma.decompress(data)
        else:
            return zlib.decompress(data)
        
    def store_block(self, block_id: str, data: Dict[str, Any]):
        """Store data block directly to disk with compression"""
        path = self._get_block_path(block_id)
        json_data = json.dumps({
            'data': data,
            'timestamp': time.time(),
            'block_id': block_id
        }).encode('utf-8')
        
        with self.lock:
            compressed_data = self._compress_data(json_data)
            with open(path, 'wb') as f:
                f.write(compressed_data)
                f.flush()
                os.fsync(f.fileno())  # Force write to disk
                
        self._check_storage_limit()
                
    def load_block(self, block_id: str) -> Optional[Dict[str, Any]]:
        """Load block from disk using memory mapping for large files and decompression"""
        path = self._get_block_path(block_id)
        if not path.exists():
            return None
            
        try:
            if path.stat().st_size > 1024 * 1024:  # For files > 1MB
                with open(path, 'rb') as f:
                    with mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) as mm:
                        compressed_data = mm.read()
                        json_data = self._decompress_data(compressed_data)
                        return json.loads(json_data.decode('utf-8'))
            else:
                with open(path, 'rb') as f:
                    compressed_data = f.read()
                    json_data = self._decompress_data(compressed_data)
                    return json.loads(json_data.decode('utf-8'))
        except Exception as e:
            print(f"Error loading block {block_id}: {e}")
            return None
                
    def delete_block(self, block_id: str):
        """Delete a block from disk"""
        path = self._get_block_path(block_id)
        if path.exists():
            path.unlink()
            
    def list_blocks(self) -> list[str]:
        """List all available blocks"""
        return [p.stem for p in self.base_path.glob("*.json")]
        
    def _get_storage_size_gb(self) -> float:
        """Get current storage size in GB"""
        total_size = sum(f.stat().st_size for f in self.base_path.glob("*.*"))
        return total_size / (1024 * 1024 * 1024)
        
    def _check_storage_limit(self):
        """Check if storage limit is exceeded and cleanup if needed"""
        if self._get_storage_size_gb() > self.max_size_gb:
            self._cleanup_old_blocks()
            
    def _cleanup_old_blocks(self):
        """Remove oldest blocks until under storage limit"""
        files = [(f, f.stat().st_mtime) for f in self.base_path.glob("*.*")]
        files.sort(key=lambda x: x[1])  # Sort by modification time
        
        while self._get_storage_size_gb() > self.max_size_gb * 0.8:  # Clean until 80% of limit
            if not files:
                break
            f, _ = files.pop(0)  # Remove oldest file
            f.unlink()
            
    def _cleanup_loop(self):
        """Background thread for periodic cleanup"""
        while True:
            time.sleep(60)  # Check every minute
            try:
                self._check_storage_limit()
            except Exception as e:
                print(f"Cleanup error: {e}")
                
    def clear_old_blocks(self, max_age_seconds: float = 3600):
        """Clear blocks older than specified age"""
        current_time = time.time()
        with self.lock:
            for f in self.base_path.glob("*.*"):
                if (current_time - f.stat().st_mtime) > max_age_seconds:
                    f.unlink()
