from __future__ import annotations
import argparse
import json
import logging
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from riemann.data import prepare_all, DATA_DIR


def main() -> None:
    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s"
    )
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--only",
        nargs="*",
        default=None,
        help="only run these loader names, e.g. --only fable",
    )
    parser.add_argument(
        "--append",
        action="store_true",
        help="append to existing corpus files instead of overwriting them",
    )
    parser.add_argument(
        "--data-dir",
        default=str(DATA_DIR),
        help="where to write pretrain_corpus.txt and ntk_finetune.jsonl",
    )
    args = parser.parse_args()
    result = prepare_all(
        data_dir=Path(args.data_dir), only=args.only, append=args.append
    )
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()