#!/usr/bin/env python3
"""
Wav2Lip installer for GOLS — sets up real lip-sync (moving lips).
===============================================================
Run from the project root:

    python models/install_wav2lip.py

What it does
------------
  1. Clones the Wav2Lip repo into  models/wav2lip/
  2. Downloads the two required checkpoints:
        • checkpoints/wav2lip_gan.pth          (the lip-sync model)
        • face_detection/detection/sfd/s3fd.pth (face detector)
  3. Installs Wav2Lip's Python dependencies (torch, opencv, librosa, …)

After it finishes, set this in backend/.env (already the default):
    WAV2LIP_PATH=models/wav2lip
    MODELS_PYTHON=python        # the interpreter that has Wav2Lip's deps

Then restart the backend and regenerate a video — the pipeline auto-detects
Wav2Lip and the avatar's lips will move in sync with the voice.

Notes
-----
  • Wav2Lip runs on CPU (slow: ~1–4 min for a short clip) or CUDA/MPS if available.
  • The checkpoint URLs below are public mirrors. If one 404s, see the printed
    fallback links or download manually into the paths shown.
"""

import os
import subprocess
import sys
import urllib.request

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DEST = os.path.join(ROOT, "models", "wav2lip")
REPO = "https://github.com/Rudrabha/Wav2Lip.git"

# Public mirrors for the checkpoints (tried in order until one works).
CHECKPOINTS = {
    os.path.join(DEST, "checkpoints", "wav2lip_gan.pth"): [
        "https://huggingface.co/Nekochu/Wav2Lip/resolve/main/wav2lip_gan.pth",
        "https://huggingface.co/camenduru/Wav2Lip/resolve/main/checkpoints/wav2lip_gan.pth",
    ],
    os.path.join(DEST, "face_detection", "detection", "sfd", "s3fd.pth"): [
        "https://huggingface.co/camenduru/Wav2Lip/resolve/main/face_detection/detection/sfd/s3fd.pth",
        "https://www.adrianbulat.com/downloads/python-fan/s3fd-619a316812.pth",
    ],
}


def sh(cmd, **kw):
    print(f"  $ {' '.join(cmd)}")
    subprocess.check_call(cmd, **kw)


def download(dest, urls):
    if os.path.isfile(dest) and os.path.getsize(dest) > 1_000_000:
        print(f"  ✓ already present: {dest}")
        return
    os.makedirs(os.path.dirname(dest), exist_ok=True)
    for url in urls:
        try:
            print(f"  ↓ {url}")
            urllib.request.urlretrieve(url, dest)
            if os.path.getsize(dest) > 1_000_000:
                print(f"  ✓ saved {dest} ({os.path.getsize(dest)//1_048_576} MB)")
                return
        except Exception as e:
            print(f"    ! failed: {e}")
    print(f"  ✗ Could not download automatically. Manually place a file at:\n      {dest}\n"
          f"    from one of:\n      " + "\n      ".join(urls))


def main():
    print("== GOLS Wav2Lip installer ==")

    # 1. Clone
    if os.path.isdir(os.path.join(DEST, ".git")) or os.path.isfile(os.path.join(DEST, "inference.py")):
        print(f"  ✓ repo already at {DEST}")
    else:
        sh(["git", "clone", "--depth", "1", REPO, DEST])

    # 2. Checkpoints
    print("Downloading checkpoints …")
    for dest, urls in CHECKPOINTS.items():
        download(dest, urls)

    # 3. Dependencies
    print("Installing Python dependencies …")
    deps = ["torch", "torchvision", "numpy", "opencv-python", "librosa==0.10.2",
            "numba", "tqdm", "scipy"]
    try:
        sh([sys.executable, "-m", "pip", "install", *deps])
    except Exception as e:
        print(f"  ! pip install failed ({e}). Install these manually: {' '.join(deps)}")

    ok = all(os.path.isfile(p) and os.path.getsize(p) > 1_000_000 for p in CHECKPOINTS)
    print("\n== Done ==")
    print("Checkpoints OK:", ok)
    print("Next: set WAV2LIP_PATH=models/wav2lip in backend/.env, restart the backend,")
    print("and regenerate a video. The lips will now move in sync with the voice.")
    if not ok:
        print("\n⚠ One or more checkpoints are missing — download them manually (see paths above).")


if __name__ == "__main__":
    main()