"""Build the two byte streams the search needs: one to search on, one held out.

Default source is text already on this machine (Python's own standard library), so no
download is needed. For a standard benchmark instead, download enwik8 yourself
(http://mattmahoney.net/dc/enwik8.zip, 36 MB), unzip it, and pass --from-file enwik8.

  python make_data.py                       # 1 MB search + 1 MB held out, local text
  python make_data.py --from-file enwik8    # same sizes, cut from enwik8
  python make_data.py --size-mb 4           # bigger streams (slower evaluations)
"""
from __future__ import annotations

import argparse
import random
import sysconfig
from pathlib import Path

MB = 1024 * 1024
MAX_SIZE_MB = 64


def gather_local_text(root: Path, wanted_bytes: int) -> bytes:
    files = sorted(path for path in root.rglob("*.py") if "test" not in path.parts)
    random.Random(0).shuffle(files)
    chunks, total = [], 0
    for path in files:
        try:
            chunk = path.read_bytes()
        except OSError:
            continue
        chunks.append(chunk)
        total += len(chunk)
        if total >= wanted_bytes:
            break
    return b"".join(chunks)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--from-file", help="cut both streams from this file")
    parser.add_argument("--size-mb", type=float, default=1.0, help="size of each stream")
    parser.add_argument("--out", default="data")
    args = parser.parse_args()

    if not 0 < args.size_mb <= MAX_SIZE_MB:
        parser.error(f"--size-mb must be between 0 and {MAX_SIZE_MB}")
    size = int(args.size_mb * MB)

    if args.from_file:
        raw = Path(args.from_file).read_bytes()[: 2 * size]
    else:
        raw = gather_local_text(Path(sysconfig.get_paths()["stdlib"]), 2 * size)
    if len(raw) < 2 * size:
        parser.error(f"source has only {len(raw)} bytes; need {2 * size}")

    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    (out / "search.bin").write_bytes(raw[:size])
    (out / "holdout.bin").write_bytes(raw[size: 2 * size])
    print(f"wrote {out / 'search.bin'} and {out / 'holdout.bin'} ({size} bytes each)")
    return 0


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