import gradio as gr
import os, uuid, time
from gradio_client import Client, handle_file
from moviepy.editor import VideoFileClip   # NEW

hf_token   = os.environ.get("TOKEN")
output_dir = "uploads/output"
os.makedirs(output_dir, exist_ok=True)

client = Client("tonyassi/vfs2-cpu", hf_token=hf_token, download_files=output_dir)

# ── helper ───────────────────────────────────────────────────────────
def preprocess_video(path: str, target_fps: int = 12,
                     target_size: int = 800, target_len: int = 4) -> str:
    """
    Returns a path to a temp-file that is
    ≤ target_len seconds, target_fps FPS,
    and resized so the longest side == target_size.
    """
    clip = VideoFileClip(path)

    # 1) trim
    if clip.duration > target_len:
        clip = clip.subclip(0, target_len)

    # 2) resize (longest side)
    w, h = clip.size
    if w >= h:
        clip = clip.resize(width=target_size)
    else:
        clip = clip.resize(height=target_size)

    # 3) FPS
    clip = clip.set_fps(target_fps)

    # 4) write to temp file
    out_path = os.path.join(
        output_dir, f"pre_{uuid.uuid4().hex}.mp4"
    )
    clip.write_videofile(
        out_path,
        codec="libx264",
        audio_codec="aac",
        fps=target_fps,
        verbose=False,
        logger=None
    )
    clip.close()
    return out_path

# ── main generate ────────────────────────────────────────────────────
def generate(input_image, input_video, gender):

    if (gender=="all"): gender = None
        
    try:
        # ⇢ preprocess video locally
        pre_video = preprocess_video(input_video)

        # ⇢ submit job to remote Space
        job = client.submit(
            input_image=handle_file(input_image),
            input_video={"video": handle_file(pre_video)},
            device='cpu',
            selector='many',
            gender=gender,
            race=None,
            order=None,
            api_name="/predict"
        )

        # wait for completion
        while not job.done():
            print("Waiting for remote job to finish...")
            time.sleep(5)

        if not job.status().success:
            print("Job failed or was cancelled:", job.status())
            return None

        video_path = job.outputs()[0]["video"]
        return video_path

    except Exception as e:
        print("Error occurred:", e)
        return None

# ── Gradio interface (unchanged) ─────────────────────────────────────
demo = gr.Interface(
    fn=generate,
    title="Deep Fake Video",
    description="""
    by [Tony Assi](https://www.tonyassi.com/)
    
    Swaps the source image face onto the target video. Videos get downsampled to 800 pixels, 4 sec, and 12 fps to cut down render time (~5 minutes). Please ❤️ this Space. tony.assi.media@gmail.com for business inquiries.
    """,
    inputs=[
        gr.Image(type="filepath", label="Source Image"),
        gr.Video(label="Target Video"),
        gr.Radio(choices=["all", "female", "male"], label="Gender", value="all"),
    ],
    outputs=gr.Video(label="Result"),
    flagging_mode="never",
    examples=[['bella.jpg', 'wizard.mp4', 'all']],
    cache_examples=True,
)

if __name__ == "__main__":
    demo.launch()