#!/usr/bin/env python3
"""生成 tuku 档案册（CBZ）集成测试夹具：真实 zipfile 产物 + 定点字节修补。

用法: python3 make_cbz_samples.py <out_dir>

产出：
  <out>/library/            书库夹具（文件夹册 + 档案册混排）
      bookA/001.png 002.png        文件夹册（2 页，PNG）
      vol2.cbz                     3 页：PNG(41x31) / PNG(40x30) / 中文名 PNG(12x34)
                                   + 非图片、目录项、0 字节条目（都应被过滤）
      Vol10.cbz                    1 页 BMP(7x9)，stored
      box/box1.cbz box2.cbz        只含档案册的子目录（书库行应写 "2 books"）
      notabook.zip                 只有文本 → 不是书
      emptypages.cbz               目录项 + 文本 → 不是书
      broken.cbz                   非 zip 字节 → 不是书
      truncated.cbz                vol2 砍半（EOCD 丢失） → 不是书
      .hidden.cbz                  隐藏名 → 不参与
  <out>/features/
      pathsafe.cbz                 安全/不安全条目混排（只应收 2 个安全图片页）
      cjk.cbz                      中文名自然序（第1页 < 第2页 < 第10页）
      big.cbz                      80 页 8x8 PNG（解压缓存 LRU 上限用）
      bomb.cbz                     CD 声明单页 96MiB（应被过滤，防 zip 炸弹）
  <out>/manifest.txt        每档案每条目一行：<档案相对路径>\t<条目名>\t<crc32>\t<size>

图片内容都是结构合法的 PNG/BMP/JPEG 头（PNG 是真图），C 测试用假解码器，
只校验解压字节与 manifest 的 CRC/尺寸一致、以及 meta 探测出的原始尺寸。
"""
import os
import struct
import sys
import zipfile
import zlib

DATE = (2024, 1, 1, 0, 0, 0)


# ---------------------------------------------------------------- 图片构造
def png_bytes(w, h, seed=0):
    raw = bytearray()
    for y in range(h):
        raw.append(0)                                   # filter: none
        for x in range(w):
            raw += bytes(((x * 7 + y * 13 + seed) & 0xFF,
                          (x * 3 + y * 5 + seed) & 0xFF,
                          (x + y + seed) & 0xFF))

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

    ihdr = struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0)   # 8-bit RGB
    return (b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", ihdr) +
            chunk(b"IDAT", zlib.compress(bytes(raw), 9)) + chunk(b"IEND", b""))


def bmp_bytes(w, h):
    stride = (w * 3 + 3) & ~3
    px = bytes(((i * 11) & 0xFF) for i in range(stride * h))
    hdr = b"BM" + struct.pack("<IHHI", 14 + 40 + len(px), 0, 0, 54)
    ih = struct.pack("<IiiHHIIiiII", 40, w, h, 1, 24, 0, len(px), 2835, 2835, 0, 0)
    return hdr + ih + px


def jpeg_header_bytes(w, h):
    """SOI + APP0 + DQT + SOF0 + EOI：头段合法（探测只看 SOF0），不做熵编码。"""
    app0 = (b"\xFF\xE0" + struct.pack(">H", 16) +
            b"JFIF\x00\x01\x01\x00\x00\x01\x00\x01\x00\x00")
    dqt = b"\xFF\xDB" + struct.pack(">H", 67) + b"\x00" + bytes(64)
    sof = b"\xFF\xC0" + struct.pack(">HBHHB", 11, 8, h, w, 1) + b"\x01\x11\x00"
    return b"\xFF\xD8" + app0 + dqt + sof + b"\xFF\xD9"


# ---------------------------------------------------------------- zip 构造
def write_zip(path, entries):
    """entries: [(name, data, zipfile 压缩方式)]"""
    with zipfile.ZipFile(path, "w") as z:
        for name, data, method in entries:
            zi = zipfile.ZipInfo(name, DATE)
            zi.compress_type = method
            z.writestr(zi, data)


def find_cd_record(data, name):
    eocd = data.rfind(b"PK\x05\x06")
    if eocd < 0:
        raise SystemExit("generator: no EOCD")
    total, _cd_size, cd_off = struct.unpack_from("<HII", data, eocd + 10)
    p = cd_off
    for _ in range(total):
        if bytes(data[p:p + 4]) != b"PK\x01\x02":
            raise SystemExit("generator: bad CD signature")
        nlen, elen, clen = struct.unpack_from("<HHH", data, p + 28)
        got = bytes(data[p + 46:p + 46 + nlen]).decode("utf-8")
        if got == name:
            return p
        p += 46 + nlen + elen + clen
    raise SystemExit("generator: entry not found: " + name)


# ---------------------------------------------------------------- 夹具
def build_library(ro):
    lib = os.path.join(ro, "library")
    os.makedirs(lib, exist_ok=True)

    os.makedirs(os.path.join(lib, "bookA"), exist_ok=True)
    with open(os.path.join(lib, "bookA", "001.png"), "wb") as f:
        f.write(png_bytes(30, 20, 1))
    with open(os.path.join(lib, "bookA", "002.png"), "wb") as f:
        f.write(png_bytes(30, 21, 2))

    write_zip(os.path.join(lib, "vol2.cbz"), [
        ("ch01/p10.png", png_bytes(40, 30, 3), zipfile.ZIP_DEFLATED),
        ("readme.txt", b"not a page\n", zipfile.ZIP_STORED),
        ("ch01/p2.png", png_bytes(41, 31, 4), zipfile.ZIP_DEFLATED),
        ("onlydir/", b"", zipfile.ZIP_STORED),
        ("\u4e2d\u6587\u540d.png", png_bytes(12, 34, 5), zipfile.ZIP_DEFLATED),
        ("empty.bin", b"", zipfile.ZIP_DEFLATED),
    ])

    write_zip(os.path.join(lib, "Vol10.cbz"), [
        ("cover.bmp", bmp_bytes(7, 9), zipfile.ZIP_STORED),
    ])

    os.makedirs(os.path.join(lib, "box"), exist_ok=True)
    write_zip(os.path.join(lib, "box", "box1.cbz"), [
        ("a.png", png_bytes(8, 8, 6), zipfile.ZIP_DEFLATED),
        ("b.png", png_bytes(8, 9, 7), zipfile.ZIP_DEFLATED),
    ])
    write_zip(os.path.join(lib, "box", "box2.cbz"), [
        ("c.jpg", jpeg_header_bytes(12, 34), zipfile.ZIP_DEFLATED),
    ])

    write_zip(os.path.join(lib, "notabook.zip"), [
        ("readme.txt", b"nothing to read\n", zipfile.ZIP_STORED),
    ])
    write_zip(os.path.join(lib, "emptypages.cbz"), [
        ("onlydir/", b"", zipfile.ZIP_STORED),
        ("readme.txt", b"x\n", zipfile.ZIP_STORED),
    ])
    with open(os.path.join(lib, "broken.cbz"), "wb") as f:
        f.write(b"NOTAZIP!" * 100)

    raw = open(os.path.join(lib, "vol2.cbz"), "rb").read()
    with open(os.path.join(lib, "truncated.cbz"), "wb") as f:
        f.write(raw[:len(raw) // 2])

    write_zip(os.path.join(lib, ".hidden.cbz"), [
        ("h.png", png_bytes(8, 8, 8), zipfile.ZIP_DEFLATED),
    ])


def build_features(ro):
    feat = os.path.join(ro, "features")
    os.makedirs(feat, exist_ok=True)

    write_zip(os.path.join(feat, "pathsafe.cbz"), [
        ("ok1.png", png_bytes(10, 10, 9), zipfile.ZIP_DEFLATED),
        ("../evil.png", png_bytes(10, 10, 10), zipfile.ZIP_DEFLATED),
        ("/abs.png", png_bytes(10, 10, 11), zipfile.ZIP_DEFLATED),
        ("a/../../b.png", png_bytes(10, 10, 12), zipfile.ZIP_DEFLATED),
        ("sub/ok2.jpg", jpeg_header_bytes(13, 17), zipfile.ZIP_DEFLATED),
        ("sub/../ok3.png", png_bytes(10, 10, 13), zipfile.ZIP_DEFLATED),
        (".hidden/x.png", png_bytes(10, 10, 14), zipfile.ZIP_DEFLATED),
        ("._mac.png", png_bytes(10, 10, 15), zipfile.ZIP_DEFLATED),
        ("sub\\win.png", png_bytes(10, 10, 16), zipfile.ZIP_DEFLATED),
        ("text.txt", b"nope\n", zipfile.ZIP_STORED),
        ("dir/", b"", zipfile.ZIP_STORED),
    ])

    write_zip(os.path.join(feat, "cjk.cbz"), [
        ("\u7b2c2\u9875.png", png_bytes(9, 9, 17), zipfile.ZIP_DEFLATED),
        ("\u7b2c10\u9875.png", png_bytes(9, 9, 18), zipfile.ZIP_DEFLATED),
        ("\u7b2c1\u9875.png", png_bytes(9, 9, 19), zipfile.ZIP_DEFLATED),
    ])

    entries = []
    for i in range(80):
        entries.append(("p%03d.png" % i, png_bytes(8, 8, 20 + i), zipfile.ZIP_DEFLATED))
    write_zip(os.path.join(feat, "big.cbz"), entries)

    bomb = os.path.join(feat, "bomb.cbz")
    write_zip(bomb, [("huge.png", png_bytes(10, 10, 30), zipfile.ZIP_STORED)])
    data = bytearray(open(bomb, "rb").read())
    rec = find_cd_record(data, "huge.png")
    struct.pack_into("<I", data, rec + 24, 0x06000000)     # 96MiB > 64MiB 上限
    with open(bomb, "wb") as f:
        f.write(bytes(data))

    raw = open(os.path.join(feat, "pathsafe.cbz"), "rb").read()
    with open(os.path.join(feat, "truncated.cbz"), "wb") as f:
        f.write(raw[:len(raw) // 2])


def write_manifest(ro):
    """每档案每条目：<相对路径>\t<条目名>\t<crc32>\t<size>（Tab 分隔，UTF-8 原名）。"""
    lines = []
    for base, _dirs, files in os.walk(ro):
        for name in sorted(files):
            if not name.lower().endswith((".cbz", ".zip")):
                continue
            path = os.path.join(base, name)
            rel = os.path.relpath(path, ro)
            try:
                with zipfile.ZipFile(path) as z:
                    for info in z.infolist():
                        if info.is_dir():
                            continue
                        payload = z.read(info)
                        lines.append("%s\t%s\t%08x\t%d" % (
                            rel, info.filename, zlib.crc32(payload) & 0xFFFFFFFF,
                            len(payload)))
            except zipfile.BadZipFile:
                continue
    with open(os.path.join(ro, "manifest.txt"), "w", encoding="utf-8") as f:
        f.write("\n".join(lines) + "\n")


def main():
    if len(sys.argv) != 2:
        raise SystemExit("usage: make_cbz_samples.py <out_dir>")
    ro = sys.argv[1]
    os.makedirs(ro, exist_ok=True)
    build_library(ro)
    build_features(ro)
    write_manifest(ro)


if __name__ == "__main__":
    main()