import os
from huggingface_hub import hf_hub_download
import cv2
import paddlehub as hub
import gradio as gr
import torch
from PIL import Image, ImageOps
import numpy as np
import imageio
from moviepy.editor import *

os.makedirs("models", exist_ok=True)
os.makedirs("data", exist_ok=True)
os.makedirs("dataout", exist_ok=True)

model_path = hf_hub_download(
    repo_id="akhaliq/lama",
    filename="best.ckpt",
    local_dir="models"  # This will download it directly into 'models' directory
)


print(f"Model downloaded to: {model_path}")

def get_frames(video_in):
    frames = []
    clip = VideoFileClip(video_in)

    # Resize and adjust FPS
    if clip.fps > 30:
        print("Video rate is over 30, resetting to 30")
        clip_resized = clip.resize(height=256)
        clip_resized.write_videofile("video_resized.mp4", fps=30)
    else:
        print("Video rate is OK")
        clip_resized = clip.resize(height=256)
        clip_resized.write_videofile("video_resized.mp4", fps=clip.fps)

    # Extract frames
    cap = cv2.VideoCapture("video_resized.mp4")
    fps = cap.get(cv2.CAP_PROP_FPS)
    print("Video fps:", fps)
    i = 0
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        frame_path = f'kang_{i}.jpg'
        cv2.imwrite(frame_path, frame)
        frames.append(frame_path)
        i += 1
    
    cap.release()
    cv2.destroyAllWindows()
    print("Video broken into frames")
    return frames, fps

def create_video(frames, fps, type_name):
    print("Building video result")
    clip = ImageSequenceClip(frames, fps=fps)
    output_file = f"{type_name}_result.mp4"
    clip.write_videofile(output_file, fps=fps)
    return output_file

def magic_lama(img_path):
    img = Image.open(img_path)
    mask = Image.open("./masks/modelscope-mask.png")
    inverted_mask = ImageOps.invert(mask)

    imageio.imwrite("./data/data.png", img)
    imageio.imwrite("./data/data_mask.png", inverted_mask)
    os.system('python predict.py model.path=/home/user/app/ indir=/home/user/app/data/ outdir=/home/user/app/dataout/ device=cpu')
    return "./dataout/data_mask.png"

def infer(video_in):
    frames_list, fps = get_frames(video_in)
    n_frame = len(frames_list)
    
    result_frames = []
    for i, frame_path in enumerate(frames_list[:n_frame]):
        lama_frame = magic_lama(frame_path)
        cleaned_frame_path = f"cleaned_frame_{i}.jpg"
        Image.open(lama_frame).save(cleaned_frame_path)
        result_frames.append(cleaned_frame_path)
        print(f"Processed frame {i + 1}/{n_frame}")

    final_video = create_video(result_frames, fps, "cleaned")
    return final_video

# Gradio Interface
inputs = gr.Video(label="Input")
outputs = gr.Video(label="Output")
title = "LaMa Video Watermark Remover"
description = (
    "<p style='text-align: center'>LaMa: Resolution-robust Large Mask Inpainting with Fourier Convolutions.<br />"
    "This demo is meant to be used as a watermark remover on Modelscope generated videos.<br />"
    "Simply upload your Modelscope video and hit Submit.</p>"
)
article = (
    "<p style='text-align: center'><a href='https://arxiv.org/abs/2109.07161' target='_blank'>Resolution-robust Large Mask Inpainting with Fourier Convolutions</a> | "
    "<a href='https://github.com/saic-mdal/lama' target='_blank'>Github Repo</a></p>"
)
examples = [
    # "./examples/modelscope-astronaut-horse.mp4",
    # "./examples/modelscope-panda.mp4",
    "./examples/modelscope-spiderman-surfing.mp4"
]

demo = gr.Interface(
    fn=infer,
    inputs=inputs,
    outputs=outputs,
    title=title,
    description=description,
    article=article,
    examples=examples
)

# Launch with prevent_thread_lock in case it's needed for async compatibility in v4
demo.launch(prevent_thread_lock=True)