#!/usr/bin/env python3
"""Download nvidia/Nemotron-Personas-Korea parquet shards for local use."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path


APP_DIR = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(APP_DIR))

from persona_dataset import download_dataset, fetch_manifest  # noqa: E402


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--force", action="store_true", help="Redownload shards that already exist.")
    parser.add_argument(
        "--limit-shards",
        type=int,
        default=None,
        help="Download only the first N shards for smoke tests. Omit for the full dataset.",
    )
    parser.add_argument("--manifest-only", action="store_true", help="Fetch manifest without downloading parquet.")
    args = parser.parse_args()
    if args.manifest_only:
        payload = fetch_manifest()
    else:
        payload = download_dataset(force=args.force, limit_shards=args.limit_shards)
    print(json.dumps(payload, indent=2, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())