import os
import shutil
import subprocess
import spaces
import gradio as gr
from pathlib import Path
from huggingface_hub import snapshot_download

# --- 1. CONFIGURATION & WEIGHT HYDRATION ---
REPO_ID = "Macfeigh/visor-weights" 
CHECKPOINT_DIR = "/app/SadTalker/checkpoints"

def hydrate_workspace():
    # Clone SadTalker and Wav2Lip if they don't exist
    if not os.path.exists("/app/SadTalker"):
        print("📥 Cloning SadTalker...")
        subprocess.run(["git", "clone", "https://github.com/OpenTalker/SadTalker.git", "/app/SadTalker"])
        
    if not os.path.exists("/app/Wav2Lip"):
        print("📥 Cloning Wav2Lip...")
        subprocess.run(["git", "clone", "https://github.com/Rudrabha/Wav2Lip.git", "/app/Wav2Lip"])

    # Download weights from your model repo
    print(f"🌡️ Pulling weights from {REPO_ID}...")
    snapshot_download(
        repo_id=REPO_ID,
        local_dir=CHECKPOINT_DIR,
        local_dir_use_symlinks=False
    )
    print("✅ Weights ready.")

# Initialize workspace before Gradio starts
hydrate_workspace()

# --- 2. GPU INFERENCE ENGINE ---
@spaces.GPU(duration=120)
def visor_execute(image_path, audio_path):
    output_dir = "/tmp/visor_output"
    if os.path.exists(output_dir):
        shutil.rmtree(output_dir)
    os.makedirs(output_dir, exist_ok=True)

    print("🚀 Executing Phase 3: Animation...")
    
    # Run SadTalker
    st_proc = subprocess.run([
        "python", "/app/SadTalker/inference.py",
        "--driven_audio", audio_path,
        "--source_image", image_path,
        "--result_dir", output_dir,
        "--still",
        "--preprocess", "full",
        "--checkpoint_dir", CHECKPOINT_DIR
    ], env={**os.environ, "PYTHONPATH": "/app/SadTalker"})

    # Find result
    generated_videos = list(Path(output_dir).glob("**/*.mp4"))
    if not generated_videos:
        return None
    
    return str(generated_videos[0])

# --- 3. INTERFACE ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🛡️ VisorFlow Core: Sovereign Intelligence Node")
    gr.Markdown("Zero-cost execution environment running on NVIDIA H200.")
    
    with gr.Row():
        with gr.Column():
            img = gr.Image(type="filepath", label="Source Portrait")
            aud = gr.Audio(type="filepath", label="Voice Command")
            btn = gr.Button("RUN PHASE 3", variant="primary")
        with gr.Column():
            out = gr.Video(label="Generated Output")

    btn.click(fn=visor_execute, inputs=[img, aud], outputs=[out])

demo.launch()