import os
import sys
import subprocess
import uuid
import torch
import gradio as gr
from pathlib import Path
from huggingface_hub import hf_hub_download

# --- Configuration ---
WAV2LIP_URL = "https://github.com/Rudrabha/Wav2Lip.git"
WAV2LIP_DIR = "/app/Wav2Lip"
CHECKPOINT_URL = "https://huggingface.co/Lykos-AI/Wav2Lip/resolve/main/wav2lip.pth"
S3FD_URL = "https://huggingface.co/Lykos-AI/Wav2Lip/resolve/main/s3fd.pth"

# --- Auto-Setup Wav2Lip 🚀 ---
if not os.path.exists(WAV2LIP_DIR):
    print(f"Cloning Wav2Lip...")
    subprocess.run(["git", "clone", WAV2LIP_URL, WAV2LIP_DIR], check=True)
    
    print("Installing lightweight dependencies...")
    subprocess.run(["pip", "install", "librosa", "opencv-python", "boto3", "requests", "tqdm", "ffmpeg-python"], check=True)

    # Setup directories
    os.makedirs(f"{WAV2LIP_DIR}/checkpoints", exist_ok=True)
    os.makedirs(f"{WAV2LIP_DIR}/face_detection/detection/sfd", exist_ok=True)

    print("Downloading checkpoints...")
    hf_hub_download(repo_id="camenduru/Wav2Lip", filename="checkpoints/wav2lip.pth", local_dir=WAV2LIP_DIR)
    hf_hub_download(repo_id="camenduru/Wav2Lip", filename="face_detection/detection/sfd/s3fd.pth", local_dir=WAV2LIP_DIR)

# Add to path
if WAV2LIP_DIR not in sys.path:
    sys.path.append(WAV2LIP_DIR)

# --- Inference Logic ---
def inference(video_path, audio_path):
    if not video_path or not audio_path:
        return None

    output_id = str(uuid.uuid4())
    output_path = f"outputs/{output_id}.mp4"
    os.makedirs("outputs", exist_ok=True)

    print(f"Starting Wav2Lip inference for {output_id}...")

    # Run Wav2Lip inference script via subprocess
    # Using --nosmooth to save CPU memory/time
    cmd = [
        "python", f"{WAV2LIP_DIR}/inference.py",
        "--checkpoint_path", f"{WAV2LIP_DIR}/checkpoints/wav2lip.pth",
        "--face", video_path,
        "--audio", audio_path,
        "--outfile", output_path,
        "--pads", "0", "10", "0", "0", # Standard padding for better results
        "--resize_factor", "2" # Resize to speed up CPU inference
    ]

    try:
        subprocess.run(cmd, check=True)
        return output_path
    except Exception as e:
        print(f"Wav2Lip Error: {e}")
        return None

# --- Gradio UI ---
with gr.Blocks() as demo:
    gr.Markdown("## 🎤 LipSync AI — Lightweight Node (Wav2Lip CPU)")
    gr.Markdown("This node uses Wav2Lip optimized for CPU. It's slower than GPU but extremely reliable.")
    gr.Markdown("### 📸 Photo-to-Video Generation")
    
    with gr.Row():
        with gr.Column():
            input_face = gr.Image(label="Face Photo (JPG/PNG)", type="filepath")
            input_audio = gr.Audio(label="Voice Audio", type="filepath")
            btn = gr.Button("🚀 Generate Lip-Sync", variant="primary")
        
        with gr.Column():
            output_video = gr.Video(label="Result Video")

    btn.click(
        fn=inference,
        inputs=[input_face, input_audio],
        outputs=[output_video],
        api_name="predict" # Same API name for the rotation engine
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft())