#!/usr/bin/env python3
"""A reference encoder and decoder for Detour files, written from the format's description, independently of
the Rust `detour` tool. It is slow (about a million letters a second) but short, so it can be read and checked.

    python3 reference.py decode [--max-length L] FILE.detour [OUTPUT]    write the word (to standard output if no OUTPUT)
    python3 reference.py encode --n N [--alphabet SYMBOLS] [--sync E] INPUT OUTPUT.detour

Standard library only (Python 3.10 or later).
"""

import argparse
import hashlib
import lzma
import struct
import sys

MAGIC, VERSION = b"DTR1", 1
BASE = {1: 0, 2: 1, 3: 3, 4: 9, 5: 33}  # the code byte of a step of weight k <= 5 and rank r is BASE[k] + r
SYNC, BIG, LIT, START, END = 251, 252, 253, 254, 255
LITERAL_PIECE = 1 << 20
XZ_FILTERS = [{"id": lzma.FILTER_LZMA2, "preset": 9 | lzma.PRESET_EXTREME, "lc": 0, "lp": 0, "pb": 0}]
XZ_MEMORY_LIMIT = 128 << 20  # a decoder refuses an xz stream that needs more memory than this to decompress


class DetourError(ValueError):
    pass


def factorial(k):
    result = 1
    for i in range(2, k + 1):
        result *= i
    return result


def rank(removed, new):
    """The rank of `new` as an ordering of `removed`; the reference order is `removed` reversed."""
    reference = list(reversed(removed))
    r = 0
    for j, letter in enumerate(new):
        c = reference.index(letter)
        r = r * (len(removed) - j) + c
        del reference[c]
    return r


def unrank(removed, r):
    k = len(removed)
    if not 0 <= r < factorial(k):
        raise DetourError(f"rank {r} is out of range for weight {k}")
    reference = list(reversed(removed))
    digits = []
    for base in range(1, k + 1):  # the last digit has base 1, the first base k
        digits.append(r % base)
        r //= base
    digits.reverse()
    out = bytearray()
    for c in digits:
        out.append(reference.pop(c))
    return bytes(out)


def default_weight(levels, n):
    """The least weight k whose level is below n - k, or None when every level is finished."""
    for k in range(1, n):
        if levels[k] < n - k:
            return k
    return None


def record(levels, k):
    levels[k] += 1
    for j in range(1, k):
        levels[j] = 0


def put_varint(out, v):
    while v >= 0x80:
        out.append((v & 0x7F) | 0x80)
        v >>= 7
    out.append(v)


def check_header(n, alphabet, sync):
    if not 2 <= n <= 16 or len(alphabet) != n or len(set(alphabet)) != n or not 0 <= sync <= 40:
        raise DetourError("bad header")


def encode_tokens(word, n, alphabet, sync=0):
    """The token stream for `word` (bytes): what a .detour file holds inside its xz container."""
    alphabet = bytes(alphabet)
    check_header(n, alphabet, sync)
    out = bytearray(MAGIC + bytes([VERSION, n]) + alphabet + bytes([sync]))
    symbols = set(alphabet)
    gap = since_sync = 0
    levels = [0] * (n + 1)
    window = None  # the current window, or None after a literal
    pending = bytearray()

    def token(code, payload=b""):
        nonlocal gap, since_sync
        put_varint(out, gap)
        gap = 0
        out.append(code)
        out.extend(payload)
        since_sync += 1

    def literal(data):
        payload = bytearray()
        put_varint(payload, len(data))
        token(LIT, bytes(payload) + bytes(data))

    def maybe_sync(offset):
        nonlocal since_sync
        if sync and since_sync >= 1 << sync:
            put_varint(out, 0)
            out.append(SYNC)
            out.extend(struct.pack("<Q", offset) + window)
            for k in range(1, n + 1):
                put_varint(out, levels[k])
            since_sync = 0

    for i, b in enumerate(word):
        pending.append(b)
        last = word[i + 1 - n:i + 1] if i + 1 >= n else b""
        if len(last) == n and set(last) <= symbols and len(set(last)) == n:
            if window is not None and len(pending) <= n:
                k = len(pending)
                r = rank(window[:k], bytes(pending))
                deviates = not (default_weight(levels, n) == k and r == 0)
                if deviates:
                    if k <= 5:
                        token(BASE[k] + r)
                    else:
                        payload = bytearray([k])
                        put_varint(payload, r)
                        token(BIG, bytes(payload))
                else:
                    gap += 1
                record(levels, k)
                window = window[k:] + bytes(pending)
                if deviates:
                    maybe_sync(i + 1)
            else:
                if len(pending) > n:
                    literal(pending[:-n])
                window = bytes(last)
                levels = [0] * (n + 1)
                token(START, window)
                maybe_sync(i + 1)
            pending = bytearray()
        elif len(pending) >= LITERAL_PIECE:
            literal(pending[:-(n - 1)])
            del pending[:-(n - 1)]
            window = None
    if pending:
        literal(pending)
    put_varint(out, gap)
    out.append(END)
    out.extend(struct.pack("<Q", len(word)) + hashlib.sha256(word).digest())
    return bytes(out)


class Reader:
    def __init__(self, data):
        self.data, self.pos = data, 0

    def take(self, k):
        if self.pos + k > len(self.data):
            raise DetourError("the stream ends early")
        chunk = self.data[self.pos:self.pos + k]
        self.pos += k
        return chunk

    def byte(self):
        return self.take(1)[0]

    def varint(self):
        """An unsigned LEB128 varint of at most 10 bytes, below 2^64, in its shortest form: a varint of two or more
        bytes may not end in the byte 0x00."""
        value = 0
        for i in range(10):
            b = self.byte()
            value |= (b & 0x7F) << (7 * i)
            if b < 0x80:
                if value >= 1 << 64:
                    break
                if i > 0 and b == 0:
                    raise DetourError("a varint is not in its shortest form")
                return value
        raise DetourError("a varint is too long")


def decode_tokens(data, limit=None):
    """The word a token stream holds, checking the stream as the Rust decoder does."""
    rd = Reader(data)
    if rd.take(4) != MAGIC:
        raise DetourError("not a Detour stream")
    if rd.byte() != VERSION:
        raise DetourError("unsupported version")
    n = rd.byte()
    if not 2 <= n <= 16:
        raise DetourError("bad header")
    alphabet = rd.take(n)
    sync = rd.byte()
    check_header(n, alphabet, sync)
    out = bytearray()
    window = None
    levels = [0] * (n + 1)
    since_sync = 0
    after_window_token = False

    def emit(letters):
        if limit is not None and len(out) + len(letters) > limit:
            raise DetourError(f"the word is longer than {limit} letters")
        out.extend(letters)

    def step(k, r):
        nonlocal window
        new = unrank(window[:k], r)
        emit(new)
        window = window[k:] + new
        record(levels, k)

    while True:
        gap = rd.varint()
        code = rd.byte()
        due = sync and after_window_token and since_sync >= 1 << sync
        if due and (code != SYNC or gap):
            raise DetourError("a SYNC token is missing here")
        if code == SYNC and not due:
            raise DetourError("a SYNC token where none is due")
        after_window_token = False
        if gap and window is None:
            raise DetourError("default steps with no current window")
        for _ in range(gap):
            k = default_weight(levels, n)
            if k is None:
                raise DetourError("a default step where every level is finished")
            step(k, 0)
        if code == SYNC:
            offset = struct.unpack("<Q", rd.take(8))[0]
            synced = rd.take(n)
            counters = [rd.varint() for _ in range(n)]
            if offset != len(out) or synced != window or counters != levels[1:n + 1]:
                raise DetourError("a SYNC token does not match the decoder's state")
            since_sync = 0
            continue
        if code == END:
            break
        since_sync += 1
        if code == START:
            window = rd.take(n)
            if sorted(window) != sorted(alphabet):
                raise DetourError("START is not a permutation of the alphabet")
            emit(window)
            levels = [0] * (n + 1)
            after_window_token = True
            continue
        if code == LIT:
            length = rd.varint()
            if length == 0:
                raise DetourError("an empty literal")
            emit(rd.take(length))
            window = None
            continue
        if code == BIG:
            k, r = rd.byte(), rd.varint()
            if not 6 <= k <= n:
                raise DetourError(f"a BIG step of weight {k} with n = {n}")
        elif code <= 152:
            k = 1 if code < 1 else 2 if code < 3 else 3 if code < 9 else 4 if code < 33 else 5
            r = code - BASE[k]
        else:
            raise DetourError(f"unknown code byte {code}")
        if window is None:
            raise DetourError("a step with no current window")
        if k > n:
            raise DetourError(f"a step of weight {k} with n = {n}")
        if default_weight(levels, n) == k and r == 0:
            raise DetourError("a step token equal to the default step")
        step(k, r)
        after_window_token = True
    length = struct.unpack("<Q", rd.take(8))[0]
    digest = rd.take(32)
    if length != len(out) or digest != hashlib.sha256(out).digest():
        raise DetourError("the decoded word does not match the length and SHA-256 in the trailer")
    if rd.pos != len(data):
        raise DetourError("data follows the end of the Detour stream")
    return bytes(out)


def read_file(path):
    """The token stream inside a .detour file: exactly one xz stream, which decompresses in at most 128 MiB of
    memory, with nothing after it."""
    decompressor = lzma.LZMADecompressor(format=lzma.FORMAT_XZ, memlimit=XZ_MEMORY_LIMIT)
    with open(path, "rb") as f:
        raw = f.read()
        if not raw.startswith(b"\xfd7zXZ\x00"):
            raise DetourError("not a valid xz stream")
        try:
            data = decompressor.decompress(raw)
        except lzma.LZMAError as error:
            if "Memory usage limit" in str(error):  # liblzma's memory limit; lzma gives no error code
                raise DetourError("the file needs more than 128 MiB to decompress") from error
            raise DetourError("not a valid xz stream") from error
    if not decompressor.eof or decompressor.unused_data:
        raise DetourError("not exactly one complete xz stream")
    return data


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    sub = parser.add_subparsers(dest="command", required=True)
    dec = sub.add_parser("decode")
    dec.add_argument("--max-length", type=int, help="refuse a word longer than this many letters")
    dec.add_argument("input")
    dec.add_argument("output", nargs="?")
    enc = sub.add_parser("encode")
    enc.add_argument("--n", type=int, required=True)
    enc.add_argument("--alphabet")
    enc.add_argument("--sync", type=int)
    enc.add_argument("input")
    enc.add_argument("output")
    args = parser.parse_args(argv)
    try:
        if args.command == "decode":
            word = decode_tokens(read_file(args.input), limit=args.max_length)
            if args.output:
                with open(args.output, "wb") as f:
                    f.write(word)
            else:
                sys.stdout.buffer.write(word)
        else:
            alphabet = (args.alphabet or "0123456789ABCDEF"[:args.n]).encode()
            sync = args.sync if args.sync is not None else (21 if args.n >= 14 else 0)
            with open(args.input, "rb") as f:
                tokens = encode_tokens(f.read(), args.n, alphabet, sync)
            with open(args.output, "wb") as f:
                f.write(lzma.compress(tokens, format=lzma.FORMAT_XZ, check=lzma.CHECK_CRC64, filters=XZ_FILTERS))
    except (DetourError, OSError) as error:
        sys.exit(f"reference.py: {error}")


if __name__ == "__main__":
    main()