import gradio as gr
import subprocess
import os
from huggingface_hub import hf_hub_download
import glob
from datetime import datetime

# Ensure 'checkpoint' directory exists
os.makedirs("checkpoint", exist_ok=True)

hf_hub_download(
    repo_id="fffiloni/X-Portrait",
    filename="model_state-415001.th",
    local_dir="checkpoint"
)

def extract_frames_with_labels(video_path, output_dir="frames"):
    # Ensure output directory exists
    os.makedirs(output_dir, exist_ok=True)
    
    # Open the video file
    video_capture = cv2.VideoCapture(video_path)
    if not video_capture.isOpened():
        raise ValueError(f"Cannot open video file: {video_path}")
    
    frame_data = []
    frame_index = 0
    
    # Loop through the video frames
    while True:
        ret, frame = video_capture.read()
        if not ret:
            break  # Exit the loop if there are no frames left to read

        # Zero-padded frame index for filename and label
        frame_label = f"{frame_index:04}"
        frame_filename = os.path.join(output_dir, f"frame_{frame_label}.jpg")
        
        # Save the frame as a .jpg file
        cv2.imwrite(frame_filename, frame)
        
        # Append the tuple (filename, label) to the list
        frame_data.append((frame_filename, frame_label))
        
        # Increment frame index
        frame_index += 1
    
    # Release the video capture object
    video_capture.release()
    
    return frame_data

# Define a function to run your script with selected inputs
def run_xportrait(
    model_config, 
    output_dir_base, 
    resume_dir, 
    seed, 
    uc_scale, 
    source_image, 
    driving_video, 
    best_frame, 
    out_frames, 
    num_mix, 
    ddim_steps
):
    # Check if the model weights are in place
    if not os.path.exists(resume_dir):
        return "Model weights not found in checkpoint directory. Please download them first."

    # Create a unique output directory name based on current date and time
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir = os.path.join(output_dir_base, f"output_{timestamp}")
    os.makedirs(output_dir, exist_ok=True)
    
    # Construct the command
    command = [
        "python3", "core/test_xportrait.py",
        "--model_config", model_config,
        "--output_dir", output_dir,
        "--resume_dir", resume_dir,
        "--seed", str(seed),
        "--uc_scale", str(uc_scale),
        "--source_image", source_image,
        "--driving_video", driving_video,
        "--best_frame", str(best_frame),
        "--out_frames", str(out_frames),
        "--num_mix", str(num_mix),
        "--ddim_steps", str(ddim_steps)
    ]
    
    # Run the command
    try:
        subprocess.run(command, check=True)
        
        # Find the generated video file in the output directory
        video_files = glob.glob(os.path.join(output_dir, "*.mp4")) + glob.glob(os.path.join(output_dir, "*.avi"))
        print(video_files)
        if video_files:
            return f"Output video saved at: {video_files[0]}", video_files[0]
        else:
            return "No video file was found in the output directory.", None
    except subprocess.CalledProcessError as e:
        return f"An error occurred: {e}", None

# Set up Gradio interface
app = gr.Interface(
    fn=run_xportrait,
    inputs=[
        gr.Textbox(value="config/cldm_v15_appearance_pose_local_mm.yaml", label="Model Config Path"),
        gr.Textbox(value="outputs", label="Output Directory"),
        gr.Textbox(value="checkpoint/model_state-415001.th", label="Resume Directory"),
        gr.Number(value=999, label="Seed"),
        gr.Number(value=5, label="UC Scale"),
        gr.Image(label="Source Image", type="filepath"),
        gr.Video(label="Driving Video"),
        gr.Number(value=36, label="Best Frame"),
        gr.Number(value=-1, label="Out Frames"),
        gr.Number(value=4, label="Number of Mix"),
        gr.Number(value=30, label="DDIM Steps")
    ],
    outputs=["text", "video"],
    title="XPortrait Model Runner",
    description="Run XPortrait with customizable parameters."
)

# Launch the Gradio app
app.launch()