utils/file_handlers.py

# utils/file_handlers.py

import json
from datetime import datetime
from pathlib import Path
import streamlit as st
from config.settings import CHAT_HISTORIES_DIR, CHAT_IMAGES_DIR

def save_chat_history(chat_history, image_data=None, filename=None):
    """Save chat history and associated image"""
    try:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        # Get the stock name from the latest analysis
        stock_name = "Unknown"
        if chat_history:
            latest_analysis = chat_history[-1].get('analysis', '')
            if "analyzing" in latest_analysis.lower():
                words = latest_analysis.split()
                for i, word in enumerate(words):
                    if word.lower() == "analyzing":
                        stock_name = words[i + 1].strip("(),.:")
        
        # Create filename with metadata
        base_filename = filename or f"{stock_name}_{timestamp}"
        
        # Save image if provided
        image_filename = None
        if image_data:
            try:
                image_filename = f"{base_filename}.jpg"
                image_path = CHAT_IMAGES_DIR / image_filename
                with open(image_path, "wb") as f:
                    f.write(image_data)
            except Exception as e:
                st.warning(f"Could not save image: {str(e)}")
        
        # Add metadata to chat history
        chat_data = {
            'metadata': {
                'stock_name': stock_name,
                'date_created': timestamp,
                'image_file': image_filename
            },
            'conversations': chat_history
        }
        
        # Save chat history
        try:
            json_filename = f"{base_filename}.json"
            filepath = CHAT_HISTORIES_DIR / json_filename
            with open(filepath, "w") as f:
                json.dump(chat_data, f)
            return json_filename
        except Exception as e:
            st.warning(f"Could not save chat history: {str(e)}")
            # Save to session state as fallback
            if 'saved_chats' not in st.session_state:
                st.session_state.saved_chats = []
            st.session_state.saved_chats.append(chat_data)
            return None
            
    except Exception as e:
        st.error(f"Error in save operation: {str(e)}")
        return None

def load_chat_history(filename):
    """Load chat history and associated image"""
    filepath = CHAT_HISTORIES_DIR / filename
    
    try:
        with open(filepath, "r") as f:
            chat_data = json.load(f)
        
        # Load associated image if it exists
        image_data = None
        if chat_data.get('metadata', {}).get('image_file'):
            image_path = CHAT_IMAGES_DIR / chat_data['metadata']['image_file']
            if image_path.exists():
                with open(image_path, "rb") as f:
                    image_data = f.read()
        
        return chat_data, image_data
    except Exception as e:
        st.error(f"Error loading chat history: {str(e)}")
        return None, None