app.py

# app.py

import streamlit as st
from services.claude_service import ClaudeService
from services.chart_analysis import ChartAnalysisService
from ui.components import (
    create_sidebar,
    show_analysis_section,
    show_chat_history,
    show_follow_up_section,
    show_save_options,
    create_expertise_selector
)
from utils.file_handlers import save_chat_history
from utils.learning_module import LearningModule
from auth.auth_manager import AuthManager
from storage.storage_manager import UserStorageManager
from pages.login import show_login_page, show_logout_button
from pages.previous_chats import show_previous_chats_tab

def init_session_state():
    if 'chat_history' not in st.session_state:
        st.session_state.chat_history = []
    if 'current_image' not in st.session_state:
        st.session_state.current_image = None
    if 'current_analysis' not in st.session_state:
        st.session_state.current_analysis = None
    if 'current_images' not in st.session_state:
        st.session_state.current_images = []
    if 'analysis_results' not in st.session_state:
        st.session_state.analysis_results = []
    if 'followups' not in st.session_state:
        st.session_state.followups = []
    if 'conversation_context' not in st.session_state:
        st.session_state.conversation_context = []
        
def reset_session_state():
    """Reset session state for a new chat"""
    st.session_state.chat_history = []
    st.session_state.current_image = None
    st.session_state.current_analysis = None
    st.session_state.current_images = []
    st.session_state.analysis_results = []
    st.session_state.conversation_context = []

def show_chart_analysis_tab(claude_service, analysis_service, storage_manager):
    """Display chart analysis functionality"""
    # Get user expertise level
    expertise_level = create_expertise_selector()
    
    # Create sidebar and get inputs
    result = create_sidebar()
    
    # Check if we need to start a new chat
    if result[0] == "new_chat":
        # Save current chat if there's content
        if st.session_state.chat_history:
            storage_manager.save_chat(
                st.session_state.chat_history,
                st.session_state.current_images,
                st.session_state.get('last_saved_chat'),
                claude_service
            )
            st.success("Previous chat saved! Starting new chat...")
        
        # Reset session state
        reset_session_state()
        st.rerun()
    
    # Unpack the sidebar results
    upload_option, uploaded_files, patterns, indicators, comparison_type = result
    
    # Main content area
    col1, col2 = st.columns([2, 1])
    
    with col1:
        st.title("📈 Stock Chart Analysis Assistant")
        
        if upload_option == "Ask Question":
            question = st.text_input(
                "What would you like to know?",
                key="main_question_input"
            )
        
        # Main analysis section
        if uploaded_files:
            # Store all uploaded images
            st.session_state.current_images = [file.getvalue() for file in uploaded_files]
            
        analyze_clicked = show_analysis_section(uploaded_files)
            
        if analyze_clicked:
            if not uploaded_files:
                st.warning("Please upload at least one chart.")
            else:
                with st.spinner("Analyzing charts..."):
                    results = analysis_service.analyze_multiple_charts(
                        st.session_state.current_images,
                        patterns,
                        indicators,
                        comparison_type,
                        expertise_level
                    )
                    
                    if results:
                        st.session_state.current_analysis = results[-1]['analysis']
                        st.session_state.chat_history.extend(results)
                        st.session_state.analysis_results = results

        # Display analysis results from session state
        if st.session_state.analysis_results:
            results = st.session_state.analysis_results
            
            # First show comparative analysis if it exists
            for idx, result in enumerate(results):
                if result.get('analysis_type') != 'Individual':
                    with st.container():
                        st.subheader(f"{result['analysis_type']} Results")
                        st.markdown(result['analysis'])
                        
                        # Add enhanced follow-up section with context
                        follow_up = show_follow_up_section(
                            key_suffix=f"comparative_{idx}",
                            previous_response=result['analysis']
                        )
                        
                        if follow_up:
                            with st.spinner("Processing follow-up..."):
                                response = analysis_service.handle_follow_up_question(
                                    follow_up,
                                    result['analysis'],
                                    st.session_state.current_images
                                )
                                if response:
                                    st.session_state.chat_history.append(response)
                                    
                                    # Show response in the chat interface
                                    st.markdown(response['analysis'])
                                    
                                    # Save to storage with smart naming
                                    storage_manager.save_chat(
                                        st.session_state.chat_history,
                                        st.session_state.current_images,
                                        None,
                                        claude_service
                                    )
            
            # Then show individual analyses
            individual_analyses = [r for r in results if r.get('analysis_type') == 'Individual']
            for idx, result in enumerate(individual_analyses):
                with st.container():
                    st.subheader(f"Analysis Results - Chart {idx + 1}")
                    st.markdown(result['analysis'])
                    
                    # Add enhanced follow-up section for each analysis
                    follow_up = show_follow_up_section(
                        key_suffix=f"individual_{idx}",
                        previous_response=result['analysis']
                    )
                    
                    if follow_up:
                        with st.spinner("Processing follow-up..."):
                            response = analysis_service.handle_follow_up_question(
                                follow_up,
                                result['analysis'],
                                st.session_state.current_images[idx] if idx < len(st.session_state.current_images) else None
                            )
                            if response:
                                st.session_state.chat_history.append(response)
                                
                                # Show response in the chat interface
                                st.markdown(response['analysis'])
                                
                                # Save to storage with smart naming
                                storage_manager.save_chat(
                                    st.session_state.chat_history,
                                    st.session_state.current_images,
                                    None,
                                    claude_service
                                )
            
            # Risk warning at the bottom
            st.warning(
                "⚠️ This analysis is AI-generated and for informational purposes only. "
                "Do not make trading decisions solely based on this information."
            )
    
    with col2:
        # Add a collapsible history section
        with st.expander("💬 View Chat History", expanded=False):
            if st.session_state.chat_history:
                # Show only last 3 interactions by default
                show_full = st.checkbox("Show full history", value=False)
                
                if show_full:
                    history_to_show = st.session_state.chat_history
                else:
                    history_to_show = st.session_state.chat_history[-3:]
                    if len(st.session_state.chat_history) > 3:
                        st.info(f"Showing last 3 of {len(st.session_state.chat_history)} interactions")
                
                # Display the selected history
                show_chat_history(history_to_show)
            else:
                st.info("No chat history yet")
        
        # Add save options in a separate expander
        with st.expander("💾 Save Analysis", expanded=False):
            save_name = show_save_options()
            if save_name and st.session_state.chat_history:
                filename = storage_manager.save_chat(
                    st.session_state.chat_history,
                    st.session_state.current_images[0] if st.session_state.current_images else None,
                    f"{save_name}.json" if save_name else None,
                    claude_service
                )
                if filename:
                    st.success(f"Chat saved as {filename}")
                else:
                    st.info("Chat saved to session state")
def show_learning_tab(learning_module):
    """Display learning center functionality"""
    st.title("📚 Trading Learning Center")
    
    # Create tabs for different learning options
    learn_tab1, learn_tab2 = st.tabs(["Structured Courses", "Custom Learning"])
    
    with learn_tab1:
        st.header("Learn Trading Step by Step")
        learning_module.display_course_selection()
        
    with learn_tab2:
        st.header("Ask Any Trading Question")
        learning_module.display_custom_learning()

def main():
    st.set_page_config(
        page_title="Stock Chart Assistant",
        layout="wide",
        initial_sidebar_state="expanded"
    )
    
    try:
        # Initialize authentication
        auth_manager = AuthManager()
        
        # Show login page if not authenticated
        if not auth_manager.is_authenticated():
            show_login_page(auth_manager)
            return
        
        # Get user's storage paths and initialize storage
        storage_paths = auth_manager.get_user_storage_paths()
        storage_manager = UserStorageManager(storage_paths)
        
        # Initialize services
        claude_service = ClaudeService()
        analysis_service = ChartAnalysisService(claude_service)
        learning_module = LearningModule(claude_service)
        
        # Initialize session state
        init_session_state()
        
        # Show logout button in sidebar
        show_logout_button(auth_manager)
        
        # Load previous context if available
        if 'chat_history' not in st.session_state:
            context = storage_manager.get_context()
            if context:
                st.session_state.chat_history = context.get('chat_history', [])
                st.session_state.current_analysis = context.get('current_analysis')
        
        # Main application tabs
        tab1, tab2, tab3 = st.tabs(["Chart Analysis", "Learning Center", "Previous Chats"])
        
        with tab1:
            show_chart_analysis_tab(claude_service, analysis_service, storage_manager)
            
        with tab2:
            show_learning_tab(learning_module)
            
        with tab3:
            show_previous_chats_tab(storage_manager)
        
        # Save context before closing
        storage_manager.save_context({
            'chat_history': st.session_state.chat_history,
            'current_analysis': st.session_state.current_analysis
        })
        
    except Exception as e:
        st.error(f"An error occurred: {str(e)}")
        st.warning("Please try refreshing the page. If the error persists, contact support.")

if __name__ == "__main__":
    main()