#!/usr/bin/env python3
"""Cross-check the kernel QR encoder (tools/qrcheck) against segno.

Strategy: segno's own bitstream differs from the kernel's in the pad bits
(segno emits a full extra zero byte when the stream is already byte-aligned),
so instead of comparing against `segno.make()` directly we feed the kernel's
own interleaved codeword stream through segno's placement primitives
(make_matrix / add_finder_patterns / add_alignment_patterns / add_codewords /
apply_mask / add_format_info / add_version_info) and compare matrices
bit-for-bit.  For auto-mask cases the chosen mask is validated against
segno's `mask_scores` penalty evaluation on the same (kernel-stream) matrix.
"""
import importlib
import subprocess
import sys

sys.path.insert(0, "/tmp/segno_pkg")
segno = importlib.import_module("segno")  # vendored reference at /tmp/segno_pkg
E = importlib.import_module("segno.encoder")
C = importlib.import_module("segno.consts")

BIN = "tools/qrcheck/target/debug/qrcheck"
ROOT = "/home/salmon/Documents/rust/rsdos"


def run(args):
    out = subprocess.run(args, capture_output=True, text=True, cwd=ROOT)
    if out.returncode != 0:
        print("RUST FAIL:", out.stderr.strip())
        sys.exit(1)
    return out.stdout


def rust_matrix(text, version, mask):
    args = [BIN, text]
    if version:
        args.append(str(version))
    if mask is not None:
        args.append(str(mask))
    lines = run(args).splitlines()
    return lines[0], lines[1:]


def rust_raw(text, version):
    """Kernel interleaved codeword stream (data+EC) as bytes."""
    hexs = run([BIN, "--raw", text, str(version)]).strip()
    return bytes.fromhex(hexs)


def _to_bits(my_cw):
    try:
        return [int(b) for byte in my_cw for b in f"{byte:08b}"]
    except (ValueError, TypeError) as exc:
        print(f"bad codeword stream {my_cw!r}: {exc}")
        sys.exit(1)


def build_matrix(version, my_cw, mask, with_info):
    """Expected matrix from the kernel's codeword stream via segno primitives.

    `with_info=True` also writes format + version info (final symbol);
    `with_info=False` leaves them light (as during mask evaluation).
    """
    n = E.calc_matrix_size(version)
    bits = _to_bits(my_cw)
    m = E.make_matrix(n, n)
    E.add_finder_patterns(m, n, n)
    E.add_alignment_patterns(m, n, n)
    E.add_codewords(m, bits, version)
    fm = E.make_matrix(n, n)
    E.add_finder_patterns(fm, n, n)
    E.add_alignment_patterns(fm, n, n)
    fm[-8][8] = 0x1
    E.apply_mask(m, E.get_data_mask_functions(False)[mask], n, n,
                 lambda i, j: fm[i][j] > 0x1)
    if with_info:
        E.add_format_info(m, version, C.ERROR_LEVEL_L, mask)
        E.add_version_info(m, version)
    # Normalize leftover (remainder) cells: 0x2/0x3 -> 0/1.
    return [[v & 0x1 if v in (2, 3) else v for v in row] for row in m]


def render(rows):
    return ["".join("#" if c else "." for c in row) for row in rows]


def compare_forced(text, version, mask):
    hdr, rows = rust_matrix(text, version, mask)
    my_cw = rust_raw(text, version)
    exp = build_matrix(version, my_cw, mask, with_info=True)
    srows = render(exp)
    tag = f"v{version} m{mask} {text[:28]!r}"
    if len(rows) != len(srows):
        print(f"FAIL {tag}: rust {len(rows)}x vs expected {len(srows)}x")
        return False
    for y, (r, s) in enumerate(zip(rows, srows)):
        if r != s:
            print(f"FAIL {tag}: row {y} differs (rust v{mask})")
            for dy in range(max(0, y - 2), min(len(rows), y + 3)):
                print(f"   rust row{dy}: {rows[dy]}")
                print(f"   exp  row{dy}: {srows[dy]}")
            return False
    print(f"OK   {tag} ({len(rows)}x{len(rows)}, {hdr})")
    return True


def segno_best_mask(version, my_cw):
    """Mask segno's penalty scoring would choose for the kernel stream."""
    n = E.calc_matrix_size(version)
    bits = _to_bits(my_cw)
    scores = []
    for mask in range(8):
        m = E.make_matrix(n, n)
        E.add_finder_patterns(m, n, n)
        E.add_alignment_patterns(m, n, n)
        E.add_codewords(m, bits, version)
        fm = E.make_matrix(n, n)
        E.add_finder_patterns(fm, n, n)
        E.add_alignment_patterns(fm, n, n)
        fm[-8][8] = 0x1
        E.apply_mask(m, E.get_data_mask_functions(False)[mask], n, n,
                     lambda i, j: fm[i][j] > 0x1)
        scores.append((sum(E.mask_scores(m, n, n)), mask))
    return min(scores)  # (score, mask); first-lowest wins like segno


def _parse_hdr(hdr):
    try:
        v = int(hdr.split()[0][1:])
        m = int(hdr.split()[1][1:])
    except (ValueError, IndexError) as exc:
        print(f"bad header {hdr!r}: {exc}")
        sys.exit(1)
    return v, m


def compare_auto(text, version=None):
    args = [BIN, text]
    if version:
        args.append(str(version))
    lines = run(args).splitlines()
    hdr = lines[0]  # "v<N> m<M>"
    rows = lines[1:]
    v, m = _parse_hdr(hdr)
    my_cw = rust_raw(text, v)
    score, s_mask = segno_best_mask(v, my_cw)
    tag = f"v{v} auto-mask {text[:28]!r}"
    if m != s_mask:
        print(f"FAIL {tag}: rust chose m{m}, segno penalty chose m{s_mask} "
              f"(score {score})")
        return False
    # Full final-matrix check with the chosen mask.
    exp = build_matrix(v, my_cw, m, with_info=True)
    srows = render(exp)
    if len(rows) != len(srows):
        print(f"FAIL {tag}: size mismatch")
        return False
    for y, (r, s) in enumerate(zip(rows, srows)):
        if r != s:
            print(f"FAIL {tag}: row {y} differs")
            for dy in range(max(0, y - 2), min(len(rows), y + 3)):
                print(f"   rust row{dy}: {rows[dy]}")
                print(f"   exp  row{dy}: {srows[dy]}")
            return False
    print(f"OK   {tag} (m{m} matches segno, {len(rows)}x{len(rows)})")
    return True


def main():
    ok = True
    texts = {
        1: "A",
        2: "hello",
        3: "the quick brown fox",
        4: "0123456789ABCDEF",
        5: "byte mode QR encoder cross-check",
        6: "x" * 100,
        7: "hello world, this is a slightly longer message for v7",
        10: "x" * 250,
        14: "x" * 400,
        20: "x" * 800,
        25: "x" * 1200,
        40: "x" * 2900,
    }
    for v in (1, 2, 3, 4, 5, 6, 7, 10, 14, 20, 25, 40):
        for m in range(8):
            ok &= compare_forced(texts[v], v, m)
    crash = (
        "!!! heap corruption in dealloc\n"
        "ptr=0x404443e0 block=0x1980000 heap=0x1980000 rsp=0x1fffe0\n"
        "mem: x=0x41414141 x=0x41414141 x=0x41414141 x=0x41414141\n"
        "mem: x=0x41414141 x=0x41414141 x=0x41414141 x=0x41414141\n"
        "stk: x=0x0 x=0x0 x=0x0 x=0x0 x=0x0 x=0x0 x=0x0 x=0x0\n"
        "hdr: b=0x1980000 s=0x20 f=0 b16=0x1980000 s16=0x20 f16=0\n"
        "RSDOS fault: #GP at rip=0x12345678 cr2=0x0 task=main\n"
    )
    ok &= compare_auto(crash)
    ok &= compare_auto("hello world", 1)
    ok &= compare_auto("byte mode QR encoder cross-check", 5)
    ok &= compare_auto("hello world, this is a slightly longer message for v7", 7)
    ok &= compare_auto("x" * 400, 14)
    ok &= compare_auto("x" * 2900, 40)
    print("ALL OK" if ok else "FAILURES")
    sys.exit(0 if ok else 1)


if __name__ == "__main__":
    main()