from pathlib import Path

from .file_utils import normalize_path, resolve_path


REQUIRED_CHECKPOINTS = (
    "checkpoints/keyframe_dub.pt",
    "checkpoints/interpolation_dub.pt",
)


def validate_required_models(models_dir):
    root = resolve_path(models_dir)
    missing = [rel for rel in REQUIRED_CHECKPOINTS if not (root / rel).is_file()]
    if missing:
        missing_text = ", ".join(missing)
        raise FileNotFoundError(
            f"KeySync model files are missing in {root}: {missing_text}. "
            "Run the KeySync Download Models node with download=True, or clone "
            "https://huggingface.co/toninio19/keysync into this folder."
        )
    return root


def download_or_check_models(models_dir, repo_id="toninio19/keysync", download=False):
    root = Path(normalize_path(models_dir))

    if download:
        try:
            from huggingface_hub import snapshot_download
        except ImportError as exc:
            raise RuntimeError(
                "huggingface_hub is required for KeySync Download Models. "
                "Install this wrapper's requirements.txt first."
            ) from exc

        root.mkdir(parents=True, exist_ok=True)
        snapshot_download(
            repo_id=repo_id,
            local_dir=str(root),
            allow_patterns=["checkpoints/*"],
            repo_type="model",
        )
    elif not root.exists():
        raise FileNotFoundError(
            f"models_dir does not exist: {root}. "
            "Run KeySync Download Models with download=True or download the models manually."
        )

    return str(validate_required_models(root))