"""LZO1X, as used by the game's compressed UE3 packages (docs/formats/upk.md).

`decompress` is a pure-Python LZO1X decoder (same semantics as lzo1x_decompress_safe; the packages
use 128 KiB blocks, so speed is fine for block-level access). `compress_stored` writes a valid LZO1X
stream that holds the data as one literal run (no matches): bigger than real compression, but any
LZO1X decoder (the game's included) reads it back exactly, and nothing else has to change.
"""
M2_MAX_OFFSET = 0x0800


class LzoError(ValueError):
    pass


def _copy_match(out, dist, n):
    start = len(out) - dist
    if start < 0:
        raise LzoError("match distance beyond output start")
    if dist >= n:
        out += out[start:start + n]
    else:  # overlapping: repeat the last `dist` bytes
        pat = bytes(out[start:])
        out += (pat * (n // dist + 1))[:n]


def decompress(src, out_len=None):
    """LZO1X block -> bytes. `out_len` (if given) is checked."""
    src = bytes(src)
    n = len(src)
    out = bytearray()
    ip = 0

    def ext(t, base):
        nonlocal ip
        while src[ip] == 0:
            t += 255
            ip += 1
        t += base + src[ip]
        ip += 1
        return t

    try:
        state = "loop"
        if src[0] > 17:
            t = src[0] - 17
            ip = 1
            if t < 4:
                state = "match_next"
            else:
                out += src[ip:ip + t]
                ip += t
                state = "first_literal_run"
        while True:
            if state == "loop":
                t = src[ip]
                ip += 1
                if t >= 16:
                    state = "match"
                    continue
                if t == 0:
                    t = ext(0, 15)
                out += src[ip:ip + t + 3]
                ip += t + 3
                state = "first_literal_run"
                continue
            if state == "first_literal_run":
                t = src[ip]
                ip += 1
                if t >= 16:
                    state = "match"
                    continue
                dist = 1 + M2_MAX_OFFSET + (t >> 2) + (src[ip] << 2)
                ip += 1
                _copy_match(out, dist, 3)
                state = "match_done"
                continue
            if state == "match_next":
                out += src[ip:ip + t]
                ip += t
                t = src[ip]
                ip += 1
                state = "match"
                continue
            if state == "match":
                if t >= 64:                       # M2
                    dist = 1 + ((t >> 2) & 7) + (src[ip] << 3)
                    ip += 1
                    _copy_match(out, dist, (t >> 5) + 1)
                elif t >= 32:                     # M3
                    t &= 31
                    if t == 0:
                        t = ext(0, 31)
                    dist = 1 + ((src[ip] | (src[ip + 1] << 8)) >> 2)
                    ip += 2
                    _copy_match(out, dist, t + 2)
                elif t >= 16:                     # M4 (distance 0 = end of stream)
                    hi = (t & 8) << 11
                    t &= 7
                    if t == 0:
                        t = ext(0, 7)
                    dist = hi + ((src[ip] | (src[ip + 1] << 8)) >> 2)
                    ip += 2
                    if dist == 0:
                        break
                    _copy_match(out, dist + 0x4000, t + 2)
                else:                             # M1 after a match: 2 bytes
                    dist = 1 + (t >> 2) + (src[ip] << 2)
                    ip += 1
                    _copy_match(out, dist, 2)
                state = "match_done"
                continue
            # match_done: the low 2 bits of the instruction's last byte = trailing literals
            t = src[ip - 2] & 3
            if t == 0:
                state = "loop"
            else:
                state = "match_next"
    except IndexError:
        raise LzoError("truncated LZO stream") from None
    if ip != n:
        raise LzoError("%d bytes after the end of the LZO stream" % (n - ip))
    if out_len is not None and len(out) != out_len:
        raise LzoError("LZO block decoded to %d bytes, expected %d" % (len(out), out_len))
    return bytes(out)


EOF = b"\x11\x00\x00"


def compress_stored(data):
    """bytes -> LZO1X stream made of one literal run (decodes back to `data`)."""
    data = bytes(data)
    n = len(data)
    if n == 0:
        return EOF
    if n <= 238:
        return bytes([n + 17]) + data + EOF
    # long literal run: 0x00, k zero bytes (255 each), last byte; length = 18 + 255*k + last
    r = n - 18
    k = (r - 1) // 255
    last = r - 255 * k
    return b"\x00" + b"\x00" * k + bytes([last]) + data + EOF