#!/usr/bin/env python3
"""
Read the prose out of a WinHelp (.hlp) file.

    python3 winhelp_extract.py path/to/file.hlp out.txt

Why this exists: the calculation rules this project needs are stated in words
inside a help file, and nothing on hand would open it. The rules themselves are
facts and formulae, which are free to use; the prose is the author's, so this
tool stays a reader and nothing it produces is copied into the application.

The format, for whoever comes next:

  header        magic 0x00035F3F, directory offset at +4
  every file    FILEHEADER = reserved:int32, used:int32, flags:uint8 (9 bytes)
  directory     B+ tree; leaf pages hold NUL-terminated names + int32 offsets
  |TOPIC        4096-byte blocks: 12-byte header then an LZ77 body
  TopicPos      paged: block * 0x4000 + offset + 12, which is how a record that
                straddles a block boundary is found again
  phrases       4.0 uses Hall compression: |PhrImage holds the phrases end to
                end and |PhrIndex holds their lengths as a bit stream

Three details cost most of the time and are worth stating plainly:

  1. The |PhrIndex bit stream is a unary run of 1-bits, then a 0, then `bits`
     more bits read LOW-ORDER FIRST; the length is (run << bits) + low + 1. The
     check that the reading is right, rather than merely plausible, is that the
     offsets end at exactly the decompressed image size the header declares.

  2. The text stream is a prefix code counted in TRAILING one-bits, not leading:
     `xxxxxxx0` one-byte phrase 0..127, `xxxxxx01` two-byte phrase
     128 + (((b >> 2) << 8) | next), `xxxxx011` a literal run, `xxxx0111`
     spaces, `xxx01111` NULs. The giveaway is that the phrase table sorts
     alphabetically in two independent runs — the first 128 entries and then the
     rest — because the first 128 are exactly those a single byte can address.

  3. A record whose stored length already equals its stated length was not worth
     compressing and is held verbatim. Expanding one of those produces fluent
     word salad, which is the worst kind of wrong: check the length first.
"""

import struct
import sys


class Hlp:
    def __init__(self, blob: bytes):
        self.b = blob
        magic, dir_start = struct.unpack_from('<Ii', blob, 0)
        if magic != 0x00035F3F:
            raise ValueError(f'not a WinHelp file (magic {magic:#x})')
        self.files = self._directory(dir_start)

    # -- internal file system ------------------------------------------------

    def _file_body(self, offset: int) -> bytes:
        reserved, used, flags = struct.unpack_from('<iiB', self.b, offset)
        return self.b[offset + 9: offset + 9 + used]

    def _directory(self, offset: int) -> dict:
        body = self._file_body(offset)
        # BTREEHEADER: magic(2) flags(2) pagesize(2) structure(16) mustbezero(2)
        # pagesplits(2) rootpage(2) mustbeneg1(2) totalpages(2) nlevels(2)
        # totalentries(4)  = 38 bytes
        magic, flags, page_size = struct.unpack_from('<HHH', body, 0)
        if magic != 0x293B:
            raise ValueError(f'bad btree magic {magic:#x}')
        root, levels = struct.unpack_from('<h', body, 26)[0], struct.unpack_from('<H', body, 32)[0]
        pages = body[38:]

        def read_page(n):
            return pages[n * page_size:(n + 1) * page_size]

        # Descend the index levels to the first leaf.
        page = root
        for _ in range(levels - 1):
            p = read_page(page)
            page = struct.unpack_from('<h', p, 6)[0]   # PreviousPage

        out = {}
        while page != -1:
            p = read_page(page)
            n_entries, prev_page, next_page = struct.unpack_from('<hhh', p, 2)
            pos = 8
            for _ in range(n_entries):
                end = p.index(b'\x00', pos)
                name = p[pos:end].decode('latin-1')
                pos = end + 1
                out[name] = struct.unpack_from('<i', p, pos)[0]
                pos += 4
            page = next_page
        return out

    def read(self, name: str):
        off = self.files.get(name)
        return None if off is None else self._file_body(off)

    # -- LZ77 ----------------------------------------------------------------

    @staticmethod
    def lz77(src: bytes) -> bytearray:
        """WinHelp's LZ77: a control byte per eight items, LSB first.

        A clear bit is one literal byte. A set bit is a 16-bit word: the low 12
        bits are (distance - 1) back into the output, the high 4 bits are
        (length - 3)."""
        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

    # -- phrase table --------------------------------------------------------

    def phrases(self):
        """The phrase list, whichever of the two schemes this file uses."""
        img = self.read('|PhrImage')
        idx = self.read('|PhrIndex')
        if img is not None and idx is not None:
            return self._phrases_hall(idx, img)

        raw = self.read('|Phrases')
        if raw is None:
            return []
        count = struct.unpack_from('<H', raw, 0)[0]
        if struct.unpack_from('<H', raw, 2)[0] == 0x0100:      # compressed
            body = self.lz77(raw[2 + 2 + 4:])
            offsets_src = raw[4:]
        else:
            body = raw
            offsets_src = raw
        # offsets are int16 from the start of the table
        offs = [struct.unpack_from('<H', raw, 4 + 2 * i)[0] for i in range(count + 1)]
        base = offs[0]
        return [body[o - base:offs[i + 1] - base] for i, o in enumerate(offs[:-1])]

    @staticmethod
    def _phrases_hall(idx: bytes, img: bytes):
        """Hall compression: |PhrImage holds the phrases end to end, |PhrIndex
        holds their lengths as a bit stream.

        Each length is a unary run of 1-bits followed by a 0, then `bits` more
        bits read low-order first: length = (run << bits) + low + 1. The result
        is checked against the decompressed image size, so a wrong reading is
        rejected rather than silently producing garbage phrases."""
        magic, n_entries, compressed_size = struct.unpack_from('<III', idx, 0)
        phr_img_size, phr_img_compressed_size = struct.unpack_from('<II', idx, 12)
        bits = struct.unpack_from('<H', idx, 24)[0] & 0x0F

        text = bytes(Hlp.lz77(img)) if phr_img_size != phr_img_compressed_size else img
        if len(text) != phr_img_size:
            return []

        data = idx[28:]
        pos = 0
        total_bits = len(data) * 8

        def bit():
            nonlocal pos
            if pos >= total_bits:
                raise EOFError
            b = (data[pos >> 3] >> (pos & 7)) & 1
            pos += 1
            return b

        offs = [0]
        acc = 0
        try:
            for _ in range(n_entries):
                run = 0
                while bit():
                    run += 1
                low = 0
                for l in range(bits):
                    if bit():
                        low |= 1 << l
                acc += (run << bits) + low + 1
                offs.append(acc)
        except EOFError:
            return []

        if offs[-1] != phr_img_size:
            return []
        return [text[offs[i]:offs[i + 1]] for i in range(n_entries)]

    # -- topic text ----------------------------------------------------------

    def topic_blocks(self, block_size=4096, compressed=True):
        topic = self.read('|TOPIC')
        if topic is None:
            return
        for off in range(0, len(topic), block_size):
            block = topic[off:off + block_size]
            if len(block) <= 12:
                continue
            body = block[12:]
            yield self.lz77(body) if compressed else body


def expand(data: bytes, phrases) -> bytes:
    """Hall phrase expansion.

    The lead byte is a prefix code counted in TRAILING one-bits, and the rest
    of the byte is the payload:

      xxxxxxx0   phrase 0..127, the common table, one byte:  n = b >> 1
      xxxxxx01   phrase 128 and up, two bytes: n = 128 + (((b >> 2) << 8) | next)
      xxxxx011   (b >> 3) + 1 literal bytes follow
      xxxx0111   (b >> 4) + 1 spaces
      xxx01111   (b >> 5) + 1 NULs

    The phrase table is sorted in two runs for exactly this reason: the first
    128 entries are the words common enough to earn a one-byte reference, the
    remainder need two. That split is what proves the reading: a table whose
    first 128 entries sort independently of the rest is not an accident.
    """
    out = bytearray()
    i, n = 0, len(data)
    while i < n:
        b = data[i]
        if b & 1 == 0:
            p = b >> 1
            i += 1
            if p < len(phrases):
                out += phrases[p]
        elif b & 7 == 3:
            count = (b >> 3) + 1
            i += 1
            out += data[i:i + count]
            i += count
        elif b & 3 == 1:
            if i + 1 >= n:
                break
            p = 128 + (((b >> 2) << 8) | data[i + 1])
            i += 2
            if p < len(phrases):
                out += phrases[p]
        elif b & 15 == 7:
            out += b' ' * ((b >> 4) + 1)
            i += 1
        elif b & 31 == 15:
            out += b'\x00' * ((b >> 5) + 1)
            i += 1
        elif b & 63 == 31:
            out += b'\x00' * ((b >> 6) + 1)
            i += 1
        else:
            i += 1
    return bytes(out)


def runs(data: bytes, minlen=24):
    """Printable runs — a fallback for a file whose framing will not reconstruct."""
    cur = bytearray()
    for byte in data:
        if 32 <= byte < 127 or byte in (9,):
            cur.append(byte)
        else:
            if len(cur) >= minlen:
                yield cur.decode('latin-1')
            cur = bytearray()
    if len(cur) >= minlen:
        yield cur.decode('latin-1')


def topics(h, block_size=4096):
    """Walk the TOPICLINK chain block by block and yield each displayable
    record's text.

    A link may run past the end of its block; the next block's header carries
    FirstTopicLink, the offset at which a fresh link starts there, so a
    straddling link's tail is the bytes before that offset. Records are
    therefore assembled across blocks rather than dropped, which is what the
    earlier attempt got wrong.
    """
    topic = h.read('|TOPIC')
    if topic is None:
        return
    ph = h.phrases()

    blocks = []
    for off in range(0, len(topic), block_size):
        blk = topic[off:off + block_size]
        if len(blk) <= 12:
            continue
        index = off // block_size
        last_link, first_link, last_header = struct.unpack_from('<iii', blk, 0)
        # TopicPos is paged: block * 0x4000 + offset + 12.
        first_offset = max(0, first_link - index * 0x4000 - 12)
        blocks.append((first_offset, bytes(Hlp.lz77(blk[12:]))))

    def emit(link: bytes):
        if len(link) < 21:
            return None
        block_len, data_len2, prev_b, next_b, data_len1 = struct.unpack_from('<iiiii', link, 0)
        record_type = link[20]
        if record_type not in (0x20, 0x23) or data_len1 < 21 or data_len1 > len(link):
            return None
        raw = link[data_len1:]
        # A record whose stored length already equals its stated length was not
        # worth phrase-compressing and is held verbatim. Running the expander
        # over one of those is what turned a table in this file into word
        # salad — the length is the tell, so check it before expanding.
        return record_type, (raw if len(raw) == data_len2 else expand(raw, ph))

    pending = b''
    for first_offset, data in blocks:
        if pending:
            # The tail of a link that began in the previous block.
            out = emit(pending + data[:first_offset])
            if out:
                yield out
            pending = b''
        pos = first_offset
        while pos + 21 <= len(data):
            block_len, data_len2, prev_b, next_b, data_len1 = struct.unpack_from('<iiiii', data, pos)
            if block_len <= 0 or data_len1 < 21 or data_len1 > block_len:
                break                      # padding to the end of the block
            if pos + block_len <= len(data):
                out = emit(data[pos:pos + block_len])
                if out:
                    yield out
                pos += block_len
            else:
                pending = data[pos:]       # straddles into the next block
                break


def main():
    path = sys.argv[1]
    out_path = sys.argv[2] if len(sys.argv) > 2 else '/dev/stdout'
    blob = open(path, 'rb').read()
    h = Hlp(blob)

    ph = h.phrases()
    print(f'internal files: {len(h.files)}, phrases: {len(ph)}', file=sys.stderr)

    with open(out_path, 'w', encoding='utf-8') as f:
        count = 0
        for rt, text in topics(h):
            s = text.decode('latin-1').replace('\x00', '\n')
            f.write(s)
            f.write('\n')
            count += 1
        print(f'records written: {count}', file=sys.stderr)


if __name__ == '__main__':
    main()