import os
import re
import smtplib
import streamlit as st
from transformers import pipeline
from typing import Dict
from gtts import gTTS
from together import Together
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from email.mime.audio import MIMEAudio
#from dotenv import load_dotenv

# Load environment variables
#load_dotenv()

# Image-to-text function
def img2txt(url: str) -> str:
    print("Initializing captioning model...")
    captioning_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
    
    print("Generating text from the image...")
    text = captioning_model(url, max_new_tokens=20)[0]["generated_text"]
    
    print(text)
    return text

# Text-to-story generation function
def txt2story(prompt: str, top_k: int, top_p: float, temperature: float) -> str:
    client = Together(api_key=os.environ.get("TOGETHER_API_KEY"))

    story_prompt = f"Write a short story of no more than 250 words based on the following prompt: {prompt}"

    stream = client.chat.completions.create(
        model="meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo",
        messages=[
            {"role": "system", "content": '''As an experienced short story writer, write a meaningful story influenced by the provided prompt.
            Ensure the story does not exceed 250 words.'''},
            {"role": "user", "content": story_prompt}
        ],
        top_k=top_k,
        top_p=top_p,
        temperature=temperature,
        stream=True
    )

    story = ''
    for chunk in stream:
        story += chunk.choices[0].delta.content

    return story

# Text-to-speech function
def txt2speech(text: str) -> None:
    print("Converting text to speech using gTTS...")
    tts = gTTS(text=text, lang='en')
    tts.save("audio_story.mp3")

# Get user preferences function
def get_user_preferences() -> Dict[str, str]:
    preferences = {
        'continent': st.selectbox("Continent", ["North America", "Europe", "Asia", "Africa", "Australia"]),
        'genre': st.selectbox("Genre", ["Science Fiction", "Fantasy", "Mystery", "Romance"]),
        'setting': st.selectbox("Setting", ["Future", "Medieval times", "Modern day", "Alternate reality"]),
        'plot': st.selectbox("Plot", ["Hero's journey", "Solving a mystery", "Love story", "Survival"]),
        'tone': st.selectbox("Tone", ["Serious", "Light-hearted", "Humorous", "Dark"]),
        'theme': st.selectbox("Theme", ["Self-discovery", "Redemption", "Love", "Justice"]),
        'conflict': st.selectbox("Conflict Type", ["Person vs. Society", "Internal struggle", "Person vs. Nature", "Person vs. Person"]),
        'twist': st.selectbox("Mystery/Twist", ["Plot twist", "Hidden identity", "Unexpected ally/enemy", "Time paradox"]),
        'ending': st.selectbox("Ending", ["Happy", "Bittersweet", "Open-ended", "Tragic"])
    }
    return preferences


# Main Streamlit application
def main():
    st.set_page_config(
        page_title="🎨 Image-to-Audio Story 🎧",
        page_icon="🖼️",
        layout="wide"
    )
    st.title("Turn the Image into Audio Story")

    # Initialize session state variables
    if "story" not in st.session_state:
        st.session_state.story = ""
    if "audio_file_path" not in st.session_state:
        st.session_state.audio_file_path = ""
    if "caption" not in st.session_state:
        st.session_state.caption = ""

    # Main content area
    col1, col2 = st.columns([2, 3])

    with col1:
        # Image upload section
        st.markdown("## 📷 Upload Image")
        uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

        # Story preferences section
        st.markdown("## 🎭 Story Preferences")
        preferences = get_user_preferences()

    with col2:
        if uploaded_file is not None:
            # Display uploaded image
            st.markdown("## 🖼️ Your Image")
            bytes_data = uploaded_file.read()
            with open("uploaded_image.jpg", "wb") as file:
                file.write(bytes_data)
            st.image(uploaded_file, use_column_width=True)

            # Process image and generate story
            if st.button("🎨 Generate Story"):
                with st.spinner("🤖 AI is working its magic..."):
                    try:
                        # Get image description
                        scenario = img2txt("uploaded_image.jpg")
                        st.session_state.caption = scenario  # Store caption in session state
                        
                        # Create story prompt
                        prompt = f"""Based on the image description: '{scenario}', 
                        create a {preferences['genre']} story set in {preferences['setting']} 
                        in {preferences['continent']}. The story should have a {preferences['tone']} 
                        tone and explore the theme of {preferences['theme']}. The main conflict 
                        should be {preferences['conflict']}. The story should have a {preferences['twist']} 
                        and end with a {preferences['ending']} ending."""
                        
                        # Generate story
                        story = txt2story(prompt, top_k=5, top_p=0.8, temperature=1.5)
                        st.session_state.story = story  # Store story in session state
                        
                        # Convert to audio
                        txt2speech(story)
                        st.session_state.audio_file_path = "audio_story.mp3"  # Store audio path in session state

                    except Exception as e:
                        st.error(f"An error occurred: {str(e)}")
                        st.warning("Please try again or contact support if the problem persists.")

        # Display results if story exists in session state
        if st.session_state.story:
            st.markdown("---")
            
            # Image caption
            with st.expander("📜 Image Caption", expanded=True):
                st.write(st.session_state.caption)
            
            # Story text
            with st.expander("📖 Generated Story", expanded=True):
                st.write(st.session_state.story)
            
            # Audio player
            with st.expander("🎧 Audio Version", expanded=True):
                st.audio(st.session_state.audio_file_path)


if __name__ == '__main__':
    main()