#!/usr/bin/env python3
"""
Pull the bitmaps out of a WinHelp file and write them as PNGs.

Some of what a help file states is stated in a picture — a table too wide for the
text flow, a diagram, a chart layout. Those are the parts that stay invisible to a
text extractor, and they are often the parts a reimplementation most needs.

Format, since it is nowhere convenient:

  |bmNN is an SHG/MRB object list
      magic     two ASCII letters, one of 'lp' 'lP' 'Lp' 'LP' — lower-case
                first letter means the objects are bitmaps or metafiles, and
                the second letter's case distinguishes single from multiple
                resolutions. All four are accepted here; the distinction does
                not change how an object is read.
      count     uint16, number of objects
      offsets   uint32 each, from the start of the file

  each object
      type      byte: 5 bitmap, 6 metafile, 8 embedded window
      packing   byte: 0 none, 1 RLE, 2 LZ77, 3 both

  the object header that follows (verified against this file rather than
  taken on trust, since the published descriptions disagree):

      +2   cdword   unused here
      +4   cdword   size after the LZ77 stage
      +8   cword    planes
      +9   cword    bits per pixel
      +10  cdword   width
      +12  cdword   height
      +14  cdword   unused here
      +16  cdword   colours used
      +18  cdword   compressed size
      +20  cdword   hotspot size
      +22  dword    offset of the pixel data from the object start
      +26  dword    offset of the hotspot data

  Packing method 3 means LZ77 first and then run-length, in that order. The
  proof that the reading is right rather than merely plausible is arithmetic:
  the LZ77 stage must produce exactly the byte count at +4, and the run-length
  stage exactly width x height at the stated depth. Both are asserted below, so
  a misparse reports itself instead of writing a black image.

A "compressed" integer is a variable-length encoding keyed on the low bit: clear
means the value is what remains after a shift of one, set means another byte or
word follows. Getting that wrong shifts everything after it, which is why the
reader below is a class with a position rather than a chain of struct calls.
"""

import struct
import sys
import zlib


class Reader:
    def __init__(self, data: bytes, pos: int = 0):
        self.d = data
        self.p = pos

    def byte(self) -> int:
        v = self.d[self.p]
        self.p += 1
        return v

    def word(self) -> int:
        v = struct.unpack_from('<H', self.d, self.p)[0]
        self.p += 2
        return v

    def dword(self) -> int:
        v = struct.unpack_from('<I', self.d, self.p)[0]
        self.p += 4
        return v

    def cword(self) -> int:
        """A compressed word: one byte if the low bit is clear, else two."""
        b = self.byte()
        if b & 1:
            return ((self.byte() << 8) | b) >> 1
        return b >> 1

    def cdword(self) -> int:
        """A compressed dword: one word if the low bit is clear, else two."""
        w = self.word()
        if w & 1:
            return ((self.word() << 16) | w) >> 1
        return w >> 1


def rle_decode(src: bytes, expected: int) -> bytearray:
    """WinHelp's run-length coding: a count byte, high bit set for a run."""
    out = bytearray()
    i = 0
    while i < len(src) and len(out) < expected:
        n = src[i]
        i += 1
        if n & 0x80:
            count = n & 0x7F
            out += src[i:i + count]
            i += count
        else:
            if i >= len(src):
                break
            out += bytes([src[i]]) * n
            i += 1
    return out


def lz77_decode(src: bytes) -> bytearray:
    """The same LZ77 the topic blocks use."""
    out = bytearray()
    i, n = 0, len(src)
    while i < n:
        ctrl = src[i]
        i += 1
        for bit in range(8):
            if i >= n:
                return out
            if not (ctrl >> bit) & 1:
                out.append(src[i])
                i += 1
                continue
            if i + 1 >= n:
                return out
            word = src[i] | (src[i + 1] << 8)
            i += 2
            dist = (word & 0x0FFF) + 1
            length = (word >> 12) + 3
            start = len(out) - dist
            if start < 0:
                return out
            for k in range(length):
                out.append(out[start + k])
    return out


def unpack(src: bytes, method: int, expected: int) -> bytes:
    if method == 0:
        return src
    if method == 1:
        return bytes(rle_decode(src, expected))
    if method == 2:
        return bytes(lz77_decode(src))
    if method == 3:
        return bytes(rle_decode(bytes(lz77_decode(src)), expected))
    raise ValueError(f'unknown packing method {method}')


def png(width: int, height: int, rows_rgb: list) -> bytes:
    """A minimal PNG writer — no dependency worth adding for this."""
    raw = b''.join(b'\x00' + row for row in rows_rgb)

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

    return (b'\x89PNG\r\n\x1a\n'
            + chunk(b'IHDR', struct.pack('>IIBBBBB', width, height, 8, 2, 0, 0, 0))
            + chunk(b'IDAT', zlib.compress(raw, 9))
            + chunk(b'IEND', b''))


def decode_bitmap(data: bytes, offset: int):
    """One object → (width, height, list of RGB rows, top row first).

    Returns None, with a reason on stderr, rather than guessing: an image that
    comes out black is worse than one that does not come out, because it looks
    like a picture of nothing rather than like a failure.
    """
    r = Reader(data, offset)
    obj_type = r.byte()
    packing = r.byte()

    r.cdword()                        # unused
    stage1_size = r.cdword()          # size after the LZ77 stage
    planes = r.cword()
    bitcount = r.cword()
    width = r.cdword()
    height = r.cdword()
    r.cdword()                        # unused
    colors_used = r.cdword()
    compressed_size = r.cdword()
    hotspot_size = r.cdword()
    data_offset = r.dword()
    r.dword()                         # hotspot offset

    if width <= 0 or height <= 0 or width > 20000 or height > 20000:
        print(f'  implausible dimensions {width}x{height}', file=sys.stderr)
        return None
    if bitcount not in (1, 4, 8, 24):
        print(f'  {bitcount} bits per pixel is not handled', file=sys.stderr)
        return None

    palette = []
    if colors_used:
        pr = Reader(data, offset + 30)
        for _ in range(colors_used):
            bb, g, rr, _z = pr.byte(), pr.byte(), pr.byte(), pr.byte()
            palette.append((rr, g, bb))
    if not palette:
        if bitcount == 1:
            palette = [(0, 0, 0), (255, 255, 255)]
        elif bitcount <= 8:
            palette = [(i, i, i) for i in range(1 << bitcount)]

    src = data[offset + data_offset:offset + data_offset + compressed_size] \
        if compressed_size else data[offset + data_offset:]

    row_bytes = ((width * bitcount + 31) // 32) * 4
    expected = row_bytes * height

    if packing == 0:
        pixels = src
    elif packing == 1:
        pixels = bytes(rle_decode(src, expected))
    elif packing == 2:
        pixels = bytes(lz77_decode(src))
    elif packing == 3:
        stage1 = bytes(lz77_decode(src))
        if stage1_size and len(stage1) != stage1_size:
            print(f'  LZ77 gave {len(stage1)} bytes, header says {stage1_size}',
                  file=sys.stderr)
            return None
        pixels = bytes(rle_decode(stage1, expected))
    else:
        print(f'  unknown packing method {packing}', file=sys.stderr)
        return None

    if len(pixels) != expected:
        print(f'  unpacked to {len(pixels)} bytes, expected {expected} '
              f'({width}x{height} at {bitcount}bpp)', file=sys.stderr)
        return None

    rows = []
    for y in range(height):
        base = y * row_bytes
        row = bytearray()
        for x in range(width):
            if bitcount == 24:
                bb = pixels[base + x * 3]
                g = pixels[base + x * 3 + 1]
                rr = pixels[base + x * 3 + 2]
                row += bytes((rr, g, bb))
                continue
            if bitcount == 1:
                idx = (pixels[base + (x >> 3)] >> (7 - (x & 7))) & 1
            elif bitcount == 4:
                byte = pixels[base + (x >> 1)]
                idx = (byte >> 4) if x % 2 == 0 else (byte & 0x0F)
            else:
                idx = pixels[base + x]
            row += bytes(palette[idx] if idx < len(palette) else (0, 0, 0))
        rows.append(bytes(row))

    rows.reverse()          # device-independent bitmaps run bottom-up
    return width, height, rows


def main():
    import hlp
    blob = open(sys.argv[1], 'rb').read()
    out_dir = sys.argv[2] if len(sys.argv) > 2 else '.'
    h = hlp.Hlp(blob)

    made = 0
    for name in sorted(h.files):
        if not name.startswith('|bm'):
            continue
        data = h.read(name)
        if data is None or len(data) < 8:
            continue

        if data[:2] not in (b'lp', b'lP', b'Lp', b'LP'):
            print(f'{name}: unexpected magic {data[:2]!r}', file=sys.stderr)
            continue
        count = struct.unpack_from('<H', data, 2)[0]
        offsets = [struct.unpack_from('<I', data, 4 + 4 * i)[0] for i in range(count)]

        for i, off in enumerate(offsets):
            try:
                res = decode_bitmap(data, off)
            except Exception as e:
                print(f'{name}[{i}]: {e}', file=sys.stderr)
                continue
            if res is None:
                print(f'{name}[{i}]: not a bitmap this reader handles', file=sys.stderr)
                continue
            w, ht, rows = res
            suffix = '' if count == 1 else f'_{i}'
            path = f'{out_dir}/{name[1:]}{suffix}.png'
            open(path, 'wb').write(png(w, ht, rows))
            print(f'{path}  {w}x{ht}')
            made += 1

    print(f'{made} images written', file=sys.stderr)


if __name__ == '__main__':
    main()