from __future__ import annotations

import sys
import subprocess
from pathlib import Path

from huggingface_hub import snapshot_download


ROOT = Path(__file__).resolve().parent

ARC2FACE_CODE = ROOT / "Arc2Face"
MODEL_DIR = ROOT / "models"

GITHUB_REPO = "https://github.com/foivospar/Arc2Face.git"
HF_REPO = "FoivosPar/Arc2Face"

REQUIRED_MODEL_FILES = [
    "encoder/config.json",
    "encoder/pytorch_model.bin",
    "arc2face/config.json",
    "arc2face/diffusion_pytorch_model.safetensors",
]


def setup_arc2face():

    # --------------------------------------------------
    # 1. Download official Arc2Face Python implementation
    # --------------------------------------------------

    package_init = ARC2FACE_CODE / "arc2face" / "__init__.py"

    if not package_init.exists():

        if ARC2FACE_CODE.exists():
            raise RuntimeError(
                f"Arc2Face directory exists but is incomplete: "
                f"{ARC2FACE_CODE}. Remove it and restart the Space."
            )

        print("Downloading Arc2Face GitHub repository...")

        subprocess.run(
            [
                "git",
                "clone",
                "--depth",
                "1",
                GITHUB_REPO,
                str(ARC2FACE_CODE),
            ],
            check=True,
        )

    # Make the repository root importable.
    repo_path = str(ARC2FACE_CODE)

    if repo_path not in sys.path:
        sys.path.insert(0, repo_path)

    if not package_init.exists():
        raise FileNotFoundError(
            f"Arc2Face Python package not found: {package_init}"
        )

    # --------------------------------------------------
    # 2. Download pretrained Arc2Face model weights
    # --------------------------------------------------

    MODEL_DIR.mkdir(parents=True, exist_ok=True)

    missing_files = [
        filename
        for filename in REQUIRED_MODEL_FILES
        if not (MODEL_DIR / filename).exists()
    ]

    if missing_files:

        print("Downloading Arc2Face pretrained weights...")

        snapshot_download(
            repo_id=HF_REPO,
            local_dir=str(MODEL_DIR),
            allow_patterns=REQUIRED_MODEL_FILES,
        )

    # --------------------------------------------------
    # 3. Verify all required files
    # --------------------------------------------------

    missing_files = [
        filename
        for filename in REQUIRED_MODEL_FILES
        if not (MODEL_DIR / filename).is_file()
    ]

    if missing_files:
        raise FileNotFoundError(
            "Arc2Face model download incomplete. Missing:\n"
            + "\n".join(missing_files)
        )

    print("Arc2Face code and pretrained weights are ready.")

    return ARC2FACE_CODE, MODEL_DIR


if __name__ == "__main__":
    setup_arc2face()