from __future__ import annotations

import argparse
from pathlib import Path
import sys

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))

from huggingface_hub import snapshot_download  # type: ignore

from personalization.config.settings import load_local_models_config


def pull_one(repo_id: str, dest: Path, force: bool = False) -> None:
    dest.mkdir(parents=True, exist_ok=True)
    if any(dest.iterdir()) and not force:
        print(f"[skip] {dest} already populated. Use --force to overwrite.")
        return
    print(f"[pull] {repo_id} -> {dest}")
    snapshot_download(repo_id=repo_id, local_dir=str(dest), local_dir_use_symlinks=False)
    print(f"[done] {repo_id}")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--target",
        choices=[
            "llm",
            "preference_extractor",
            "embed_qwen3",
            "embed_nemotron",
            "embedders",
            "reranker_qwen3",
            "rerankers",
            "all",
        ],
        default="all",
    )
    parser.add_argument("--force", action="store_true")
    args = parser.parse_args()

    cfg = load_local_models_config()
    if args.target in ("llm", "all"):
        pull_one(cfg.llm.hf_id, ROOT / cfg.llm.local_path, force=args.force)
    if args.target in ("preference_extractor", "all"):
        pull_one(
            cfg.preference_extractor.hf_id,
            ROOT / cfg.preference_extractor.local_path,
            force=args.force,
        )
    if args.target in ("embed_qwen3", "embedders", "all") and cfg.embedding and cfg.embedding.qwen3:
        pull_one(
            cfg.embedding.qwen3.hf_id,
            ROOT / cfg.embedding.qwen3.local_path,
            force=args.force,
        )
    if args.target in ("embed_nemotron", "embedders", "all") and cfg.embedding and cfg.embedding.nemotron:
        pull_one(
            cfg.embedding.nemotron.hf_id,
            ROOT / cfg.embedding.nemotron.local_path,
            force=args.force,
        )
    if args.target in ("reranker_qwen3", "rerankers", "all") and cfg.reranker and cfg.reranker.qwen3_8b:
        pull_one(
            cfg.reranker.qwen3_8b.hf_id,
            ROOT / cfg.reranker.qwen3_8b.local_path,
            force=args.force,
        )


if __name__ == "__main__":
    main()