#!/usr/bin/env python3
"""生成 tuku zip 读取器的宿主测试样例（真实 zipfile + 定点字节修补）。

用法: python3 make_zip_samples.py <out_dir>

生成：
  basic.zip       5 条目：stored 文本 / deflate 大二进制 / 中文名 deflate /
                  目录项 / 0 字节 deflate
  comment.zip     单条目 + 归档注释（EOCD 注释长度路径）
  empty.zip       0 条目
  truncated.zip   basic.zip 砍半（EOCD 丢失）
  badmagic.zip    非 zip 字节
  badoffset.zip   首条目 local header 偏移改成 0xFFFFFF00（越界）
  badcrc.zip      首条目 CD 里 CRC 翻转（内容解出但校验失败）
  method99.zip    首条目压缩方式改成 99，CD 与局部头都改
  badlocal.zip    首条目局部头签名改成 "BAD!"
  bigsize.zip     deflate 条目 CD comp_size 改成 8 MiB（数据越界）
  zip64claim.zip  EOCD 总条目数改成 0xFFFF（zip64 标记）
  bigrand.zip     单条目 256 KiB 不可压 deflate（压缩流 > 64 KiB，覆盖分块 inflate）
"""
import os
import struct
import sys
import zipfile

PAGE2 = (b"MEOW" * 10 + bytes(range(256))) * 180          # 53280 字节，可压缩
HELLO = b"hello zip\n" * 3
CJK_NAME = "sub/dir/\u4e2d\u6587\u540d.txt"
CJK_DATA = "\u4e2d\u6587\u5185\u5bb9\n".encode("utf-8")
BIGRAND_LEN = 262144


def lcg_bytes(n):
    """与 test_zip.c 的 build_rand 同一个 LCG，用于 256 KiB 不可压样例。"""
    x = 123456789
    out = bytearray(n)
    for i in range(n):
        x = (1103515245 * x + 12345) & 0x7FFFFFFF
        out[i] = (x >> 16) & 0xFF
    return bytes(out)

EOCD_SIG = b"PK\x05\x06"
CD_SIG = b"PK\x01\x02"


def write_basic(path):
    with zipfile.ZipFile(path, "w") as z:
        z.writestr(zipfile.ZipInfo("hello.txt", (2024, 1, 1, 0, 0, 0)),
                   HELLO, compress_type=zipfile.ZIP_STORED)
        z.writestr(zipfile.ZipInfo("page2.png", (2024, 1, 1, 0, 0, 0)),
                   PAGE2, compress_type=zipfile.ZIP_DEFLATED)
        z.writestr(zipfile.ZipInfo(CJK_NAME, (2024, 1, 1, 0, 0, 0)),
                   CJK_DATA, compress_type=zipfile.ZIP_DEFLATED)
        z.writestr(zipfile.ZipInfo("onlydir/", (2024, 1, 1, 0, 0, 0)),
                   b"", compress_type=zipfile.ZIP_STORED)
        z.writestr(zipfile.ZipInfo("empty.bin", (2024, 1, 1, 0, 0, 0)),
                   b"", compress_type=zipfile.ZIP_DEFLATED)


def write_comment(path):
    with zipfile.ZipFile(path, "w") as z:
        z.writestr("note.txt", b"commented\n", compress_type=zipfile.ZIP_DEFLATED)
        z.comment = b"meow archive comment"


def load(path):
    with open(path, "rb") as f:
        return bytearray(f.read())


def save(path, data):
    with open(path, "wb") as f:
        f.write(bytes(data))


def eocd_info(data):
    i = data.rfind(EOCD_SIG)
    if i < 0:
        raise SystemExit("generator: no EOCD in fixture")
    total, cd_size, cd_off = struct.unpack_from("<HII", data, i + 10)
    return i, total, cd_size, cd_off


def cd_records(data):
    _, total, _cd_size, cd_off = eocd_info(data)
    out = []
    p = cd_off
    for _ in range(total):
        if bytes(data[p:p + 4]) != CD_SIG:
            raise SystemExit("generator: bad CD signature")
        nlen, elen, clen = struct.unpack_from("<HHH", data, p + 28)
        loff = struct.unpack_from("<I", data, p + 42)[0]
        name = bytes(data[p + 46:p + 46 + nlen]).decode("utf-8")
        out.append({"off": p, "name": name, "local": loff})
        p += 46 + nlen + elen + clen
    return out


def find_record(data, name):
    for r in cd_records(data):
        if r["name"] == name:
            return r
    raise SystemExit("generator: entry not found: " + name)


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

    basic = os.path.join(out, "basic.zip")
    write_basic(basic)
    write_comment(os.path.join(out, "comment.zip"))
    with zipfile.ZipFile(os.path.join(out, "empty.zip"), "w"):
        pass

    raw = load(basic)
    save(os.path.join(out, "truncated.zip"), raw[:len(raw) // 2])
    save(os.path.join(out, "badmagic.zip"), b"NOTAZIP!" * 100)

    d = load(basic)
    hello = find_record(d, "hello.txt")
    struct.pack_into("<I", d, hello["off"] + 42, 0xFFFFFF00)
    save(os.path.join(out, "badoffset.zip"), d)

    d = load(basic)
    hello = find_record(d, "hello.txt")
    crc = struct.unpack_from("<I", d, hello["off"] + 16)[0]
    struct.pack_into("<I", d, hello["off"] + 16, crc ^ 0xFFFFFFFF)
    save(os.path.join(out, "badcrc.zip"), d)

    d = load(basic)
    hello = find_record(d, "hello.txt")
    struct.pack_into("<H", d, hello["off"] + 10, 99)
    struct.pack_into("<H", d, hello["local"] + 8, 99)
    save(os.path.join(out, "method99.zip"), d)

    d = load(basic)
    hello = find_record(d, "hello.txt")
    d[hello["local"]:hello["local"] + 4] = b"BAD!"
    save(os.path.join(out, "badlocal.zip"), d)

    d = load(basic)
    page2 = find_record(d, "page2.png")
    struct.pack_into("<I", d, page2["off"] + 20, 0x00800000)
    save(os.path.join(out, "bigsize.zip"), d)

    d = load(basic)
    eocd, _total, _cd_size, _cd_off = eocd_info(d)
    struct.pack_into("<H", d, eocd + 10, 0xFFFF)
    save(os.path.join(out, "zip64claim.zip"), d)

    with zipfile.ZipFile(os.path.join(out, "bigrand.zip"), "w") as z:
        z.writestr("rand.bin", lcg_bytes(BIGRAND_LEN),
                   compress_type=zipfile.ZIP_DEFLATED)


if __name__ == "__main__":
    main()