import streamlit as st
import os
from PIL import Image
import numpy as np
import torch
from chatbot import Chatbot
import time

# Set environment variables to force CPU usage
os.environ['CUDA_VISIBLE_DEVICES'] = ''
os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'

# Function to save uploaded file
def save_uploaded_file(uploaded_file):
    try:
        if not os.path.exists('uploads'):
            os.makedirs('uploads')
        with open(os.path.join('uploads', uploaded_file.name), 'wb') as f:
            f.write(uploaded_file.getbuffer())
        return True
    except Exception as e:
        st.error(f"Error: {e}")
        return False

# Function to show dashboard content
def show_dashboard():
    st.title("👗 Fashion Recommender System")
    st.write("Welcome to our Fashion Recommender System! Upload an image or describe what you're looking for to get personalized fashion recommendations.")
    
    # Initialize chatbot with loading state
    with st.spinner("Loading fashion assistant..."):
        try:
            chatbot = Chatbot()
            chatbot.load_data()
            st.success("✅ Fashion assistant loaded successfully!")
        except Exception as e:
            st.error(f"❌ Error initializing chatbot: {str(e)}")
            st.info("The system is running in limited mode. Some features may not be available.")
            return

    # Sidebar for uploaded image
    st.sidebar.header("📸 Image Upload")
    uploaded_file = st.sidebar.file_uploader("Choose an image", type=['jpg', 'jpeg', 'png'])
    
    if uploaded_file is not None:
        if save_uploaded_file(uploaded_file):
            display_image = Image.open(uploaded_file)
            st.sidebar.image(display_image, caption='Uploaded Image', use_column_width=True)
            
            # Process image and get recommendations
            with st.spinner("Analyzing your image and finding recommendations..."):
                try:
                    image_path = os.path.join("uploads", uploaded_file.name)
                    caption = chatbot.generate_image_caption(image_path)
                    
                    st.write("### 🖼️ Image Analysis")
                    col1, col2 = st.columns([1, 2])
                    with col1:
                        st.image(display_image, width=200)
                    with col2:
                        st.write("**Generated Caption:**")
                        st.info(caption)
                    
                    # Get recommendations based on caption
                    bot_response, recommended_products = chatbot.generate_response(caption)
                    
                    st.write("### 💫 Recommended Products")
                    if recommended_products:
                        cols = st.columns(3)
                        for i, product in enumerate(recommended_products[:3]):
                            with cols[i]:
                                product_info = chatbot.get_product_info(product['corpus_id'])
                                if product_info:
                                    st.image(
                                        product_info['image'], 
                                        caption=product_info['name'],
                                        width=150
                                    )
                                    st.write(f"**{product_info['name']}**")
                                    st.caption(f"Category: {product_info['category']}")
                                    st.caption(f"Type: {product_info['article_type']}")
                                else:
                                    st.info("Product info not available")
                    else:
                        st.warning("No products found matching your image.")
                        
                except Exception as e:
                    st.error(f"Error processing image: {str(e)}")
        else:
            st.error("Error in uploading the file.")

    # Main content area for chat
    st.write("---")
    st.write("### 💬 Chat with Fashion Assistant")
    st.write("Describe what you're looking for or ask for fashion advice!")
    
    # Initialize chat history
    if "messages" not in st.session_state:
        st.session_state.messages = []
    
    # Display chat messages from history
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])
    
    # Chat input
    if prompt := st.chat_input("What are you looking for today?"):
        # Add user message to chat history
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.markdown(prompt)
        
        # Generate response
        with st.chat_message("assistant"):
            with st.spinner("Finding the perfect fashion items..."):
                try:
                    bot_response, recommended_products = chatbot.generate_response(prompt)
                    
                    # Display bot response
                    st.markdown(bot_response)
                    
                    # Display recommended products
                    if recommended_products:
                        st.write("**🎯 Recommended for you:**")
                        
                        # Display products in columns
                        product_cols = st.columns(3)
                        for i, product in enumerate(recommended_products[:3]):
                            with product_cols[i]:
                                product_info = chatbot.get_product_info(product['corpus_id'])
                                if product_info:
                                    st.image(
                                        product_info['image'], 
                                        caption=product_info['name'],
                                        width=150
                                    )
                                    st.write(f"**{product_info['name']}**")
                                    st.caption(f"Category: {product_info['category']}")
                                    st.caption(f"Type: {product_info['article_type']}")
                                    st.caption(f"Season: {product_info['season']}")
                                else:
                                    st.info("Product info not available")
                    else:
                        st.info("No specific products found. Try describing what you're looking for in more detail!")
                        
                    # Add assistant response to chat history
                    st.session_state.messages.append({"role": "assistant", "content": bot_response})
                    
                except Exception as e:
                    error_msg = "Sorry, I encountered an error while processing your request. Please try again."
                    st.error(error_msg)
                    st.session_state.messages.append({"role": "assistant", "content": error_msg})

    # Clear chat button
    if st.button("Clear Chat History"):
        st.session_state.messages = []
        st.rerun()

# Main Streamlit app
def main():
    # Set page configuration
    st.set_page_config(
        page_title="Fashion Recommender System",
        page_icon="👗",
        layout="wide",
        initial_sidebar_state="expanded"
    )
    
    # Add custom CSS
    st.markdown("""
    <style>
    .stChatMessage {
        padding: 1rem;
        border-radius: 0.5rem;
        margin-bottom: 1rem;
    }
    .product-card {
        border: 1px solid #ddd;
        border-radius: 10px;
        padding: 1rem;
        margin: 0.5rem 0;
        background-color: #f9f9f9;
    }
    </style>
    """, unsafe_allow_html=True)

    # Show dashboard content
    show_dashboard()

# Run the main app
if __name__ == "__main__":
    main()