wd666/utils/storage.py

"""
Storage Manager - Handle local persistence of configuration, history/reports, and assets.
"""
import os
import json
import time
from typing import List, Dict, Any
from pathlib import Path

# Constants
STORAGE_DIR = ".storage"
CONFIG_FILE = "config.json"
HISTORY_DIR = "history"
ASSETS_DIR = "assets"

class StorageManager:
    def __init__(self):
        self.root_dir = Path(STORAGE_DIR)
        self.config_path = self.root_dir / CONFIG_FILE
        self.history_dir = self.root_dir / HISTORY_DIR
        self.assets_dir = self.root_dir / ASSETS_DIR
        
        # Ensure directories exist
        self.root_dir.mkdir(exist_ok=True)
        self.history_dir.mkdir(exist_ok=True)
        self.assets_dir.mkdir(exist_ok=True)
        
    def save_config(self, config_data: Dict[str, Any]):
        """Save UI configuration to file"""
        try:
            with open(self.config_path, 'w', encoding='utf-8') as f:
                json.dump(config_data, f, indent=2, ensure_ascii=False)
        except Exception as e:
            print(f"Error saving config: {e}")

    def load_config(self) -> Dict[str, Any]:
        """Load UI configuration from file"""
        if not self.config_path.exists():
            return {}
        try:
            with open(self.config_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            print(f"Error loading config: {e}")
            return {}

    def save_asset(self, uploaded_file) -> str:
        """Save an uploaded file (e.g., background image) into assets directory.

        Args:
            uploaded_file: a file-like object (Streamlit UploadedFile) or bytes-like

        Returns:
            The saved file path as string, or None on failure.
        """
        try:
            # Determine filename
            if hasattr(uploaded_file, 'name'):
                filename = uploaded_file.name
            else:
                filename = f"asset_{int(time.time())}"

            # sanitize
            safe_name = "".join([c for c in filename if c.isalnum() or c in (' ', '.', '_', '-')]).strip().replace(' ', '_')
            dest = self.assets_dir / f"{int(time.time())}_{safe_name}"

            # Write bytes
            with open(dest, 'wb') as out:
                # Streamlit UploadedFile has getbuffer()
                if hasattr(uploaded_file, 'getbuffer'):
                    out.write(uploaded_file.getbuffer())
                else:
                    # try reading
                    data = uploaded_file.read()
                    if isinstance(data, str):
                        data = data.encode('utf-8')
                    out.write(data)

            return str(dest)
        except Exception as e:
            print(f"Error saving asset: {e}")
            return None

    def save_history(self, session_type: str, topic: str, content: str, metadata: Dict[str, Any] = None):
        """
        Save a session report/history
        
        Args:
            session_type: 'council' or 'debate'
            topic: The main topic
            content: The full markdown report or content
            metadata: Additional info (model used, date, etc)
        """
        timestamp = int(time.time())
        date_str = time.strftime("%Y-%m-%d %H:%M:%S")
        
        # Create a safe filename
        safe_topic = "".join([c for c in topic[:20] if c.isalnum() or c in (' ', '_', '-')]).strip().replace(' ', '_')
        filename = f"{timestamp}_{session_type}_{safe_topic}.json"
        
        data = {
            "id": str(timestamp),
            "timestamp": timestamp,
            "date": date_str,
            "type": session_type,
            "topic": topic,
            "content": content,
            "metadata": metadata or {}
        }
        
        try:
            with open(self.history_dir / filename, 'w', encoding='utf-8') as f:
                json.dump(data, f, indent=2, ensure_ascii=False)
            return True
        except Exception as e:
            print(f"Error saving history: {e}")
            return False

    def list_history(self) -> List[Dict[str, Any]]:
        """List all history items (metadata only)"""
        items = []
        if not self.history_dir.exists():
            return []
            
        for file in self.history_dir.glob("*.json"):
            try:
                with open(file, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    # Return summary info
                    items.append({
                        "id": data.get("id"),
                        "date": data.get("date"),
                        "type": data.get("type"),
                        "topic": data.get("topic"),
                        "filename": file.name
                    })
            except Exception:
                continue
                
        # Sort by timestamp desc
        return sorted(items, key=lambda x: x.get("date", ""), reverse=True)

    def load_history_item(self, filename: str) -> Dict[str, Any]:
        """Load full content of a history item"""
        path = self.history_dir / filename
        if not path.exists():
            return None
        try:
            with open(path, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception:
            return None

    # ==================== Session Cache (Resume Functionality) ====================
    def save_session_state(self, key: str, data: Dict[str, Any]):
        """Save temporary session state for recovery"""
        try:
            # We use a dedicated cache file per key
            cache_file = self.root_dir / f"{key}_cache.json"
            data["_timestamp"] = int(time.time())
            with open(cache_file, 'w', encoding='utf-8') as f:
                json.dump(data, f, indent=2, ensure_ascii=False)
        except Exception as e:
            print(f"Error saving session cache: {e}")

    def load_session_state(self, key: str) -> Dict[str, Any]:
        """Load temporary session state"""
        cache_file = self.root_dir / f"{key}_cache.json"
        if not cache_file.exists():
            return None
        try:
            with open(cache_file, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception:
            return None

    def clear_session_state(self, key: str):
        """Clear temporary session state"""
        cache_file = self.root_dir / f"{key}_cache.json"
        if cache_file.exists():
            try:
                os.remove(cache_file)
            except Exception:
                pass