import os
import time
import datetime
import argparse
import torch
import gradio as gr
from huggingface_hub import snapshot_download
from diffueraser.diffueraser import DiffuEraser
from propainter.inference import Propainter, get_device

# Handle MoviePy import and flag
try:
    from moviepy.editor import VideoFileClip
    _HAS_MOVIEPY = True
except ImportError:
    _HAS_MOVIEPY = False

# ------------------------------------------------------------
# Optional Hugging Face Spaces GPU integration
# ------------------------------------------------------------
try:
    import spaces  
    _HAS_SPACES = True
except Exception:
    spaces = None
    _HAS_SPACES = False

_IN_SPACES = bool(
    os.getenv("SPACE_ID")
    or os.getenv("SPACE_REPO_NAME")
    or os.getenv("HF_SPACE_ID")
    or os.getenv("SYSTEM") == "spaces"
)

def maybe_spaces_gpu(duration: int = 100):
    def _decorator(fn):
        if _HAS_SPACES and _IN_SPACES:
            return spaces.GPU(duration=duration)(fn)
        return fn
    return _decorator

# ------------------------------------------------------------
# Download weights
# ------------------------------------------------------------
subfolders = ["diffuEraser", "stable-diffusion-v1-5", "PCM_Weights", "propainter", "sd-vae-ft-mse"]
for subfolder in subfolders:
    os.makedirs(os.path.join("weights", subfolder), exist_ok=True)

snapshot_download(repo_id="lixiaowen/diffuEraser", local_dir="./weights/diffuEraser")
snapshot_download(repo_id="stable-diffusion-v1-5/stable-diffusion-v1-5", local_dir="./weights/stable-diffusion-v1-5")
snapshot_download(repo_id="wangfuyun/PCM_Weights", local_dir="./weights/PCM_Weights")
snapshot_download(repo_id="camenduru/ProPainter", local_dir="./weights/propainter")
snapshot_download(repo_id="stabilityai/sd-vae-ft-mse", local_dir="./weights/sd-vae-ft-mse")

# ------------------------------------------------------------
# Model init
# ------------------------------------------------------------
base_model_path = "weights/stable-diffusion-v1-5"
vae_path = "weights/sd-vae-ft-mse"
diffueraser_path = "weights/diffuEraser"
propainter_model_dir = "weights/propainter"

device = get_device()
ckpt = "2-Step"
video_inpainting_sd = DiffuEraser(device, base_model_path, vae_path, diffueraser_path, ckpt=ckpt)
propainter = Propainter(propainter_model_dir, device=device)

# ------------------------------------------------------------
# Helper: Normalize Video
# ------------------------------------------------------------
def normalize_video(input_path: str, output_path: str, fps: float = None, max_duration: float = None):
    if not _HAS_MOVIEPY:
        return None
    clip = VideoFileClip(input_path)
    if max_duration is not None:
        clip = clip.subclip(0, min(max_duration, clip.duration))
    if fps is None:
        fps = float(clip.fps)
    clip.write_videofile(output_path, codec="libx264", audio=False, fps=fps)
    clip.close()
    return fps

# ------------------------------------------------------------
# Inference Function
# ------------------------------------------------------------
@maybe_spaces_gpu(duration=600)
def infer(input_video, input_mask):
    video_length = 600
    mask_dilation_iter = 8
    max_img_size = 1080
    ref_stride = 4
    neighbor_length = 35
    subvideo_length = 95
    
    save_path = "results"
    os.makedirs(save_path, exist_ok=True)
    
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    priori_path = os.path.join(save_path, f"priori_{timestamp}.mp4")
    output_path = os.path.join(save_path, f"diffueraser_result_{timestamp}.mp4")

    norm_video_path = input_video
    norm_mask_path = input_mask

    if _HAS_MOVIEPY:
        norm_video_path = os.path.join(save_path, f"video_norm_{timestamp}.mp4")
        norm_mask_path = os.path.join(save_path, f"mask_norm_{timestamp}.mp4")
        try:
            normalize_video(input_video, norm_video_path, fps=30.0, max_duration=None)
            normalize_video(input_mask, norm_mask_path, fps=30.0, max_duration=None)
        except Exception:
            norm_video_path = input_video
            norm_mask_path = input_mask

    start_time = time.time()
    propainter.forward(
        norm_video_path, norm_mask_path, priori_path,
        video_length=video_length, ref_stride=ref_stride,
        neighbor_length=neighbor_length, subvideo_length=subvideo_length,
        mask_dilation=mask_dilation_iter
    )

    video_inpainting_sd.forward(
        norm_video_path, norm_mask_path, priori_path, output_path,
        max_img_size=max_img_size, video_length=video_length,
        mask_dilation_iter=mask_dilation_iter, guidance_scale=None
    )
    
    print(f"Inference time: {time.time() - start_time:.4f} s")
    torch.cuda.empty_cache()
    return output_path

# ------------------------------------------------------------
# Gradio UI
# ------------------------------------------------------------
with gr.Blocks() as demo:
    gr.Markdown("# DiffuEraser: Video Inpainting")
    with gr.Row():
        with gr.Column():
            input_video = gr.Video(label="Input Video")
            input_mask = gr.Video(label="Input Mask")
            submit_btn = gr.Button("Submit")
        with gr.Column():
            video_result = gr.Video(label="Result")
            
    submit_btn.click(fn=infer, inputs=[input_video, input_mask], outputs=[video_result])


demo.queue().launch(
    show_error=True,
    share=True,
    footer_links=["gradio", "settings"],  # equivalent to old show_api=False
)