# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Make memcmp sort your numbers correctly.

Encode signed integers and floats so that comparing the bytes (unsigned,
byte by byte, like memcmp) gives numeric order: flip the sign bit of an
integer; for a float, invert all bits if negative, else set the sign bit.
Then composite string keys, with and without an escape and a terminator.
"""
import argparse, json, random, struct

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--values", type=int, default=10_000, help="Values per test.")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
random.seed(args.seed)
MASK = (1 << 64) - 1
def enc_i64(n): return struct.pack(">Q", (n & MASK) ^ (1 << 63))
def enc_f64(x):
    b = struct.unpack(">Q", struct.pack(">d", x))[0]
    b = (~b & MASK) if (b >> 63) else (b | (1 << 63))
    return struct.pack(">Q", b)
ints = [random.randint(-10**12, 10**12) for _ in range(args.values)]
floats = [random.uniform(-1e6, 1e6) for _ in range(args.values)]
naive = lambda n: struct.pack(">q", n)
naive_f = lambda x: struct.pack(">d", x)

def concat(parts): return b"".join(s.encode() for s in parts)
def escaped(parts): return b"".join(s.encode().replace(b"\x00", b"\x00\xff") + b"\x00\x01" for s in parts)
keys = [("ab", "c"), ("a", "bc"), ("a", "b"), ("ab", ""), ("a", "")]
rows = [
    {"test": "integers, plain two's complement", "sorted_correctly": sorted(ints, key=naive) == sorted(ints)},
    {"test": "integers, sign bit flipped", "sorted_correctly": sorted(ints, key=enc_i64) == sorted(ints)},
    {"test": "floats, plain IEEE bytes", "sorted_correctly": sorted(floats, key=naive_f) == sorted(floats)},
    {"test": "floats, sign rule", "sorted_correctly": sorted(floats, key=enc_f64) == sorted(floats)},
    {"test": "composite keys, plain concatenation unique", "sorted_correctly": len({concat(k) for k in keys}) == len(keys)},
    {"test": "composite keys, escaped + terminated unique", "sorted_correctly": len({escaped(k) for k in keys}) == len(keys)},
    {"test": "composite keys, escaped + terminated order", "sorted_correctly": sorted(keys, key=escaped) == sorted(keys)},
]
for r in rows:
    r["result"] = "yes" if r.pop("sorted_correctly") else "no"
    print(f"{r['test']:<46} {r['result']}")
print("('ab','c') and ('a','bc') concatenated:", concat(("ab", "c")), concat(("a", "bc")), "-> the same bytes")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"tests": rows}}, f)