import streamlit as st
import os
from image_to_story_app_huggingface_model import image2text, text2story, story2audio

# Set page configuration
st.set_page_config(
    page_title="Image to Story Generator",
    page_icon="📚",
    layout="wide"
)

# Add custom CSS
st.markdown("""
    <style>
        .stApp {
            max-width: 1200px;
            margin: 0 auto;
        }
        .upload-text {
            text-align: center;
            padding: 20px;
        }
    </style>
""", unsafe_allow_html=True)

def main():
    st.title("📚 Image to Story Generator")
    st.write("Upload an image and I'll create a story about it!")

    # File uploader
    uploaded_file = st.file_uploader("Choose an image...", type=['png', 'jpg', 'jpeg'])

    if uploaded_file is not None:
        # Display the uploaded image
        st.image(uploaded_file, caption='Uploaded Image', use_container_width=True)
        
        # Save the uploaded file temporarily
        temp_path = "temp_image.jpg"
        with open(temp_path, "wb") as f:
            f.write(uploaded_file.getbuffer())

        if st.button('Generate Story'):
            try:
                with st.spinner('Generating caption from image...'):
                    caption = image2text(temp_path)
                    st.write("**Generated Caption:**", caption)

                with st.spinner('Creating story...'):
                    story = text2story(caption)
                    st.write("**Generated Story:**", story)

                with st.spinner('Creating audio narration...'):
                    output_file = "generated_story.flac"
                    audio_success = story2audio(story, output_file)
                    
                    if audio_success and os.path.exists(output_file):
                        with open(output_file, 'rb') as audio_file:
                            audio_bytes = audio_file.read()
                        if len(audio_bytes) > 0:
                            st.audio(audio_bytes, format='audio/flac')
                            st.download_button(
                                label="Download Audio",
                                data=audio_bytes,
                                file_name="story_narration.flac",
                                mime="audio/flac"
                            )
                        else:
                            st.error("Audio file is empty. Please try again.")
                    else:
                        st.error("Failed to generate audio. The story might be too long.")
            
            except Exception as e:
                st.error(f"An error occurred: {str(e)}")
            
            finally:
                # Cleanup
                if os.path.exists(temp_path):
                    os.unlink(temp_path)
                if os.path.exists(output_file):
                    os.unlink(output_file)

    st.markdown("---")
    st.markdown("Made with ❤️ using Hugging Face models")

if __name__ == "__main__":
    main()