#!/usr/bin/env python3
"""Pick the mst-test-suite cases to vendor into hubble-commit.

Selection is deterministic and biased:

- every case involving the empty tree (subset 000) bc edge cases are fun
- the cases with the most structural churn (nodes created plus deleted)
- the cases with the most record ops, where op ordering matters most
- an even sample of the rest

usage: scripts/select-mst-cases.py <suite-checkout> <fixtures-dir> [count]
"""

import json
import shutil
import sys
from pathlib import Path


def load(path):
    with path.open() as f:
        return json.load(f)


def main():
    if len(sys.argv) < 3:
        sys.exit(__doc__)
    suite = Path(sys.argv[1])
    out = Path(sys.argv[2])
    want = int(sys.argv[3]) if len(sys.argv) > 3 else 300

    cases = sorted((suite / "tests").rglob("*.json"))
    if not cases:
        sys.exit(f"no cases under {suite / 'tests'}")

    scored = []
    for path in cases:
        case = load(path)
        results = case["results"]
        churn = len(results["created_nodes"]) + len(results["deleted_nodes"])
        ops = len(results["record_ops"])
        empty = "_000" in path.stem  # either side is the empty tree
        scored.append((path, churn, ops, empty))

    # the empty-tree cases are the most valuable single group, but there are
    # 255 of them, so take an even stride rather than all of them
    empties = [p for p, _, _, empty in scored if empty]
    stride = max(1, len(empties) // max(1, want // 2))
    picked = set(empties[::stride])
    for key in (lambda s: -s[1], lambda s: -s[2]):
        for path, *_ in sorted(scored, key=key)[: want // 4]:
            picked.add(path)
    remaining = [p for p, *_ in scored if p not in picked]
    if len(picked) < want and remaining:
        stride = max(1, len(remaining) // max(1, want - len(picked)))
        picked.update(remaining[::stride][: want - len(picked)])

    if out.exists():
        shutil.rmtree(out)
    for path in sorted(picked):
        dest = out / path.relative_to(suite)
        dest.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(path, dest)

    cars = out / "cars"
    shutil.copytree(suite / "cars", cars)

    total = sum(f.stat().st_size for f in out.rglob("*") if f.is_file())
    print(f"{len(picked)} cases, {total / 1024:.0f} kB into {out}")


if __name__ == "__main__":
    main()