from __future__ import annotations

"""
Standalone model download script.

Usage:
  python download_model_new.py
"""

from pathlib import Path

# =========================
# User Config
# Modify these variables directly, then run:
#   python download_model_new.py
# =========================
DOWNLOAD_MODEL_ID = "kuohao/gemma-4-26B-A4B-it-FP8"
DOWNLOAD_SAVE_DIR = "./models/gemma-4-26B-A4B-it-FP8"
DOWNLOAD_CACHE_DIR = "./modelscope_cache"
DOWNLOAD_REVISION = ""

def resolve_path(raw: str, base_dir: Path) -> Path:
    path = Path(raw).expanduser()
    if path.is_absolute():
        return path.resolve()
    return (base_dir / path).resolve()


def main() -> None:
    try:
        from modelscope.hub.snapshot_download import snapshot_download
    except Exception as exc:
        raise RuntimeError(
            "Missing dependencies. Please install first:\n"
            "  pip install -r requirements.txt"
        ) from exc

    script_dir = Path(__file__).resolve().parent
    model_id = DOWNLOAD_MODEL_ID.strip()
    model_dir_raw = DOWNLOAD_SAVE_DIR.strip()
    cache_dir_raw = DOWNLOAD_CACHE_DIR.strip()
    revision = DOWNLOAD_REVISION.strip()

    if not model_id:
        raise ValueError("DOWNLOAD_MODEL_ID is empty.")
    if not model_dir_raw:
        raise ValueError("DOWNLOAD_SAVE_DIR is empty.")
    if not cache_dir_raw:
        raise ValueError("DOWNLOAD_CACHE_DIR is empty.")

    model_dir = resolve_path(model_dir_raw, script_dir)
    cache_dir = resolve_path(cache_dir_raw, script_dir)
    model_dir.parent.mkdir(parents=True, exist_ok=True)
    cache_dir.mkdir(parents=True, exist_ok=True)

    print(f"[INFO] model_id={model_id}")
    print(f"[INFO] model_dir={model_dir}")
    print(f"[INFO] cache_dir={cache_dir}")
    if revision:
        print(f"[INFO] revision={revision}")

    kwargs = {
        "model_id": model_id,
        "local_dir": str(model_dir),
        "cache_dir": str(cache_dir),
    }
    if revision:
        kwargs["revision"] = revision

    downloaded_path = snapshot_download(**kwargs)
    print(f"[OK] download complete: {downloaded_path}")


if __name__ == "__main__":
    main()