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",
    "RIFE_trained_model_v4.13.2.zip"
], check=True)

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

@spaces.GPU()
def run_rife(
    input_video, 
    model_dir, 
    multi, 
    exp, 
    fps, 
    scale, 
    uhd, 
    fp16, 
    skip, 
    montage, 
    png_mode, 
    ext
):
    """
    Constructs the command line arguments based on Gradio inputs 
    and runs the inference_video.py script via subprocess.
    """
    
    if input_video is None:
        raise gr.Error("Please upload a video first.")

    # 1. Define Output Filename
    output_name = f"output_{multi}X.{ext}"
    
    # 2. Build the Command
    cmd = ["python3", SCRIPT_NAME]
    
    # --video
    cmd.extend(["--video", input_video])
    
    # --output
    cmd.extend(["--output", output_name])
    
    # --multi (Multiplier)
    cmd.extend(["--multi", str(int(multi))])
    
    # --exp 
    # Only add exp if it is not default, or if specific logic requires it. 
    # Usually multi overrides exp in RIFE logic, but we pass it if set.
    if exp != 1:
        cmd.extend(["--exp", str(int(exp))])

    # --fps (Target FPS)
    if fps > 0:
        cmd.extend(["--fps", str(int(fps))])
        
    # --scale (Resolution scale)
    # Check against float 1.0
    if scale != 1.0:
        cmd.extend(["--scale", str(scale)])

    # --ext (Extension)
    cmd.extend(["--ext", ext])

    # --model (Model directory)
    if model_dir and model_dir.strip() != "":
        cmd.extend(["--model", model_dir])

    # --- Boolean Flags ---
    
    if uhd:
        cmd.append("--UHD")
        
    if fp16:
        cmd.append("--fp16")
        
    if skip:
        cmd.append("--skip")
        
    if montage:
        cmd.append("--montage")
        
    if png_mode:
        cmd.append("--png")

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

    # 3. Run the Subprocess
    try:
        # We use a large timeout because video processing takes time
        process = subprocess.run(cmd, capture_output=True, text=True)
        
        # Log stdout/stderr
        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}")

        # 4. Return Result
        if png_mode:
            gr.Info("Processing complete. Output is a folder of PNGs (Video preview unavailable for PNG mode).")
            return None 
        
        if os.path.exists(output_name):
            return output_name
        else:
            raise gr.Error("Output file was not found. Check console for details.")

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


# --- Gradio UI Layout ---

with gr.Blocks(title="RIFE Video Interpolation") as app:
    gr.Markdown("# 🚀 RIFE: Real-Time Intermediate Flow Estimation")
    gr.Markdown("Upload a video to increase its frame rate (smoothness) using AI.")

    with gr.Row():
        # --- Left Column: Inputs & Settings ---
        with gr.Column(scale=1):
            input_vid = gr.Video(label="Input Video", sources=["upload"])
            
            with gr.Group():
                gr.Markdown("### 🎯 Core Parameters")
                
                with gr.Row():
                    multi_param = gr.Dropdown(
                        choices=["2", "4", "8", "16", "32"], 
                        value="2", 
                        label="Interpolation Multiplier (--multi)",
                        info="How many times to multiply the frames. 2X doubles the FPS (e.g., 30fps -> 60fps). 4X quadruples it."
                    )
                    ext_param = gr.Dropdown(
                        choices=["mp4", "avi", "mov", "mkv"], 
                        value="mp4", 
                        label="Output Format (--ext)",
                        info="The file extension for the generated video."
                    )
                
                model_param = gr.Textbox(
                    value="train_log", 
                    label="Model Directory (--model)", 
                    placeholder="train_log",
                    info="Path to the folder containing the trained model files (e.g., 'train_log' or 'rife-v4.6')."
                )

            with gr.Accordion("⚙️ Advanced Settings", open=False):
                gr.Markdown("Fine-tune the inference process.")
                
                with gr.Row():
                    scale_param = gr.Slider(
                        minimum=0.1, maximum=1.0, value=1.0, step=0.1, 
                        label="Input Scale (--scale)", 
                        info="1.0 = Original resolution. Set to 0.5 to reduce memory usage for 4K video inputs."
                    )
                    fps_param = gr.Number(
                        value=0, 
                        label="Force Target FPS (--fps)", 
                        info="0 = Auto-calculate based on multiplier. Enter a number (e.g., 60) to force a specific output frame rate."
                    )
                    exp_param = gr.Number(
                        value=1, 
                        label="Exponent Power (--exp)", 
                        info="Alternative to Multiplier. Sets multiplier to 2^exp. (Usually left at 1 if Multiplier is set)."
                    )

                with gr.Row():
                    uhd_chk = gr.Checkbox(
                        label="UHD Mode (--UHD)", 
                        value=False,
                        info="Optimized for 4K video. Equivalent to setting scale=0.5 manually."
                    )
                    fp16_chk = gr.Checkbox(
                        label="FP16 Mode (--fp16)", 
                        value=True, 
                        info="Uses half-precision floating point. Faster and uses less VRAM with minimal quality loss."
                    )
                
                with gr.Row():
                    skip_chk = gr.Checkbox(
                        label="Skip Static Frames (--skip)", 
                        value=False,
                        info="If the video has frames that don't move, skip processing them to save time."
                    )
                    montage_chk = gr.Checkbox(
                        label="Montage (--montage)", 
                        value=False,
                        info="Creates a video with the Original on the Left and Interpolated on the Right for comparison."
                    )
                    png_chk = gr.Checkbox(
                        label="Output as PNGs (--png)", 
                        value=False,
                        info="Outputs a sequence of images instead of a video file. (Video Preview will be disabled)."
                    )

            btn_run = gr.Button("✨ Start Interpolation", variant="primary", size="lg")

        # --- Right Column: Output ---
        with gr.Column(scale=1):
            output_vid = gr.Video(label="Interpolated Result")
            gr.Markdown("**Note:** Processing time depends on video length, resolution, and your GPU speed.")

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

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