# storage/fallback_storage.py

import streamlit as st
from datetime import datetime

class FallbackStorageManager:
    """Fallback storage using Streamlit session state when persistent storage isn't available"""
    
    @staticmethod
    def init_session_storage():
        """Initialize session storage"""
        if 'fallback_storage' not in st.session_state:
            st.session_state.fallback_storage = {
                'chats': [],
                'images': {},
                'context': {}
            }

    @staticmethod
    def save_chat(chat_data, images=None):
        """Save chat to session state"""
        FallbackStorageManager.init_session_storage()
        
        chat_id = f"chat_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        
        # Store chat data
        chat_data['id'] = chat_id
        st.session_state.fallback_storage['chats'].append(chat_data)
        
        # Store images
        if images:
            st.session_state.fallback_storage['images'][chat_id] = images
            
        return chat_id

    @staticmethod
    def get_all_chats():
        """Get all chats from session state"""
        FallbackStorageManager.init_session_storage()
        return sorted(
            st.session_state.fallback_storage['chats'],
            key=lambda x: x.get('timestamp', ''),
            reverse=True
        )

    @staticmethod
    def get_chat(chat_id):
        """Get specific chat with images"""
        FallbackStorageManager.init_session_storage()
        
        # Find chat
        chat = next(
            (chat for chat in st.session_state.fallback_storage['chats']
             if chat.get('id') == chat_id),
            None
        )
        
        if chat:
            # Add images if they exist
            chat['images'] = st.session_state.fallback_storage['images'].get(chat_id, [])
            
        return chat

    @staticmethod
    def save_context(context_data):
        """Save context to session state"""
        FallbackStorageManager.init_session_storage()
        st.session_state.fallback_storage['context'] = context_data

    @staticmethod
    def get_context():
        """Get context from session state"""
        FallbackStorageManager.init_session_storage()
        return st.session_state.fallback_storage.get('context')