#!/usr/bin/env python3
"""make_mem_samples.py — 为 test_reader_membudget.c 生成解码夹具（M22 / M27）。

用法：python3 make_mem_samples.py <输出目录>
依赖：Pillow。没有 Pillow 时测试自动 SKIP JPEG 相关断言（BMP/PNG 小夹具由 C 直接写）。

M27 追加「6 类典型页」与「回退路径」夹具，用来断言门控模型与真缩解后端同口径：
  典型（应全部走后端、不占位）：
    base2000.jpg   2000x3000 baseline
    prog2000.jpg   2000x3000 progressive
    restart2000.jpg 2000x3000 progressive + restart markers
    page1600.png   1600x2400 PNG
    page1600.bmp   1600x2400 BMP
    mp4.jpg        2320x1740（4MP）
    mp12.jpg       4000x3000（12MP）
  回退（后端拒绝、按旧 stb 门控占位）：
    cmyk.jpg       2000x3000 CMYK（4 分量）
    adam7.png      1600x2400 Adam7 交错
    mono1.bmp      1600x2400 1bpp
  M22 旧断言仍用：
    page.jpg / big4.jpg / prog.jpg（1000x1500 / 2000x2000 / 1000x1500 progressive）
"""
import os
import struct
import sys
import zlib

from PIL import Image, ImageDraw


def pattern(w, h, mode="RGB"):
    im = Image.new(mode, (w, h), (248, 247, 244) if mode == "RGB" else 200)
    d = ImageDraw.Draw(im)
    for k in range(0, min(w, h) - 200, 70):
        d.rectangle([20 + k, 20 + k, 160 + k, 120 + k], outline=(10, 10, 12), width=3)
    return im


def chunk(tag, data):
    return (struct.pack(">I", len(data)) + tag + data +
            struct.pack(">I", zlib.crc32(tag + data) & 0xFFFFFFFF))


def write_adam7(path, img):
    """手写 Adam7 交错 PNG（Pillow 不支持写交错）。RGB8。"""
    w, h = img.size
    px = img.load()
    raw = bytearray()
    passes = [(0, 0, 8, 8), (4, 0, 8, 8), (0, 4, 4, 8), (2, 0, 4, 4),
              (0, 2, 2, 4), (1, 0, 2, 2), (0, 1, 1, 2)]
    for (x0, y0, dx, dy) in passes:
        xs = list(range(x0, w, dx))
        ys = list(range(y0, h, dy))
        if not xs or not ys:
            continue
        for y in ys:
            raw.append(0)
            for x in xs:
                r, g, b = px[x, y][:3]
                raw += bytes((r, g, b))
    ihdr = struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 1)
    data = (b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", ihdr) +
            chunk(b"IDAT", zlib.compress(bytes(raw))) + chunk(b"IEND", b""))
    with open(path, "wb") as f:
        f.write(data)


def main():
    if len(sys.argv) != 2:
        print(__doc__)
        return 2
    out = sys.argv[1]
    os.makedirs(out, exist_ok=True)

    # ---- M22 旧断言夹具 ----
    pattern(1000, 1500).save(os.path.join(out, "page.jpg"), "JPEG", quality=60, subsampling=2)
    pattern(2000, 2000).save(os.path.join(out, "big4.jpg"), "JPEG", quality=60, subsampling=2)
    pattern(1000, 1500).save(os.path.join(out, "prog.jpg"), "JPEG", quality=60, subsampling=2,
                             progressive=True)

    # ---- M27 6 类典型页 ----
    big = pattern(2000, 3000)
    big.save(os.path.join(out, "base2000.jpg"), "JPEG", quality=60, subsampling=2)
    big.save(os.path.join(out, "prog2000.jpg"), "JPEG", quality=60, subsampling=2, progressive=True)
    big.save(os.path.join(out, "restart2000.jpg"), "JPEG", quality=60, subsampling=2,
             progressive=True, restart_marker_blocks=8)
    pattern(1600, 2400).save(os.path.join(out, "page1600.png"), "PNG")
    pattern(1600, 2400).save(os.path.join(out, "page1600.bmp"), "BMP")
    pattern(2320, 1740).save(os.path.join(out, "mp4.jpg"), "JPEG", quality=60, subsampling=2)
    pattern(4000, 3000).save(os.path.join(out, "mp12.jpg"), "JPEG", quality=60, subsampling=2)

    # ---- M27 回退路径 ----
    big.convert("CMYK").save(os.path.join(out, "cmyk.jpg"), "JPEG", quality=60)
    write_adam7(os.path.join(out, "adam7.png"), pattern(1600, 2400))
    pattern(1600, 2400).convert("1").save(os.path.join(out, "mono1.bmp"), "BMP")
    return 0


if __name__ == "__main__":
    sys.exit(main())