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

def img2txt(url: str) -> str:
    captioning_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
    text = captioning_model(url, max_new_tokens=20)[0]["generated_text"]
    return text

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": "user", "content": story_prompt}],
        top_k=top_k,
        top_p=top_p,
        temperature=temperature,
        stream=True
    )
    story = ''.join(chunk.choices[0].delta.content for chunk in stream)
    return story

def txt2speech(text: str) -> None:
    tts = gTTS(text=text, lang='en')
    tts.save("audio_story.mp3")

def send_story_email(recipient_email: str, story_text: str, audio_file_path: str) -> bool:
    try:
        smtp_server = os.environ.get("SMTP_SERVER")
        smtp_port = int(os.environ.get("SMTP_PORT", 587))
        sender_email = os.environ.get("SENDER_EMAIL")
        sender_password = os.environ.get("SENDER_PASSWORD")

        msg = MIMEMultipart()
        msg['From'] = sender_email
        msg['To'] = recipient_email
        msg['Subject'] = "Your Generated Story"
        
        msg.attach(MIMEText(f"Here's your generated story:\n\n{story_text}\n\nEnjoy!", 'plain'))

        with open(audio_file_path, 'rb') as audio_file:
            audio_part = MIMEAudio(audio_file.read(), _subtype='mp3')
            audio_part.add_header('Content-Disposition', 'attachment', filename=os.path.basename(audio_file_path))
            msg.attach(audio_part)

        with smtplib.SMTP(smtp_server, smtp_port) as server:
            server.starttls()
            server.login(sender_email, sender_password)
            server.send_message(msg)
        
        return True
        
    except Exception as e:
        print(f"Error sending email: {str(e)}")
        return False

def validate_email(email: str) -> bool:
    pattern = r'^[\w\.-]+@[\w\.-]+\.\w+$'
    return re.match(pattern, email) is not None

def main():
    st.set_page_config(page_title="🎨 Image-to-Audio Story 🎧", layout="wide")
    
    st.title("Turn the Image into Audio Story")

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

    if uploaded_file is not None:
        st.image(uploaded_file)
        
        if st.button("🎨 Generate Story"):
            scenario = img2txt(uploaded_file)  # Process the uploaded image
            prompt = f"Create a story based on: {scenario}"
            story = txt2story(prompt, top_k=5, top_p=0.8, temperature=1.5)
            txt2speech(story)

            st.session_state.story = story
            
            st.audio("audio_story.mp3")  # Use Streamlit's audio player
            
            email = st.text_input("Enter your email address:")
            if st.button("📤 Send to Email"):
                if validate_email(email):
                    if send_story_email(email, story, "audio_story.mp3"):
                        st.success(f"Email sent to: {email}")
                    else:
                        st.error("❌ Failed to send email.")
                else:
                    st.error("Please enter a valid email address.")

if __name__ == '__main__':
    main()