import streamlit as st
from transformers import pipeline
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
from transformers import AutoProcessor, AutoModel
from gtts import gTTS
import os
import io


# Initialize session state for storing data
if 'scenario' not in st.session_state:
    st.session_state.scenario = None
if 'scenario_zh' not in st.session_state:
    st.session_state.scenario_zh = None
if 'story' not in st.session_state:
    st.session_state.story = None
if 'story_zh' not in st.session_state:
    st.session_state.story_zh = None
if 'audio_generated' not in st.session_state:
    st.session_state.audio_generated = False
if 'audio_data' not in st.session_state:
    st.session_state.audio_data = None


# function part
# img2text
def img2text(url):
    image_to_text_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
    text = image_to_text_model(url)[0]["generated_text"]
    return text


# Translation function EN to ZH
def translate_to_chinese(text):
    translator = pipeline("translation", model="steve-tong/opus-mt-en-zh-hk")
    translation = translator(text)[0]["translation_text"]
    return translation


# text2story
def text2story(text):
    # Initialize the text generation pipeline
    generator = pipeline('text-generation', model='gpt2')

    # Create a prompt for the story
    prompt = f"Create a short story about this scene: {text}\n\nStory:"

    # Generate the story
    story = generator(prompt,
                      max_length=100,
                      num_return_sequences=1,
                      temperature=0.7)[0]['generated_text']

    # Clean up the story by removing the prompt
    story = story.replace(prompt, "").strip()
    return story


def text2audio(text, lang='zh'):
    try:
        # Create a gTTS object with Chinese language
        tts = gTTS(text=text, lang=lang)
        # Save to BytesIO object
        audio_bytes = io.BytesIO()
        tts.write_to_fp(audio_bytes)
        audio_bytes.seek(0)  # Reset the pointer to the start

        return {
            'audio': audio_bytes,
            'sampling_rate': 24000  # gTTS default sampling rate
        }
    except Exception as e:
        st.error(f"音頻製作出左問題: {str(e)}")
        return None


# Apply custom CSS for modern, stylish kid-friendly UI
st.set_page_config(page_title="故事魔法", page_icon="✨", layout="wide")

st.markdown("""
<style>
  /* Modern, stylish kid-friendly design */
  @import url('https://fonts.googleapis.com/css2?family=Quicksand:wght@400;600;700&display=swap');
  @import url('https://fonts.googleapis.com/css2?family=Noto+Sans+HK:wght@400;500;700&display=swap');

  :root {
    --primary-color: #6C63FF;
    --secondary-color: #41B883;
    --accent-color: #FF6B6B;
    --background-light: #F7F9FC;
    --text-dark: #2E3A59;
    --shadow: 0 10px 20px rgba(0,0,0,0.08);
    --border-radius: 16px;
  }

  .stApp {
    background: linear-gradient(135deg, #F4F9FF, #EEFAFF);
    font-family: 'Noto Sans HK', sans-serif;
    color: var(--text-dark);
  }

  .main .block-container {
    max-width: 1000px;
    padding-top: 2rem;
    padding-bottom: 2rem;
  }

  /* Modern headers */
  h1, h2, h3 {
    font-family: 'Noto Sans HK', sans-serif;
    font-weight: 700;
    color: var(--primary-color);
  }

  h1 {
    font-size: 2.5rem;
    text-align: center;
    margin-bottom: 0;
  }

  h2 {
    font-size: 1.8rem;
    margin-bottom: 1rem;
  }

  h3 {
    font-size: 1.4rem;
    margin-bottom: 0.8rem;
  }

  /* Subtitle */
  .subtitle {
    text-align: center;
    color: #6B7897;
    font-size: 1.2rem;
    margin-bottom: 2rem;
  }

  /* Card containers */
  .stCard {
    background: white;
    border-radius: var(--border-radius);
    padding: 1.5rem;
    box-shadow: var(--shadow);
    margin-bottom: 1.5rem;
  }

  /* Accent borders for stages */
  .css-nahz7x, .css-ocqkz7, .css-4z1n4l {
    border-left: it 5px solid var(--primary-color) !important;
  }

  .css-1r6slb0, .css-1ubpcwi {
    border-left: 5px solid var(--secondary-color) !important;
  }

  .css-pkbazv, .css-5rimss {
    border-left: 5px solid var(--accent-color) !important;
  }

  /* Custom file uploader */
  .stFileUploader > div > div {
    background: var(--background-light);
    border: 2px dashed #D0D8E6;
    border-radius: 12px;
    padding: 20px;
    transition: all 0.3s ease;
  }

  .stFileUploader > div > div:hover {
    border-color: var(--primary-color);
  }

  /* Uploaded image styling */
  .stImage img {
    border-radius: 12px;
    box-shadow: var(--shadow);
  }

  /* Stage icons */
  .stage-icon {
    font-size: 1.6rem;
    margin-right: 10px;
    vertical-align: middle;
  }

  /* Response styling */
  .stText {
    font-size: 1.1rem;
    line-height: 1.7;
    background: var(--background-light);
    padding: 1rem;
    border-radius: 12px;
    border-left: 4px solid var(--secondary-color);
    margin: 1rem 0;
    box-shadow: 0 5px 15px rgba(0,0,0,0.05);
  }

  /* Button styling */
  .stButton > button {
    background: var(--secondary-color) !important;
    color: white !important;
    border: none !important;
    border-radius: 50px !important;
    padding: 0.6rem 1.5rem !important;
    font-size: 1.1rem !important;
    font-weight: 600 !important;
    font-family: 'Noto Sans HK', sans-serif !important;
    transition: all 0.3s ease !important;
    box-shadow: 0 5px 15px rgba(65, 184, 131, 0.3) !important;
  }

  .stButton > button:hover {
    background: #37A574 !important;
    transform: translateY(-3px) !important;
    box-shadow: 0 8px 20px rgba(65, 184, 131, 0.4) !important;
  }

  .stButton > button:active {
    transform: translateY(0) !important;
  }

  /* Audio player styling */
  audio {
    width: 100%;
    border-radius: 50px;
    height: 40px;
  }

  /* Emoji animation */
  @keyframes bounce {
    0%, 100% { transform: translateY(0); }
    50% { transform: translateY(-15px); }
  }

  .emoji {
    font-size: 1.8rem;
    display: inline-block;
    animation: bounce 2s infinite;
    margin: 0 8px;
  }

  .emoji:nth-child(2) {
    animation-delay: 0.2s;
  }

  .emoji:nth-child(3) {
    animation-delay: 0.4s;
  }

  .emoji:nth-child(4) {
    animation-delay: 0.6s;
  }

  /* Welcome message */
  .welcome-message {
    text-align: center;
    padding: 3rem 1.5rem;
  }

  .welcome-icon {
    font-size: 4rem;
    margin-bottom: 1rem;
  }
</style>
""", unsafe_allow_html=True)

# App header with Cantonese
st.title("✨ 故事魔法")
st.markdown("<p class='subtitle'>上載一張圖片，睇下佢點變成一個神奇嘅故事！</p>",
            unsafe_allow_html=True)

# File uploader with Cantonese
with st.container():
    st.subheader("揀一張靚相啦！")
    uploaded_file = st.file_uploader("", key="upload")

if uploaded_file is not None:
    # Save uploaded file
    bytes_data = uploaded_file.getvalue()
    temp_file_path = uploaded_file.name
    with open(temp_file_path, "wb") as file:
        file.write(bytes_data)

    # Display image
    st.image(uploaded_file, use_column_width=True)

    # Reset session state if a new file is uploaded (detect by checking if there's no scenario yet)
    if st.session_state.scenario is None:
        # Stage 1: Image to Text
        with st.container():
            st.markdown("<h3><span class='stage-icon'>🔍</span> 圖片解讀中</h3>", unsafe_allow_html=True)
            
            # Generate caption if not already done
            st.session_state.scenario = img2text(temp_file_path)
            
            # Display English caption
            st.text("英文描述: " + st.session_state.scenario)
            
            # Translate the caption to Chinese
            st.session_state.scenario_zh = translate_to_chinese(st.session_state.scenario)
            
            # Display Chinese caption
            st.text("中文描述: " + st.session_state.scenario_zh)

        # Stage 2: Text to Story
        with st.container():
            st.markdown("<h3><span class='stage-icon'>📝</span> 故事創作中</h3>", unsafe_allow_html=True)
            
            # Generate story if not already done
            st.session_state.story = text2story(st.session_state.scenario)
            
            # Display English story
            st.text("英文故事: " + st.session_state.story)
            
            # Translate the story to Chinese
            st.session_state.story_zh = translate_to_chinese(st.session_state.story)
            
            # Display Chinese story
            st.text("中文故事: " + st.session_state.story_zh)
    else:
        # Display saved results from session state
        with st.container():
            st.markdown("<h3><span class='stage-icon'>🔍</span> 圖片解讀中</h3>", unsafe_allow_html=True)
            st.text("英文描述: " + st.session_state.scenario)
            st.text("中文描述: " + st.session_state.scenario_zh)

        with st.container():
            st.markdown("<h3><span class='stage-icon'>📝</span> 故事創作中</h3>", unsafe_allow_html=True)
            st.text("英文故事: " + st.session_state.story)
            st.text("中文故事: " + st.session_state.story_zh)

    # Stage 3: Story to Audio data
    with st.container():
        st.markdown("<h3><span class='stage-icon'>🔊</span> 故事準備朗讀中</h3>", unsafe_allow_html=True)

        # Play button with Cantonese text
        if st.button("🔊 播放故事"):
            # Only generate audio if not already done
            if not st.session_state.audio_generated:
                st.session_state.audio_data = text2audio(st.session_state.story_zh, lang='zh')
                st.session_state.audio_generated = True
                
            # Play the audio
            if st.session_state.audio_data:
                st.audio(st.session_state.audio_data['audio'],
                        format="audio/wav",
                        start_time=0,
                        sample_rate=st.session_state.audio_data['sampling_rate'])
            else:
                st.error("哎呀！再試多次啦！")

    # Cleanup: Remove the temporary file when the user is done
    if os.path.exists(temp_file_path):
        os.remove(temp_file_path)
else:
    # Clear session state when no file is uploaded
    st.session_state.scenario = None
    st.session_state.scenario_zh = None
    st.session_state.story = None
    st.session_state.story_zh = None
    st.session_state.audio_generated = False
    st.session_state.audio_data = None
    
    # Welcome message in Cantonese
    st.markdown("""
    <div class="welcome-message">
        <div class="welcome-icon">✨</div>
        <h2>歡迎嚟到故事魔法！</h2>
        <p style="font-size: 1.2rem; color: #6B7897; max-width: 500px; margin: 0 auto 30px;">
            上載一張你鍾意嘅相片，我哋嘅魔法師會幫你變出一個好好玩嘅故事！
        </p>
        <div>
            <span class="emoji">🚀</span>
            <span class="emoji">🦄</span>
            <span class="emoji">🔮</span>
            <span class="emoji">🌈</span>
        </div>
    </div>
    """, unsafe_allow_html=True)