from __future__ import annotations

"""
Standalone model download script.

Usage:
  python model_download.py
"""

import os
from pathlib import Path

from dotenv import load_dotenv

DEFAULT_ENV_FILE = ".env"
DEFAULT_MODEL_ID = "Qwen/Qwen3.5-9B"
DEFAULT_MODEL_DIR = "./models/Qwen3.5-9B"
DEFAULT_CACHE_DIR = "./modelscope_cache"


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 load_runtime_env(script_dir: Path) -> Path:
    env_name = (os.getenv("ENV_FILE", DEFAULT_ENV_FILE) or DEFAULT_ENV_FILE).strip()
    env_path = (script_dir / env_name).resolve()
    if env_path.exists():
        load_dotenv(env_path)
    return env_path


def env_flag(name: str, default: bool = False) -> bool:
    raw = os.getenv(name, "")
    if not raw:
        return default
    return raw.strip().lower() in {"1", "true", "yes", "on"}


def download_model(
    model_id: str,
    model_dir: Path,
    cache_dir: Path,
    revision: str = "",
    skip_if_exists: bool = False,
) -> Path:
    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

    model_dir.parent.mkdir(parents=True, exist_ok=True)
    cache_dir.mkdir(parents=True, exist_ok=True)

    if skip_if_exists and model_dir.exists() and any(model_dir.iterdir()):
        print(f"[INFO] model already exists, skip download: {model_dir}")
        return model_dir.resolve()

    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}")
    return Path(downloaded_path).resolve()


def main() -> None:
    script_dir = Path(__file__).resolve().parent
    env_path = load_runtime_env(script_dir)
    model_id = os.getenv("DOWNLOAD_MODEL_ID", os.getenv("MODEL_SOURCE", DEFAULT_MODEL_ID)).strip()
    model_dir_raw = os.getenv("DOWNLOAD_SAVE_DIR", os.getenv("MODEL_DIR", DEFAULT_MODEL_DIR)).strip()
    cache_dir_raw = os.getenv("DOWNLOAD_CACHE_DIR", os.getenv("MODELSCOPE_CACHE", DEFAULT_CACHE_DIR)).strip()
    revision = os.getenv("DOWNLOAD_REVISION", "").strip()
    skip_if_exists = env_flag("SKIP_MODEL_DOWNLOAD_IF_EXISTS", True)

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

    model_dir = resolve_path(model_dir_raw, script_dir)
    cache_dir = resolve_path(cache_dir_raw, script_dir)
    print(f"[INFO] env_file={env_path}")
    download_model(
        model_id=model_id,
        model_dir=model_dir,
        cache_dir=cache_dir,
        revision=revision,
        skip_if_exists=skip_if_exists,
    )


if __name__ == "__main__":
    main()