import streamlit as st
from PIL import Image
import io
import traceback

try:
    from image_to_text import generate_caption
except ImportError as e:
    st.error(f"Error importing generate_caption function: {str(e)}")
    st.error("Please make sure all required libraries are installed.")
    st.stop()

def set_custom_style():
    st.markdown("""
    <style>
    .stApp {
        max-width: 800px;
        margin: 0 auto;
        padding: 1rem;
    }
    .title {
        color: #1e3a8a;
        font-size: 2.5rem;
        font-weight: bold;
        margin-bottom: 1rem;
        text-align: center;
    }
    .stButton>button {
        background-color: #3b82f6;
        color: white;
        font-weight: bold;
        padding: 0.5rem 2rem;
        border-radius: 5px;
        border: none;
        transition: background-color 0.3s ease;
    }
    .stButton>button:hover {
        background-color: #2563eb;
    }
    .caption-container {
        background-color: #f3f4f6;
        padding: 1rem;
        border-radius: 5px;
        margin-top: 1rem;
    }
    .caption-text {
        font-size: 1.2rem;
        font-style: italic;
        color: #4a5568;
    }
    .stAlert {
        background-color: #fee2e2;
        color: #991b1b;
        padding: 0.5rem;
        border-radius: 5px;
        margin-top: 1rem;
    }
    </style>
    """, unsafe_allow_html=True)

def main():
    set_custom_style()
    
    st.markdown('<h1 class="title">Image Captioning App</h1>', unsafe_allow_html=True)

    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
    
    if uploaded_file is not None:
        try:
            image = Image.open(uploaded_file)
            st.image(image, caption='Uploaded Image', use_column_width=True)
            
            if st.button('Generate Caption'):
                with st.spinner('Generating caption...'):
                    try:
                        caption = generate_caption(image)
                        st.markdown('<div class="caption-container">', unsafe_allow_html=True)
                        st.markdown("### Generated Caption:")
                        st.markdown(f'<p class="caption-text">{caption}</p>', unsafe_allow_html=True)
                        st.markdown('</div>', unsafe_allow_html=True)
                    except Exception as e:
                        st.error(f"An error occurred while generating the caption: {str(e)}")
                        st.error("Detailed error traceback:")
                        st.code(traceback.format_exc())
        except Exception as e:
            st.error(f"An error occurred while processing the image: {str(e)}")
            st.error("Detailed error traceback:")
            st.code(traceback.format_exc())

if __name__ == "__main__":
    main()