import os
import tempfile
import subprocess
import streamlit as st
import logging

# Clear Streamlit cache
st.cache_data.clear()
st.cache_resource.clear()

# Logging setup
logging.basicConfig(level=logging.INFO)

# Constants
CHECKPOINT_PATH = "checkpoints/wav2lip_gan.pth"
OUTPUT_VIDEO_PATH = "./results/result_voice.mp4"

# App config
st.set_page_config(page_title="Wav2Lip Lip-Sync App", layout="centered")
st.title("CodeHubb Lip-Sync")
st.write("Upload a face video and an audio file to generate a lip-synced output using the Wav2Lip model.")

# Upload video file
video_file = st.file_uploader("🎞️ Upload Video File (mp4)", type=["mp4"])

# Audio file upload
st.markdown("---")
audio_file = st.file_uploader("🎵 Upload Audio File (mp3, wav)", type=["mp3", "wav"])

# Expander for settings
with st.expander("⚙️ Advanced Settings"):
    st.markdown(f"**Model Checkpoint:** `{CHECKPOINT_PATH}`")

# Determine if button should be active
audio_ready = audio_file is not None
generate_btn = st.button("🚀 Generate Lip-Synced Video", key="generate_button", disabled=not (video_file and audio_ready))

# State management
if "processing" not in st.session_state:
    st.session_state.processing = False

# Result flag
video_generated = False

if generate_btn and not st.session_state.processing:
    st.session_state.processing = True
    with st.spinner("Processing... This may take a few moments ⏳"):
        with tempfile.TemporaryDirectory() as tempdir:
            video_path = os.path.join(tempdir, video_file.name)

            with open(video_path, "wb") as f:
                f.write(video_file.read())

            audio_path = os.path.join(tempdir, audio_file.name)
            with open(audio_path, "wb") as f:
                f.write(audio_file.read())

            # Run Wav2Lip inference
            try:
                result = subprocess.run(
                    ["python", "inference.py",
                     "--checkpoint_path", CHECKPOINT_PATH,
                     "--face", video_path,
                     "--audio", audio_path],
                    capture_output=True,
                    text=True
                )

                if result.returncode == 0:
                    st.success("✅ Processing complete!")
                    st.video(OUTPUT_VIDEO_PATH)

                    with open(OUTPUT_VIDEO_PATH, "rb") as f:
                        video_data = f.read()

                    video_generated = True

                else:
                    st.error("❌ Error during processing:")
                    st.code(result.stderr, language="bash")

            except Exception as e:
                st.error(f"An unexpected error occurred: {e}")

    st.session_state.processing = False

# Show download button only if video was successfully generated
if video_generated:
    st.download_button(
        label="⬇️ Download Lip-Synced Video",
        data=video_data,
        file_name="result_voice.mp4",
        mime="video/mp4"
    )