#!/usr/bin/env python3
"""Split unfinished JSONL records into balanced temporary retry shards."""

from __future__ import annotations

import argparse
import json
from pathlib import Path


def record_id(value: dict) -> str:
    """Return the stable trajectory identifier used by audit manifests."""

    sample_id = value.get("trajectory_id") or value.get("sample_id")
    if not isinstance(sample_id, str) or not sample_id:
        raise ValueError("JSONL record has no trajectory_id/sample_id")
    return sample_id


def main() -> None:
    """Write unfinished input records round-robin across the requested shards."""

    parser = argparse.ArgumentParser()
    parser.add_argument("input", type=Path)
    parser.add_argument("completed", type=Path)
    parser.add_argument("output_dir", type=Path)
    parser.add_argument("--shards", type=int, default=3)
    args = parser.parse_args()
    if args.shards < 1:
        raise SystemExit("--shards must be positive")

    completed_ids: set[str] = set()
    with args.completed.open(encoding="utf-8") as handle:
        for line in handle:
            if line.strip():
                completed_ids.add(record_id(json.loads(line)))

    args.output_dir.mkdir(parents=True, exist_ok=True)
    handles = [
        (args.output_dir / f"part_{index}.input.jsonl").open("w", encoding="utf-8")
        for index in range(args.shards)
    ]
    counts = [0] * args.shards
    try:
        unfinished = 0
        with args.input.open(encoding="utf-8") as handle:
            for line in handle:
                if not line.strip():
                    continue
                value = json.loads(line)
                if record_id(value) in completed_ids:
                    continue
                shard = unfinished % args.shards
                handles[shard].write(line)
                counts[shard] += 1
                unfinished += 1
    finally:
        for handle in handles:
            handle.close()

    print(
        json.dumps(
            {
                "completed": len(completed_ids),
                "unfinished": sum(counts),
                "shard_counts": counts,
            },
            indent=2,
        )
    )


if __name__ == "__main__":
    main()