"""Save a small JSON document by replacement. Python 3.10+, single writer."""
import argparse
import json
import os
from pathlib import Path
import tempfile


def atomic_write_json(path, data):
    """Replace a file in a trusted local directory; parent must already exist."""
    target = Path(path).absolute()
    # Serialize first: invalid data must not touch the destination.
    text = json.dumps(data, ensure_ascii=False, allow_nan=False, indent=2) + "\n"
    temporary = None
    try:
        with tempfile.NamedTemporaryFile(
            mode="w", encoding="utf-8", newline="\n",
            dir=target.parent, prefix=f".{target.name}.",
            suffix=".tmp", delete=False,
        ) as handle:
            temporary = Path(handle.name)
            handle.write(text)
            handle.flush()
            os.fsync(handle.fileno())
        # Close before replacement, including on Windows.
        os.replace(temporary, target)
    finally:
        if temporary is not None:
            temporary.unlink(missing_ok=True)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("output", type=Path, help="demo JSON path; replaced if present")
    args = parser.parse_args()
    old = {"completed_batches": 1, "status": "ready"}
    atomic_write_json(args.output, old)
    before = args.output.read_bytes()
    try:
        atomic_write_json(args.output, {"loss": float("nan")})
    except ValueError:
        assert args.output.read_bytes() == before
        print("Rejected NaN; previous file unchanged")
    new = {"completed_batches": 2, "status": "ready"}
    atomic_write_json(args.output, new)
    assert json.loads(args.output.read_text(encoding="utf-8")) == new
    print("OK: completed_batches=2")


if __name__ == "__main__":
    main()