import gradio as gr
import subprocess
import os
import spaces

# Download the file
subprocess.run([
    "wget",
    "--no-check-certificate",
    "https://drive.google.com/uc?id=1mj9lH6Be7ztYtHAr1xUUGT3hRtWJBy_5",
    "-O",
    "RIFE_trained_model_v4.13.2.zip"
], check=True)

# Unzip the downloaded file
subprocess.run([
    "unzip",
    "-o", 
    "RIFE_trained_model_v4.13.2.zip"
], check=True)

# The name of your script
SCRIPT_NAME = "inference_video.py"

@spaces.GPU(duration=120)
def run_rife(
    input_video, 
    frame_multiplier, 
    time_exponent, 
    fixed_fps, 
    video_scale, 
    enable_uhd_mode, 
    remove_duplicate_frames, 
    create_montage,
    progress=gr.Progress(track_tqdm=True),
):
    if input_video is None:
        raise gr.Error("Please upload a video first.")

    ext = "mp4"
    model_dir = "train_log"
    
    video_path_wo_ext = os.path.splitext(os.path.basename(input_video))[0]
    output_base_name = "{}_{}X_{}fps.{}".format(video_path_wo_ext, int(frame_multiplier), int(fixed_fps), ext)
    
    cmd = ["python3", SCRIPT_NAME]
    cmd.extend(["--video", input_video])
    cmd.extend(["--output", output_base_name])
    cmd.extend(["--multi", str(int(frame_multiplier))])
    
    if time_exponent != 1:
        cmd.extend(["--exp", str(int(time_exponent))])

    if fixed_fps > 0:
        cmd.extend(["--fps", str(int(fixed_fps))])
        gr.Warning("Will not merge audio because using fps flag!")
        
    if video_scale != 1.0:
        cmd.extend(["--scale", str(video_scale)])

    cmd.extend(["--ext", ext])
    cmd.extend(["--model", model_dir])
    cmd.append("--fp16") 

    if enable_uhd_mode:
        cmd.append("--UHD")
        
    if remove_duplicate_frames:
        cmd.append("--skip")
        
    if create_montage:
        cmd.append("--montage")

    print(f"Executing command: {' '.join(cmd)}")

    try:
        process = subprocess.run(cmd, capture_output=True, text=True)
        
        if process.stdout:
            print("STDOUT:", process.stdout)
        if process.stderr:
            print("STDERR:", process.stderr)

        if process.returncode != 0:
            raise gr.Error(f"Inference failed. Error: {process.stderr}")

        final_output_file = f"{output_base_name}.{ext}"
        
        if os.path.exists(final_output_file):
            return final_output_file
        else:
            if os.path.exists(output_base_name):
                return output_base_name
            raise gr.Error(f"Output file not found. Expected: {final_output_file}")

    except Exception as e:
        raise gr.Error(f"An error occurred: {str(e)}")


# --- Gradio UI Layout ---

with gr.Blocks(title="Professional FPS Booster") as app:
    gr.Markdown("# ⚡ RIFE: High-Performance FPS Booster")
    gr.Markdown("Maximize video fluidity and frame rate using deep flow estimation.")
    gr.Markdown("⚠️ **Notice:** Keep input videos under 60 seconds for frame interpolation to prevent GPU task aborts.")

    with gr.Row():
        # --- Left Column: Inputs & Settings ---
        with gr.Column(scale=1):
            input_vid = gr.Video(label="🎬 Input Source Video", sources=["upload"])
            
            with gr.Group():
                gr.Markdown("### 🚀 Performance Settings")
                
                multi_param = gr.Dropdown(
                    choices=["2", "3", "4", "5", "6"], 
                    value="2", 
                    label="⚡ Frame Multiplier",
                    info="2X = Double FPS (e.g. 30 -> 60). Higher multipliers create more intermediate frames."
                )

            with gr.Accordion("🛠️ Advanced Configuration", open=False):
                gr.Markdown("Control rendering parameters.")
                
                with gr.Row():
                    scale_param = gr.Dropdown(
                        choices=[0.25, 0.5, 1.0, 2.0, 4.0],
                        value=1.0,
                        label="📉 Render Scale",
                        info="1.0 = Original Resolution. Reduce to 0.5 for faster processing on 4K content."
                    )
                    fps_param = gr.Number(
                        value=0, 
                        label="🎯 Force Output FPS", 
                        info="0 = Auto-calculate. Set to 30 or 60 to lock the framerate. Audio will be removed when forcing FPS"
                    )
                    exp_param = gr.Number(
                        value=1, 
                        label="🔢 Exponent Power", 
                        info="Alternative multiplier calculation (2^exp)."
                    )

                with gr.Row():
                    uhd_chk = gr.Checkbox(
                        label="💎 UHD/4K Optimization", 
                        value=False,
                        info="Enable for high-resolution sources."
                    )
                    skip_chk = gr.Checkbox(
                        label="⏩ Skip Static Frames", 
                        value=False,
                        info="Bypass processing for non-moving scenes to save time."
                    )
                    montage_chk = gr.Checkbox(
                        label="🆚 Split-Screen Comparison", 
                        value=False,
                        info="Output video showing Original vs. Processed."
                    )

            btn_run = gr.Button("🚀 RENDER HIGH FPS VIDEO", variant="primary", size="lg")

        # --- Right Column: Output ---
        with gr.Column(scale=1):
            output_vid = gr.Video(label="⚡ High FPS Result")
            gr.Markdown("**Status:** Rendering time depends on input resolution and duration.")

    # --- Bind Logic ---
    btn_run.click(
        fn=run_rife,
        inputs=[
            input_vid, 
            multi_param, 
            exp_param, 
            fps_param, 
            scale_param, 
            uhd_chk, 
            skip_chk, 
            montage_chk
        ],
        outputs=output_vid
    )

if __name__ == "__main__":
    app.launch(
        theme=gr.themes.Soft(),
        mcp_server=True,
    )