#!/usr/bin/env python3
"""Reassemble an annotate.py run's flat per-position output back into
game-shaped containers: one record per source game, holding its
ordered list of annotated positions.

Why this exists: extract_positions.py's positions.jsonl (fen, score,
depth, pv, exact per row) is a flat list -- the row order follows each
shard's script order, which happens to group consecutive rows by game
(write_shard_script writes one game's positions fully before moving to
the next), but there's no game_id recorded, so that grouping is
unrecoverable from positions.jsonl alone. That matters for anything
that needs to avoid splitting one game's (highly correlated,
near-duplicate) positions across a train/holdout boundary.

This script recovers it deterministically, without re-running any
search: re-load the same PGN file with the same max_games (pure
parsing, no engine involved -- cheap), re-shard it the same way
annotate.py did (same n_shards, same deterministic round-robin), and
walk each game's known position count (positions_for_game) to slice
the already-extracted flat rows back into per-game chunks.

Output shape -- deliberately the same "a game containing an ordered
list of annotated positions" shape a self-play PGN with per-move
score/PV comments already has, so a future loader for self-play games
can produce the same games.jsonl schema and both sources mix freely at
training-set-assembly time:

    {"game_id": "...", "result": "1-0", "positions": [
        {"fen": ..., "depth": ..., "score": ..., "score_raw": ...,
         "pv": ..., "exact": ..., "exact_reason": ...},
        ...
    ]}

Usage:
    python3 build_games.py --rundir runs/sd10_twic_2M \
        --pgn ../pgn/twic_filtered.pgn --max-games 30000 --shards 16 \
        --out runs/sd10_twic_2M/games.jsonl
"""
import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from annotate import load_games, shard_games, positions_for_game


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--rundir", required=True,
                     help="directory holding positions.jsonl (extract_positions.py output)")
    ap.add_argument("--positions", default=None,
                     help="defaults to <rundir>/positions.jsonl")
    ap.add_argument("--pgn", nargs="+", required=True,
                     help="same PGN file(s) annotate.py was originally run against")
    ap.add_argument("--max-games", type=int, required=True,
                     help="same --max-games (or games actually used) as the original run")
    ap.add_argument("--shards", type=int, default=16,
                     help="same --shards as the original run")
    ap.add_argument("--prefilter-quiet", action="store_true",
                     help="must match whether the original annotate.py run "
                          "used --prefilter-quiet -- changes positions_for_game's "
                          "per-game ply count, which this script relies on to "
                          "re-slice the flat positions.jsonl")
    ap.add_argument("--out", required=True)
    args = ap.parse_args()

    rundir = Path(args.rundir)
    positions_path = Path(args.positions) if args.positions else rundir / "positions.jsonl"

    print(f"Loading all rows from {positions_path} ...")
    flat_rows = []
    with open(positions_path) as f:
        for line in f:
            flat_rows.append(json.loads(line))
    print(f"  {len(flat_rows)} rows")

    print(f"Re-loading {args.pgn} (max_games={args.max_games}) to recover game boundaries ...")
    games = load_games(args.pgn, args.max_games)
    print(f"  {len(games)} games")
    shards = shard_games(games, args.shards)

    # Per-shard row ranges: shard_NN's rows are a contiguous slice of
    # flat_rows in the same order extract_positions.py wrote them
    # (shard 00 first, then 01, ...), each shard's own length equal to
    # its total position count.
    shard_lengths = []
    for shard in shards:
        shard_lengths.append(sum(
            1 for g in shard
            for _ in positions_for_game(g, args.prefilter_quiet)))
    if sum(shard_lengths) != len(flat_rows):
        sys.exit(f"row count mismatch: recomputed shards expect "
                  f"{sum(shard_lengths)} total positions, "
                  f"positions.jsonl has {len(flat_rows)} -- wrong --pgn/"
                  f"--max-games/--shards for this run?")

    n_games_written = 0
    n_positions_written = 0
    with open(args.out, "w") as out:
        row_cursor = 0
        for shard_idx, shard in enumerate(shards):
            for game_idx, game in enumerate(shard):
                n_positions = sum(1 for _ in positions_for_game(game, args.prefilter_quiet))
                game_rows = flat_rows[row_cursor:row_cursor + n_positions]
                row_cursor += n_positions

                rec = {
                    "game_id": f"{shard_idx:02d}_{game_idx:05d}",
                    "result": game.headers.get("Result", "*"),
                    "positions": game_rows,
                }
                out.write(json.dumps(rec) + "\n")
                n_games_written += 1
                n_positions_written += len(game_rows)

        assert row_cursor == len(flat_rows)

    print(f"Wrote {args.out}: {n_games_written} games, "
          f"{n_positions_written} total positions")


if __name__ == "__main__":
    main()