storage/storage_manager.py

# storage/storage_manager.py

import json
from pathlib import Path
from datetime import datetime
import shutil
import streamlit as st

class UserStorageManager:
    def __init__(self, storage_paths):
        self.paths = storage_paths
        self.ensure_directories()

    def ensure_directories(self):
        """Ensure all required directories exist"""
        try:
            for path in self.paths.values():
                path.mkdir(parents=True, exist_ok=True)
        except Exception as e:
            st.error(f"Error creating directories: {str(e)}")

    def save_chat(self, chat_data, images=None, filename=None, claude_service=None):
        """Save chat with associated images and smart naming"""
        try:
            timestamp = datetime.now().isoformat()
            formatted_timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            
            # Generate smart name if available, with fallback handling
            chat_title = None
            try:
                if claude_service and isinstance(chat_data, list):
                    from utils.smart_naming import generate_chat_name
                    chat_title = generate_chat_name(chat_data, claude_service)
            except Exception as e:
                st.warning(f"Could not generate smart name: {str(e)}. Using default naming.")
            
            chat_id = filename or f"chat_{formatted_timestamp}"
            if filename:
                chat_id = filename.replace('.json', '')
            
            # Save chat data
            chat_file = self.paths["chats"] / f"{chat_id}.json"
            
            # Add metadata
            chat_metadata = {
                'id': chat_id,
                'timestamp': timestamp,
                'formatted_timestamp': formatted_timestamp,
                'title': chat_title or f"Trading Analysis {formatted_timestamp}",
                'data': chat_data
            }
            
            # Save chat metadata
            with open(chat_file, "w") as f:
                json.dump(chat_metadata, f, indent=4)
            
            # Save associated images
            if images:
                image_dir = self.paths["images"] / chat_id
                image_dir.mkdir(exist_ok=True)
                
                if isinstance(images, list):
                    for idx, img_data in enumerate(images):
                        img_file = image_dir / f"image_{idx}.png"
                        with open(img_file, "wb") as f:
                            f.write(img_data)
                else:
                    img_file = image_dir / "image_0.png"
                    with open(img_file, "wb") as f:
                        f.write(images)
                
                # Update chat metadata with image references
                chat_metadata['image_paths'] = [str(p) for p in image_dir.glob("*.png")]
                with open(chat_file, "w") as f:
                    json.dump(chat_metadata, f, indent=4)

            return chat_id
            
        except Exception as e:
            st.error(f"Error saving chat: {str(e)}")
            return None

    def get_all_chats(self):
        """Get all user's chats with metadata"""
        try:
            chats = []
            # Sort chat files by modification time for most recent first
            chat_files = sorted(
                self.paths["chats"].glob("*.json"),
                key=lambda x: x.stat().st_mtime,
                reverse=True
            )
            for chat_file in chat_files:
                try:
                    with open(chat_file, "r") as f:
                        chat_data = json.load(f)
                        # Load associated images if they exist
                        chat_id = chat_data.get('id', chat_file.stem)
                        image_dir = self.paths["images"] / chat_id
                        if image_dir.exists():
                            images = []
                            for img_file in sorted(image_dir.glob("*.png")):
                                with open(img_file, "rb") as f:
                                    images.append(f.read())
                            chat_data['images'] = images
                        chats.append(chat_data)
                except Exception as e:
                    st.warning(f"Error loading chat {chat_file.name}: {str(e)}")
                    continue
            return chats
        except Exception as e:
            st.error(f"Error loading chats: {str(e)}")
            return []

    def get_chat(self, chat_id):
        """Get specific chat with images"""
        try:
            chat_file = self.paths["chats"] / f"{chat_id}.json"
            if not chat_file.exists():
                return None
                
            with open(chat_file, "r") as f:
                chat_data = json.load(f)
                
            # Load associated images
            image_dir = self.paths["images"] / chat_id
            if image_dir.exists():
                images = []
                for img_file in sorted(image_dir.glob("*.png")):
                    with open(img_file, "rb") as f:
                        images.append(f.read())
                chat_data['images'] = images
                
            return chat_data
        except Exception as e:
            st.error(f"Error loading chat: {str(e)}")
            return None

    def save_context(self, context_data):
        """Save conversation context"""
        try:
            context_file = self.paths["context"] / "context.json"
            with open(context_file, "w") as f:
                json.dump(context_data, f, indent=4)
        except Exception as e:
            st.error(f"Error saving context: {str(e)}")

    def get_context(self):
        """Get saved conversation context"""
        try:
            context_file = self.paths["context"] / "context.json"
            if context_file.exists():
                with open(context_file, "r") as f:
                    return json.load(f)
            return None
        except Exception as e:
            st.error(f"Error loading context: {str(e)}")
            return None

    def clear_context(self):
        """Clear conversation context"""
        try:
            context_file = self.paths["context"] / "context.json"
            if context_file.exists():
                context_file.unlink()
        except Exception as e:
            st.error(f"Error clearing context: {str(e)}")

    def delete_chat(self, chat_id):
        """Delete a chat and its associated images"""
        try:
            # Delete chat file
            chat_file = self.paths["chats"] / f"{chat_id}.json"
            if chat_file.exists():
                chat_file.unlink()
            
            # Delete associated images
            image_dir = self.paths["images"] / chat_id
            if image_dir.exists():
                shutil.rmtree(image_dir)
            
            return True
        except Exception as e:
            st.error(f"Error deleting chat: {str(e)}")
            return False

    def get_storage_stats(self):
        """Get storage usage statistics"""
        try:
            stats = {
                'total_chats': 0,
                'total_images': 0,
                'storage_used': 0  # in bytes
            }
            
            # Count chats
            stats['total_chats'] = len(list(self.paths["chats"].glob("*.json")))
            
            # Count images
            for image_dir in self.paths["images"].glob("*"):
                if image_dir.is_dir():
                    stats['total_images'] += len(list(image_dir.glob("*.png")))
            
            # Calculate storage used
            for path_type, path in self.paths.items():
                if path.exists():
                    for file_path in path.rglob("*"):
                        if file_path.is_file():
                            stats['storage_used'] += file_path.stat().st_size
            
            return stats
        except Exception as e:
            st.error(f"Error getting storage stats: {str(e)}")
            return None