import os
import gradio as gr
import shutil
import subprocess
from pathlib import Path

# Model repo URL (update to your actual repo if it's private or elsewhere)
MODEL_REPO_URL = "https://github.com/yourname/Wan2.2-Animate-14B.git"
CKPT_DIR = "Wan2.2-Animate-14B"
CKPT_PATH = f"./{CKPT_DIR}/process_checkpoint"
PROCESS_SCRIPT = f"{CKPT_DIR}/modules/animate/preprocess/preprocess_data.py"
GENERATE_SCRIPT = "generate.py"

def clone_model_repo():
    if not os.path.exists(CKPT_DIR):
        print("Cloning model repository...")
        run_command(f"git clone {MODEL_REPO_URL} {CKPT_DIR}")
    else:
        print("Model repo already exists.")

def run_command(command):
    print(f"Running command:\n{command}\n")
    result = subprocess.run(command, shell=True, capture_output=True, text=True)
    print(result.stdout)
    print(result.stderr)
    return result.stdout + "\n" + result.stderr

def process_and_generate(video, image, mode):
    if not video or not image:
        return "Error: Please upload both video and reference image."

    # Clone model repo if not exists
    clone_model_repo()

    # Setup temp directories
    temp_base = f"./temp_runs/{mode.lower()}"
    input_video_path = os.path.join(temp_base, "video.mp4")
    input_image_path = os.path.join(temp_base, "image.jpeg")
    process_results_path = os.path.join(temp_base, "process_results")

    if os.path.exists(temp_base):
        shutil.rmtree(temp_base)
    os.makedirs(temp_base, exist_ok=True)

    # Save inputs
    video.save(input_video_path)
    image.save(input_image_path)
    os.makedirs(process_results_path, exist_ok=True)

    # Build preprocess command
    preprocess_cmd = f"""
    python {PROCESS_SCRIPT} \
        --ckpt_path {CKPT_PATH} \
        --video_path {input_video_path} \
        --refer_path {input_image_path} \
        --save_path {process_results_path} \
        --resolution_area 1280 720
    """