import gradio as gr
import cv2
import numpy as np
import subprocess
import os
import tempfile

# === Step 1: Stack videos vertically ===
def stack_videos(main_path, meme_path, output_path):
    cap_main = cv2.VideoCapture(main_path)
    cap_meme = cv2.VideoCapture(meme_path)

    width, height = 1080, 1920
    main_height = int(height * 0.7)
    meme_height = height - main_height

    fps = int(cap_main.get(cv2.CAP_PROP_FPS)) or 30
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

    while True:
        ret_main, frame_main = cap_main.read()
        ret_meme, frame_meme = cap_meme.read()

        if not ret_main:
            break
        if not ret_meme:
            cap_meme.set(cv2.CAP_PROP_POS_FRAMES, 0)
            ret_meme, frame_meme = cap_meme.read()

        frame_main = cv2.resize(frame_main, (width, main_height))
        frame_meme = cv2.resize(frame_meme, (width, meme_height))

        stacked = np.vstack((frame_main, frame_meme))
        out.write(stacked)

    cap_main.release()
    cap_meme.release()
    out.release()
    return output_path

# === Step 2: Mix audio ===
def mix_audio(main_path, meme_path, output_path):
    cmd = [
        'ffmpeg',
        '-i', main_path,
        '-i', meme_path,
        '-filter_complex',
        '[0:a]volume=0.9[a0];[1:a]volume=0.1[a1];[a0][a1]amix=inputs=2:duration=first[aout]',
        '-map', '[aout]',
        '-y',
        output_path
    ]
    subprocess.run(cmd, check=True)
    return output_path

# === Step 3: Combine stacked video + mixed audio ===
def combine_video_audio(video_path, audio_path, output_path):
    cmd = [
        'ffmpeg',
        '-i', video_path,
        '-i', audio_path,
        '-c:v', 'copy',
        '-c:a', 'aac',
        '-shortest',
        '-y',
        output_path
    ]
    subprocess.run(cmd, check=True)
    return output_path

# === Gradio interface function ===
def create_reel(main_video, meme_video):
    try:
        with tempfile.TemporaryDirectory() as tmpdir:
            main_path = os.path.join(tmpdir, "main.mp4")
            meme_path = os.path.join(tmpdir, "meme.mp4")

            # Save uploaded files
            with open(main_path, "wb") as f:
                f.write(main_video.read())
            with open(meme_path, "wb") as f:
                f.write(meme_video.read())

            stacked_video_path = os.path.join(tmpdir, "stacked.mp4")
            mixed_audio_path = os.path.join(tmpdir, "audio.mp3")
            final_output_path = os.path.join(tmpdir, "final_reel.mp4")

            stack_videos(main_path, meme_path, stacked_video_path)
            mix_audio(main_path, meme_path, mixed_audio_path)
            combine_video_audio(stacked_video_path, mixed_audio_path, final_output_path)

            return final_output_path

    except subprocess.CalledProcessError as e:
        return f"FFmpeg error: {e}"
    except Exception as e:
        return f"Error: {e}"

# === Gradio UI ===
gr.Interface(
    fn=create_reel,
    inputs=[
        gr.Video(label="Main Video (Top 70%)"),
        gr.Video(label="Meme Video (Bottom 30%)")
    ],
    outputs=gr.Video(label="Final Reel (Stacked + Mixed Audio)"),
    title="📱 Reel Generator",
    description="Stack your main video (top 70%) with a meme video (bottom 30%). Audio: 90% main, 10% meme. Perfect for Reels/Shorts."
).launch()