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

# --- 1. BOOTSTRAP ENVIRONMENT ---
def setup_repos():
    if not os.path.exists("/app/SadTalker"):
        print("📥 Cloning Repositories...")
        subprocess.run(["git", "clone", "https://github.com/OpenTalker/SadTalker.git", "/app/SadTalker"])
        subprocess.run(["git", "clone", "https://github.com/Rudrabha/Wav2Lip.git", "/app/Wav2Lip"])
        # Fix BasicSR compatibility
        subprocess.run(["find", "/usr/local/lib/python3.10/site-packages/basicsr", "-name", "degradations.py", "-exec", "sed", "-i", "s/functional_tensor/functional/g", "{}", "+"])

setup_repos()

# --- 2. THE GPU-ACCELERATED CORE ---
@spaces.GPU(duration=120) # Grants H200 access for 2 minutes per click
def generate(image, audio):
    # Setup paths
    workspace = Path("/tmp/visor_workspace")
    workspace.mkdir(parents=True, exist_ok=True)
    
    img_path = workspace / "input.jpg"
    aud_path = workspace / "input.mp3"
    
    # Gradio provides file paths directly
    shutil.copy(image, img_path)
    shutil.copy(audio, aud_path)

    # Note: On HF Spaces, you should use their 'checkpoints' or 
    # use 'gdown' to pull your weights into /app/SadTalker/checkpoints
    # For testing, SadTalker will auto-download if folder is empty.

    print("🎬 Running Animation...")
    subprocess.run([
        "python", "/app/SadTalker/inference.py",
        "--driven_audio", str(aud_path),
        "--source_image", str(img_path),
        "--result_dir", "/tmp/results",
        "--still", "--preprocess", "full"
    ], env={**os.environ, "PYTHONPATH": "/app/SadTalker"})

    # Return the first mp4 found
    result_video = list(Path("/tmp/results").glob("**/*.mp4"))
    return result_video[0] if result_video else None

# --- 3. GRADIO INTERFACE (The Frontend) ---
with gr.Blocks(title="VisorFlow Core") as demo:
    gr.Markdown("# 🛡️ VisorFlow Core: ZeroGPU Edition")
    with gr.Row():
        with gr.Column():
            input_img = gr.Image(type="filepath", label="Source Image")
            input_aud = gr.Audio(type="filepath", label="Driven Audio")
            run_btn = gr.Button("Execute Phase 3", variant="primary")
        with gr.Column():
            output_video = gr.Video(label="Generated Intelligence")

    run_btn.click(fn=generate, inputs=[input_img, input_aud], outputs=[output_video])

demo.launch()