#!/usr/bin/env python3
"""Parallelize mine_king_mates.py across N shards by game-boundary byte
offset, since a full 642k-game scan is CPU-bound (python-chess replay +
SAN parsing per game) and this box has 24 cores. Finds "[Event " line
starts via a fast raw byte scan (not a full PGN parse) to get exact game
boundaries, splits into N even-sized shards by game count, launches one
mine_king_mates.py subprocess per shard with --byte-start/--byte-end, then
concatenates the resulting CSVs (one header, rows from all shards).
"""
import argparse
import csv
import os
import subprocess
import sys
import time


def find_game_offsets(pgn_file):
    offsets = []
    pos = 0
    with open(pgn_file, "rb") as f:
        for line in f:
            if line.startswith(b"[Event "):
                offsets.append(pos)
            pos += len(line)
    return offsets, pos  # pos ends up as total file size


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("pgn_file")
    ap.add_argument("out_dir")
    ap.add_argument("--num-shards", type=int, default=8)
    ap.add_argument("--extra-args", default="",
                     help="Extra args passed through to mine_king_mates.py, "
                          "e.g. '--control-ratio 1.0'")
    args = ap.parse_args()

    os.makedirs(args.out_dir, exist_ok=True)
    script_dir = os.path.dirname(os.path.abspath(__file__))
    mine_script = os.path.join(script_dir, "mine_king_mates.py")

    print("Scanning for game boundaries...", file=sys.stderr)
    t0 = time.time()
    offsets, file_size = find_game_offsets(args.pgn_file)
    print(f"Found {len(offsets)} games, {file_size} bytes, "
          f"in {time.time() - t0:.1f}s", file=sys.stderr)

    n = args.num_shards
    games_per_shard = len(offsets) // n
    shard_bounds = []
    for i in range(n):
        start_idx = i * games_per_shard
        end_idx = (i + 1) * games_per_shard if i < n - 1 else len(offsets)
        byte_start = offsets[start_idx]
        byte_end = offsets[end_idx] if end_idx < len(offsets) else file_size
        shard_bounds.append((byte_start, byte_end))

    procs = []
    for i, (byte_start, byte_end) in enumerate(shard_bounds):
        out_csv = os.path.join(args.out_dir, f"shard_{i}.csv")
        log_file = os.path.join(args.out_dir, f"shard_{i}.log")
        cmd = [
            sys.executable, mine_script, args.pgn_file, out_csv,
            "--byte-start", str(byte_start),
            "--byte-end", str(byte_end),
        ] + args.extra_args.split()
        print(f"Launching shard {i}: bytes [{byte_start}, {byte_end}) -> {out_csv}",
              file=sys.stderr)
        with open(log_file, "w") as logf:
            p = subprocess.Popen(cmd, stdout=logf, stderr=subprocess.STDOUT)
        procs.append((i, p, out_csv))

    failed = []
    for i, p, out_csv in procs:
        rc = p.wait()
        print(f"Shard {i} exited with code {rc}", file=sys.stderr)
        if rc != 0:
            failed.append(i)

    if failed:
        print(f"FAILED shards: {failed} -- not merging, check logs in {args.out_dir}",
              file=sys.stderr)
        sys.exit(1)

    merged_path = os.path.join(args.out_dir, "merged.csv")
    header = None
    total_rows = 0
    with open(merged_path, "w", newline="") as out_f:
        writer = csv.writer(out_f)
        for i, _, out_csv in procs:
            with open(out_csv, newline="") as in_f:
                reader = csv.reader(in_f)
                this_header = next(reader)
                if header is None:
                    header = this_header
                    writer.writerow(header)
                elif this_header != header:
                    print(f"WARNING: shard {i} header mismatch: {this_header}",
                          file=sys.stderr)
                for row in reader:
                    writer.writerow(row)
                    total_rows += 1

    print(f"Merged {total_rows} rows from {n} shards into {merged_path}",
          file=sys.stderr)


if __name__ == "__main__":
    main()