import uvicorn
import os
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import FileResponse
import subprocess
import uuid
import os

def validate_checkpoint(path):
    if not os.path.exists(path):
        raise Exception(f"Checkpoint missing: {path}")
    if os.path.getsize(path) < 100 * 1024 * 1024:  # <100MB = corrupted
        raise Exception(f"Checkpoint corrupted or incomplete: {path}")

# Validate
validate_checkpoint("Wav2Lip/Wav2Lip.pth")
validate_checkpoint("Wav2Lip/wav2lip_gan.pth")

app = FastAPI(title="Wav2Lip API (HuggingFace)")

# Clone repo once
if not os.path.exists("Wav2Lip"):
    subprocess.run(["git", "clone", "https://github.com/Rudrabha/Wav2Lip.git"])

# Download pretrained model once
if not os.path.exists("Wav2Lip/Wav2Lip.pth"):
    subprocess.run([
        "wget",
        "-O",
        "Wav2Lip/Wav2Lip.pth",
        "https://huggingface.co/spaces/gopal/wav2lip-model/resolve/main/Wav2Lip.pth"
    ])

@app.post("/generate")
async def generate(
    face: UploadFile = File(...),
    audio: UploadFile = File(...)
):
    session = str(uuid.uuid4())
    os.makedirs(session)

    face_path = f"{session}/face.{face.filename.split('.')[-1]}"
    audio_path = f"{session}/audio.{audio.filename.split('.')[-1]}"
    output_path = f"{session}/result.mp4"

    with open(face_path, "wb") as f:
        f.write(await face.read())

    with open(audio_path, "wb") as f:
        f.write(await audio.read())

    cmd = [
        "python",
        "Wav2Lip/inference.py",
        "--checkpoint_path",
        "Wav2Lip/Wav2Lip.pth",
        "--face",
        face_path,
        "--audio",
        audio_path,
        "--outfile",
        output_path
    ]

    subprocess.run(cmd)

    return FileResponse(output_path, media_type="video/mp4", filename="output.mp4")

@app.get("/")
def home():
    return {"message": "Wav2Lip FastAPI is running on HuggingFace!"}

if __name__ == "__main__":
    uvicorn.run("app:app", host="0.0.0.0", port=7860)