# import part
import streamlit as st
from transformers import pipeline

# 儿童友好主题配置
def set_child_theme():
    st.markdown("""
    <style>
        .stApp {
            background-color: #F9F3E6;
        }
        .stButton>button {
            background-color: #FFB6C1 !important;
            color: white !important;
            border-radius: 20px !important;
            padding: 10px 24px !important;
            font-size: 20px !important;
            font-family: 'Comic Sans MS';
        }
        h1, h2, h3 {
            color: #6A4C93 !important;
            font-family: 'Comic Sans MS';
        }
        .st-emotion-cache-1kyxreq {
            justify-content: center;
        }
    </style>
    """, unsafe_allow_html=True)

# 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

# text2story（儿童内容优化版）
def text2story(scenario):
    child_prompt = f"Create a children's story suitable for 3-10 years old about {scenario}. " \
                   "Include talking animals and magical elements. " \
                   "Use simple words and happy ending. " \
                   "Story structure: Beginning -> Problem -> Solution -> Happy ending.\n\n"
    
    pipe = pipeline("text-generation", 
                    model="pranavpsv/genre-story-generator-v2",
                    max_new_tokens=200,
                    truncation=True,
                    temperature=0.7,  # 降低随机性
                    top_p=0.9,
                    repetition_penalty=1.2)
    
    full_story = pipe(child_prompt + scenario,
                     do_sample=True,
                     num_return_sequences=1)[0]['generated_text']
    
    # 过滤不合适内容
    safe_story = full_story.replace("violence", "").replace("scary", "").replace("died", "solved")
    
    # 确保以快乐结局结束
    happy_endings = ["happily ever after", "big smile", "learned a good lesson"]
    if not any(ending in safe_story.lower() for ending in happy_endings):
        safe_story += " And they all lived happily ever after!"
    
    # 添加表情符号装饰
    return "🌈 " + safe_story[:safe_story.rfind(".")+1] + " 🎉"

# text2audio（儿童语音优化）
def text2audio(story_text):
    pipe = pipeline("text-to-audio", 
                   model="Matthijs/mms-tts-eng",
                   vocoder="edresson/wavegrad-ljspeech")
    audio_data = pipe(story_text,
                    # 调整语音参数
                    speaker_embeddings={"pitch_scale": 1.5, "energy_scale": 1.2})
    return audio_data

def main():
    set_child_theme()
    st.set_page_config(page_title="Magic Story Maker", page_icon="🧚")
    
    st.title("🧚✨ Magic Picture Story Maker ✨🧚")
    
    with st.expander("👀 How to use?"):
        st.write("""
        1. Upload any picture (pets, toys, drawings)
        2. Watch magic describe the picture
        3. Get a magical story
        4. Listen to the fairy tale!
        """)
    
    uploaded_file = st.file_uploader("📷 Take a photo of your favorite thing...", type=["jpg", "png"])

    if uploaded_file is not None:
        # 显示上传的图片（圆形裁剪）
        st.markdown(f"""
        <div style="display: flex; justify-content: center;">
            <img src="data:image/png;base64,{uploaded_file.getvalue().encode("base64").decode()}" 
                 style="border-radius: 50%; width: 200px; height: 200px; object-fit: cover;">
        </div>
        """, unsafe_allow_html=True)

        # Processing stages
        stages = {
            "🔮 Reading Magic Picture...": img2text,
            "📖 Writing Fairy Tale...": text2story,
            "🎧 Preparing Story Song...": text2audio
        }

        results = {}
        current_stage = None
        try:
            for stage_name, stage_func in stages.items():
                current_stage = stage_name
                with st.spinner(f"{stage_name}"):
                    if stage_name == "🔮 Reading Magic Picture...":
                        results["scenario"] = stage_func(uploaded_file.name)
                    elif stage_name == "📖 Writing Fairy Tale...":
                        results["story"] = stage_func(results["scenario"])
                    else:
                        results["audio"] = stage_func(results["story"])
                
                # 添加阶段完成动画
                st.success(f"✅ {stage_name.replace('...','')} Complete! ✨")

        except Exception as e:
            st.error(f"Oops! Magic wand broke at {current_stage} 🪄💥 Please try again!")
            return

        # 结果显示
        st.balloons()
        with st.container():
            st.subheader("🧚 Your Magic Story")
            st.markdown(f"""
            <div style="
                background: white;
                padding: 20px;
                border-radius: 15px;
                box-shadow: 0 4px 6px rgba(0,0,0,0.1);
                font-family: 'Comic Sans MS';
                font-size: 18px;
                color: #4B4453;
            ">
            {results['story']}
            </div>
            """, unsafe_allow_html=True)
            
            # 自动播放音频
            st.subheader("🎶 Story Song")
            st.audio(results['audio']['audio'], 
                    format="audio/wav",
                    sample_rate=results['audio']['sampling_rate'])
            
            # 下载按钮
            st.download_button("📥 Save Fairy Tale", 
                             results['story'], 
                             file_name="magic_story.txt",
                             mime="text/plain")

if __name__ == "__main__":
    main()