import os
from pathlib import Path

from .file_utils import ensure_exists, is_windows_drive_path, normalize_path, resolve_path


def _resolve_checkpoint(models_dir, checkpoint):
    checkpoint_text = str(checkpoint).strip().strip('"')
    if os.path.isabs(checkpoint_text) or is_windows_drive_path(checkpoint_text):
        return resolve_path(checkpoint_text)
    return resolve_path(Path(normalize_path(models_dir)) / checkpoint_text)


def load_keysync_model_config(keysync_repo_path, models_dir, keyframes_ckpt, interpolation_ckpt):
    repo_path = ensure_exists(keysync_repo_path, "KeySync repo path", is_dir=True)
    infer_script = repo_path / "scripts" / "infer_raw.sh"
    if not infer_script.is_file():
        raise FileNotFoundError(
            f"scripts/infer_raw.sh not found. Please clone antonibigata/keysync into keysync_repo_path: {repo_path}"
        )

    models_root = ensure_exists(models_dir, "models_dir", is_dir=True)
    keyframes_path = _resolve_checkpoint(models_root, keyframes_ckpt)
    interpolation_path = _resolve_checkpoint(models_root, interpolation_ckpt)

    if not keyframes_path.is_file():
        raise FileNotFoundError(f"keyframes_ckpt does not exist: {keyframes_path}")
    if not interpolation_path.is_file():
        raise FileNotFoundError(f"interpolation_ckpt does not exist: {interpolation_path}")

    return {
        "keysync_repo_path": str(repo_path),
        "models_dir": str(models_root),
        "keyframes_ckpt": str(keyframes_path),
        "interpolation_ckpt": str(interpolation_path),
        "infer_script": str(infer_script),
    }