"""
Disk-based Storage Manager for Virtual RAM
Uses filesystem chunks as virtual memory pages
"""

import os
import mmap
import json
from typing import Optional, Dict, Any
import threading
from pathlib import Path
import logging

class DiskStorage:
    def __init__(self, base_path: str = "storage/virtual_ram", page_size: int = 4096):
        """Initialize disk-based storage for virtual RAM"""
        self.base_path = Path(base_path)
        self.page_size = page_size
        self.page_lock = threading.Lock()
        self.page_map: Dict[int, str] = {}  # Maps virtual addresses to file paths
        
        # Create storage directory
        os.makedirs(self.base_path, exist_ok=True)
        
        # Load existing page map if any
        self._load_page_map()
        
    def _load_page_map(self):
        """Load existing page map from disk"""
        map_file = self.base_path / "page_map.json"
        if map_file.exists():
            with open(map_file, 'r') as f:
                self.page_map = {int(k): v for k, v in json.load(f).items()}
                
    def _save_page_map(self):
        """Save page map to disk"""
        map_file = self.base_path / "page_map.json"
        with open(map_file, 'w') as f:
            json.dump(self.page_map, f)
            
    def allocate_page(self, virt_addr: int) -> bool:
        """Allocate a new page at the given virtual address"""
        with self.page_lock:
            if virt_addr in self.page_map:
                return False
                
            # Create page file
            page_file = self.base_path / f"page_{virt_addr:016x}.bin"
            
            try:
                # Create empty page file
                with open(page_file, 'wb') as f:
                    f.write(b'\0' * self.page_size)
                    
                # Update page map
                self.page_map[virt_addr] = str(page_file)
                self._save_page_map()
                return True
                
            except Exception as e:
                logging.error(f"Failed to allocate page: {e}")
                return False
                
    def read_page(self, virt_addr: int) -> Optional[bytes]:
        """Read data from a page"""
        page_file = self.page_map.get(virt_addr)
        if not page_file:
            return None
            
        try:
            with open(page_file, 'rb') as f:
                with mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) as mm:
                    return mm.read()
        except Exception as e:
            logging.error(f"Failed to read page: {e}")
            return None
            
    def write_page(self, virt_addr: int, data: bytes) -> bool:
        """Write data to a page"""
        if len(data) > self.page_size:
            return False
            
        page_file = self.page_map.get(virt_addr)
        if not page_file:
            if not self.allocate_page(virt_addr):
                return False
            page_file = self.page_map[virt_addr]
            
        try:
            with open(page_file, 'r+b') as f:
                with mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_WRITE) as mm:
                    mm.write(data.ljust(self.page_size, b'\0'))
            return True
        except Exception as e:
            logging.error(f"Failed to write page: {e}")
            return False
            
    def free_page(self, virt_addr: int) -> bool:
        """Free a page"""
        with self.page_lock:
            page_file = self.page_map.get(virt_addr)
            if not page_file:
                return False
                
            try:
                # Remove page file
                os.remove(page_file)
                # Update page map
                del self.page_map[virt_addr]
                self._save_page_map()
                return True
            except Exception as e:
                logging.error(f"Failed to free page: {e}")
                return False
                
    def get_total_memory(self) -> int:
        """Get total allocated memory size"""
        return len(self.page_map) * self.page_size
        
    def get_free_space(self) -> int:
        """Get available disk space for virtual memory"""
        return os.statvfs(self.base_path).f_bavail * os.statvfs(self.base_path).f_frsize