import streamlit as st
import tempfile
import os
import numpy as np
import cv2
import moviepy.editor as mp
from model import generate_stick_animation, add_camera_effects, add_background, add_sound

# Streamlit UI
st.set_page_config(page_title="AI Stick Animation", layout="wide")
st.title("🎭 AI Stick Animation Generator")

# User Inputs
story = st.text_area("📝 Enter Your Story:", "John is walking in a dark forest...")
character_name = st.text_input("🎭 Enter Main Character Name:", "John")

# Extra Features Selection
background_choice = st.selectbox("🌆 Choose Background", ["None", "Dark Forest", "Haunted House", "Battlefield"])
camera_effect = st.selectbox("📽️ Choose Camera Effect", ["None", "Shake", "Zoom", "Slow Motion"])
sound_effect = st.selectbox("🎶 Choose Background Music", ["None", "Horror", "Action", "Comedy"])

# Generate Animation Button
if st.button("🚀 Generate Animation"):
    with st.spinner("AI is generating animation..."):
        output_frames = generate_stick_animation(story, character_name)

        # Apply Extra Features
        output_frames = add_camera_effects(output_frames, camera_effect)
        output_frames = add_background(output_frames, background_choice)

        # Save animation as MP4
        temp_dir = tempfile.mkdtemp()
        output_path = os.path.join(temp_dir, "animation.mp4")
        clip = mp.ImageSequenceClip(output_frames, fps=24)  # High FPS for smooth motion
        clip.write_videofile(output_path, codec="libx264", fps=24)

        # Add Sound Effect
        final_video = add_sound(output_path, sound_effect)

        # Display Video
        st.video(final_video)
        st.download_button("📥 Download Animation", open(final_video, "rb"), file_name="stick_animation.mp4")