import streamlit as st
import requests
import time
import os
from moviepy.editor import VideoFileClip
from heygen import upload_photo, create_video, get_video_status, download_video, list_talking_photos, split_video, AUDIO_URL
from utils import *


# Streamlit App Configuration
st.set_page_config(
    page_title="Number to Video",
    page_icon="🎥",
    initial_sidebar_state="expanded",
)

# Session State Initialization
if "talking_photo_id" not in st.session_state:
    st.subheader("Photo Requirements")
    st.write("""
    - The face should be intact and clearly visible.
    - Preferably use real human faces.
    - Ensure only one face is visible in the photo.
    - Face resolution should be larger than 200x200 pixels.
    - The photo should be in JPEG or PNG format.
    """)
    st.image("sample.jpg", use_container_width=True)
    st.session_state.talking_photo_id = None

if "video_id" not in st.session_state:
    st.session_state.video_id = None

if "video_url" not in st.session_state:
    st.session_state.video_url = None

# Sidebar
st.sidebar.title("Settings")
uploaded_photo = st.sidebar.file_uploader("Upload a photo", type=["jpeg", "png"])
upload = st.sidebar.button("Upload Photo", use_container_width=True)

if uploaded_photo and upload:
    file_path = f"uploaded_photo.{uploaded_photo.type.split('/')[1]}"
    with open(file_path, "wb") as f:
        f.write(uploaded_photo.read())

    file_type = "image/jpeg" if uploaded_photo.type in ["image/jpg", "image/jpeg"] else "image/png"
    st.session_state.talking_photo_id = upload_photo(file_path, file_type)
    st.info("Photo uploaded successfully.")
    st.info("Now generate the video.")

# Video Creation
if st.session_state.talking_photo_id:
    create_video_button = st.sidebar.button("Generate Video", use_container_width=True, key="generate_video")
    if create_video_button:
        title="Number to Video"
        st.session_state.video_id = create_video(title, st.session_state.talking_photo_id, AUDIO_URL)
        # st.sidebar.write(f"Video ID: {st.session_state.video_id}")

# Video Status Check
if st.session_state.video_id:
    with st.spinner("Generating video... This may take 2-5 minutes."):
        while True:
            video_status = get_video_status(st.session_state.video_id)
            st.sidebar.info(f"Video Status: {video_status['status']}")

            if video_status["status"] == "completed":
                st.session_state.video_url = video_status["video_url"]
                st.sidebar.info("Video generation completed.")
                break
            elif video_status["status"] == "failed":
                st.sidebar.warning("Video generation failed. Please try again.")
                break

            time.sleep(20)

# Video Download and Display
if st.session_state.video_url:
    video_id = st.session_state.video_id
    video_url = st.session_state.video_url
    if not os.path.exists(f"{video_id}.mp4"):
        download_video(video_id, video_url)
        split_video(video_id)
        video_path = f"{video_id}.mp4"
    # download_video(video_id, video_url)
    # split_video(video_id)

    # Input number
    number = st.text_input("Enter a number between 1 and 99,99,999", "24")
    number = int(number)
    col1, col2 = st.columns(2)

    with col1:
        with st.container(border=True):
            # st.subheader("Photo avatar")
            # st.divider()

            # st.image("uploaded_photo.jpg", use_container_width=True)

            # Generate clip sequence based on the number
            clips = generate_clip_sequence(number)
            if clips:
                st.write("Generated Clips Sequence:")
                st.divider()
                for i, clip in enumerate(clips):
                    st.write(f"Clip {i + 1}: {clip}")
                submit = st.button("Create Video", use_container_width=True, key="create_video")

    with col2:
        with st.container(border=True):
            st.subheader("Output")
            output_file = f"{video_id}_{number}_output.mp4"
            if submit:
                combine_video(number, video_id, output_file)
                st.video(output_file, autoplay=True)