#!/usr/bin/env python3
# gen_samples.py — 本地造 R2 阅读器测试样本（EPUB/FB2/CBZ/MOBI7）。不联网、不拉真实书。
# 输出到 /tmp/rd_samples/。MOBI7 用手写 PalmDoc(LZ77) 压缩 + PDB 容器（对齐 libmobi mobi_decompress_lz77）。
import os, struct, zlib, base64, json

OUT = "/tmp/rd_samples"
os.makedirs(OUT, exist_ok=True)

def png(w, h, rgb):
    """最小 valid PNG（zlib deflate，RGB 8bit）。"""
    raw = b""
    for _ in range(h):
        raw += b"\x00" + bytes(rgb) * w
    def chunk(t, d):
        return struct.pack(">I", len(d)) + t + d + struct.pack(">I", zlib.crc32(t + d) & 0xffffffff)
    ihdr = struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0)
    return (b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", ihdr)
            + chunk(b"IDAT", zlib.compress(raw)) + chunk(b"IEND", b""))

# ---------- EPUB ----------
def make_epub():
    from zipfile import ZipFile, ZIP_DEFLATED
    cover = png(20, 20, (200, 60, 60))
    ch1 = """<?xml version="1.0" encoding="UTF-8"?>
<html><head><title>第一章</title></head><body>
<h1>第一章 测试</h1>
<p>这是一段<em>带格式</em>的正文，包含实体 &amp; 和标点符号。English words mixed here.</p>
<p>第二段：测试坏标签容错 <b>未闭合 <i>粗斜体</p>
<p><img src="../images/cover.png" alt="cover"/></p>
<p>&#20197;&#19978;是数值实体引用的中文。</p>
</body></html>"""
    ch2 = """<?xml version="1.0"?>
<html><head><title>第二章</title></head><body>
<h1>第二章 测试</h1>
<p>第二章正文，用于验证多章节 spine 顺序。</p>
<p>结尾</p>
</body></html>"""
    container = """<?xml version="1.0"?>
<container version="1.0" xmlns="urn:oasis:names:tc:opendocument:xmlns:container">
 <rootfiles><rootfile full-path="OEBPS/content.opf" media-type="application/oebps-package+xml"/></rootfiles>
</container>"""
    opf = """<?xml version="1.0"?>
<package xmlns="http://www.idpf.org/2007/opf" version="3.0" unique-identifier="bookid">
 <metadata xmlns:dc="http://purl.org/dc/elements/1.1/">
  <dc:title>测试书</dc:title><dc:creator>MeowOS</dc:creator>
 </metadata>
 <manifest>
  <item id="ncx" href="toc.ncx" media-type="application/x-dtbncx+xml"/>
  <item id="ch1" href="text/ch1.xhtml" media-type="application/xhtml+xml"/>
  <item id="ch2" href="text/ch2.xhtml" media-type="application/xhtml+xml"/>
  <item id="cover-img" href="images/cover.png" media-type="image/png"/>
 </manifest>
 <spine toc="ncx"><itemref idref="ch1"/><itemref idref="ch2"/></spine>
</package>"""
    ncx = """<?xml version="1.0"?>
<ncx xmlns="http://www.daisy.org/z3986/2005/ncx/" version="2005-1">
 <navMap>
  <navPoint id="n1"><navLabel><text>第一章</text></navLabel><content src="text/ch1.xhtml"/></navPoint>
  <navPoint id="n2"><navLabel><text>第二章</text></navLabel><content src="text/ch2.xhtml"/></navPoint>
 </navMap>
</ncx>"""
    with ZipFile(os.path.join(OUT, "sample.epub"), "w", ZIP_DEFLATED) as z:
        z.writestr("mimetype", "application/epub+zip")
        z.writestr("META-INF/container.xml", container)
        z.writestr("OEBPS/content.opf", opf)
        z.writestr("OEBPS/toc.ncx", ncx)
        z.writestr("OEBPS/text/ch1.xhtml", ch1.encode("utf-8"))
        z.writestr("OEBPS/text/ch2.xhtml", ch2.encode("utf-8"))
        z.writestr("OEBPS/images/cover.png", cover)

# ---------- FB2 ----------
def make_fb2():
    img_data = base64.b64encode(png(12, 12, (30, 120, 200))).decode()
    fb2 = f"""<?xml version="1.0" encoding="UTF-8"?>
<FictionBook xmlns="http://www.gribuser.ru/xml/fictionbook/2.0">
  <description><title-info><book-title>FB2测试</book-title><author><first-name>Meo</first-name><last-name>OS</last-name></author></title-info></description>
  <body>
   <section><title><p>第一章</p></title>
    <p>这是 FB2 段落，含实体 &amp; 与 &#27491;&#25991;。</p>
    <p><image l:href="#img1"/></p>
   </section>
   <section><title><p>第二章</p></title>
    <p>第二章内容。</p>
   </section>
  </body>
  <binary id="img1" content-type="image/png">{img_data}</binary>
</FictionBook>"""
    open(os.path.join(OUT, "sample.fb2"), "wb").write(fb2.encode("utf-8"))

# ---------- CBZ ----------
def make_cbz():
    from zipfile import ZipFile, ZIP_STORED
    with ZipFile(os.path.join(OUT, "sample.cbz"), "w", ZIP_STORED) as z:
        z.writestr("01.png", png(30, 30, (10, 200, 10)))
        z.writestr("02.png", png(40, 20, (200, 10, 10)))
        z.writestr("03.png", png(50, 50, (10, 10, 220)))

# ---------- MOBI7 ----------
def palmdoc_compress(data):
    """把字节串压成 PalmDoc LZ77（对话 libmobi mobi_decompress_lz77）。
    最简单的正确编码：全部字面量用「长度 1..8 的 run」。"""
    out = bytearray()
    i = 0
    n = len(data)
    while i < n:
        run = min(8, n - i)
        out.append(run)          # 0x01..0x08 → 拷贝 run 个字面字节
        out += data[i:i+run]
        i += run
    return bytes(out)

def make_mobi():
    text = (
        "<html><body>"
        "<h1>MOBI 测试</h1>"
        "<p>这是 Mobi 正文，用于验证 libmobi rawml 抽取与粗解析。</p>"
        "<p>第二段，含 entity &amp; 与标点。</p>"
        "</body></html>"
    )
    print("  mobi text bytes:", len(text.encode("utf-8")))
    raw = text.encode("utf-8")
    # PalmDoc 压缩后的文本记录（text_length = 未压缩长度）
    comp = palmdoc_compress(raw)
    # PDB 记录：
    # record0 = PalmDoc 头(16) + MOBI 头(即 mobi_header) + 其余记录 offset 表等，需手排。
    # 这里用 libmobi 认识的最小结构：
    # record0 内含: compression=2(uncompressed)? 实际 PalmDoc 压缩 compression=2(PalmDoc)？看 mobi.h
    # 我们按 MOBI_HEADER 经典布局构造。
    import io
    # --- 构造 record0 ---
    # PalmDoc/EOF 头 16 字节（MOBIRecord0Header）
    rec0_hdr = struct.pack("<HHIHHI",
        2,          # compression = 2 -> PalmDoc LZ77? 0=none,1=?,2=PalmDoc
        0,          # unused
        len(raw),   # text_length (uncompressed)
        1,          # record_count (text records)
        4096,       # record_size
        0,          # encryption type
        )          # 12 bytes? 实际是 16 字节: compression,unused,text_length,record_count,record_size,encryption,unknown
    # MOBIRecord0Header 实际 16 字节: compression(2),unused(2),text_length(4),record_count(2),record_size(2),encryption(2),unknown(2)
    rec0_hdr = struct.pack(">HHIHHHH", 2, 0, len(raw), 1, 4096, 0, 0)
    # MOBI 头（MOBIMobiHeader），经典字段（我们只需让 libmobi 能 parse）
    mobi_header = bytearray()
    mobi_header += b"MOBI"                       # identifier
    mobi_header += struct.pack(">I", 0x100)      # header length (0x100)
    mobi_header += struct.pack(">I", 0x10000000) # mobi type
    mobi_header += struct.pack(">I", 0)          # text encoding (0 = cp1252, 65001=utf8)
    mobi_header += struct.pack(">I", 0)          # unique id
    mobi_header += struct.pack(">I", 0)          # file version
    mobi_header += struct.pack(">I", 0)          # orthographic index
    mobi_header += struct.pack(">I", 0)          # inflection index
    mobi_header += struct.pack(">I", 0)          # index names
    mobi_header += struct.pack(">I", 0)          # index keys
    mobi_header += struct.pack(">I", 0)          # extra index 0..5
    mobi_header += struct.pack(">I", 0)
    mobi_header += struct.pack(">I", 0)
    mobi_header += struct.pack(">I", 0)
    mobi_header += struct.pack(">I", 0)
    mobi_header += struct.pack(">I", 0)
    mobi_header += struct.pack(">I", 0)          # first non-book index
    mobi_header += struct.pack(">I", 0)          # full name offset
    mobi_header += struct.pack(">I", 0)          # full name length
    mobi_header += struct.pack(">I", 0)          # locale
    mobi_header += struct.pack(">I", 0)          # input language
    mobi_header += struct.pack(">I", 0)          # output language
    mobi_header += struct.pack(">I", 0)          # min version
    mobi_header += struct.pack(">I", 0)          # first image index
    mobi_header += struct.pack(">I", 0)          # huffman record offset
    mobi_header += struct.pack(">I", 0)          # huffman record count
    mobi_header += struct.pack(">I", 0)          # huffman table offset
    mobi_header += struct.pack(">I", 0)          # huffman table length
    mobi_header += struct.pack(">I", 0)          # EXTH flags
    mobi_header += struct.pack(">I", 32)         # 32 unknown bytes
    mobi_header += b"\x00" * 32
    mobi_header += struct.pack(">I", 0)          # unknown 0..4 (to 0x100)
    mobi_header += b"\x00" * (0x100 - len(mobi_header))
    # 内容记录：文本记录 + 一个 0xE98E0B0D EOF 记录
    rec1 = comp
    eof = struct.pack(">II", 0xE98E0B0D, 0)   # EOF 记录
    records = [rec0_hdr + bytes(mobi_header), rec1, eof]
    # 构造 PDB
    name = "TESTMOBI"
    pdb = bytearray()
    name_block = name.encode("utf-8")[:31].ljust(32, b"\x00")
    pdb += name_block
    pdb += struct.pack(">H", 0)     # attributes
    pdb += struct.pack(">H", 0)     # version
    pdb += struct.pack(">I", 0)     # creation date
    pdb += struct.pack(">I", 0)     # modification
    pdb += struct.pack(">I", 0)     # backup
    pdb += struct.pack(">I", 0)     # modification number
    pdb += struct.pack(">I", 0)     # appInfo
    pdb += struct.pack(">I", 0)     # sortInfo
    pdb += b"BOOK"                  # type
    pdb += b"MOBI"                  # creator
    pdb += struct.pack(">I", 0)     # uniqueID seed
    pdb += struct.pack(">I", 0)     # nextRecordListID
    pdb += struct.pack(">H", len(records))  # numRecords
    # record 表（8 字节 each）: offset(4), attr(1), uniqueID(3)
    data_start = 78 + 8 * len(records) + 2   # + appInfo? PDB 有 appInfo+sortInfo? 这里偏移从 78 起步
    # 实际 PDB: 头 78 字节，然后 record list，然后 appInfo (2 bytes 0,1), sortInfo(2 bytes) 可缺省
    # 我们用 title 记录头: 78 + 8*n; 记录区从 (78+8*n+? ) 起；这里不加 appInfo 让 offset 简单
    offs = []
    cur = 78 + 8 * len(records)
    for r in records:
        offs.append(cur)
        cur += len(r)
    for i, o in enumerate(offs):
        pdb += struct.pack(">I", o)          # offset (4)
        pdb += struct.pack("<B", 0)          # attr (1)
        pdb += (i & 0xFFFFFF).to_bytes(3, "little")   # unique ID (3 字节)
    for r in records:
        pdb += r
    # 补 appInfo/sortInfo? libmobi 可能要求 record0 后紧跟 appInfo。先不加，看解析结果。
    path = os.path.join(OUT, "sample.mobi")
    open(path, "wb").write(bytes(pdb))
    print("  wrote", path, len(pdb), "bytes")

make_epub(); make_fb2(); make_cbz(); make_mobi()
print("samples written to", OUT)
for f in sorted(os.listdir(OUT)):
    print(" ", f, os.path.getsize(os.path.join(OUT, f)), "bytes")