#!/usr/bin/env python

# This is a modified version of a program written by David Fifield <david@bamsoftware.com>
# http://www.bamsoftware.com/hacks/deflate.html
# That program is in the public domain.
#
# Modifications have been made to return data instead of write out a file
# and convert to python3

import struct
import zlib
import math

# For ZLIB and DEFLATE, see RFCs 1950 and 1951 respectively.


# An abstraction to build up byte strings from bits.
class BitBuffer:
    def __init__(self):
        self.buf = []
        self.work = 0
        self.p = 0

    def push(self, val, nbits):
        while nbits > self.lack():
            tmp = val % (1 << self.lack())
            self.work |= tmp << self.p
            nbits -= self.lack()
            val //= 1 << self.lack()
            self.buf.append(self.work)
            self.work = 0
            self.p = 0
        val %= 1 << nbits
        self.work |= val << self.p
        self.p += nbits
        if self.lack() == 0:
            self.buf.append(self.work)
            self.work = 0
            self.p = 0

    def push_rev(self, val, nbits):
        m = (1 << nbits) >> 1
        while m > 0:
            if val & m:
                self.push(0b1, 1)
            else:
                self.push(0b0, 1)
            m >>= 1

    def lack(self):
        return 8 - self.p

    def is_even(self):
        return self.p == 0

    def get_buf(self):
        return bytes(self.buf)


# For a given length, return the 3-tuple
#   (length code, extra bits, number of extra bits)
# This is based on the table in RFC 1951, section 3.2.5.
def length_code_for(n):
    if n < 3:
        return None, None, None
    if n < 11:
        baselen = 3
        basecode = 257
        nbits = 0
    elif n < 19:
        baselen = 11
        basecode = 265
        nbits = 1
    elif n < 35:
        baselen = 19
        basecode = 269
        nbits = 2
    elif n < 67:
        baselen = 35
        basecode = 273
        nbits = 3
    elif n < 131:
        baselen = 67
        basecode = 277
        nbits = 4
    elif n < 258:
        baselen = 131
        basecode = 281
        nbits = 5
    else:
        return 285, 258, 0

    m = 1 << nbits
    return (n - baselen) // m + basecode, nbits, (n - baselen) % m


# This is a specialized version of DEFLATE that can only compress strings
# containing only \x00 bytes.
#
# The return value is a 3-tuple (preamble, n, postamble). The deflated stream is
# constructed by writing the preamble, then n \x00 bytes, then the postamble.
def deflate_zeroes(n):
    # The rest of this algorithm doesn't work for 1032 bytes or less.
    # The png still has the correct size, but compression ratio is bad. The use case
    # is for compressing large files, so we are not worried about 1KB images.
    n = max(n, 1033)

    pre_bits = BitBuffer()

    # BFINAL    0 Not final block
    pre_bits.push(0b0, 1)
    # BTYPE    10 Dynamic Huffman codes
    pre_bits.push(0b10, 2)
    # HLIT  11101 (decimal 29) 286 literal/length codes
    pre_bits.push(0b11101, 5)
    # HDIST 00001 (decimal 1) 2 distance codes (we only use one of them)
    pre_bits.push(0b00001, 5)
    # HCLEN  1110 (decimal 14) 18 code length codes (we need access to
    #             the length-1 code).
    pre_bits.push(0b1110, 4)

    # Figure out how much is left over after compressing 258 zeroes at a time.
    # n - 1 because we must emit a literal byte first.
    num_excess = (n - 1) % 258
    excess_code, excess_nbits, excess_bits = length_code_for(num_excess)

    # The Huffman tree for code lengths. Give us code lengths for 0, 1, 2, 3,
    # and 18.
    #  00 -> 1
    #  01 -> 3
    #  10 -> 18
    # 110 -> 0
    # 111 -> 2
    pre_bits.push(0b000, 3)  # 16
    pre_bits.push(0b000, 3)  # 17
    pre_bits.push(0b010, 3)  # 18
    pre_bits.push(0b011, 3)  # 0
    pre_bits.push(0b000, 3)  # 8
    pre_bits.push(0b000, 3)  # 7
    pre_bits.push(0b000, 3)  # 9
    pre_bits.push(0b000, 3)  # 6
    pre_bits.push(0b000, 3)  # 10
    pre_bits.push(0b000, 3)  # 5
    pre_bits.push(0b000, 3)  # 11
    pre_bits.push(0b000, 3)  # 4
    pre_bits.push(0b000, 3)  # 12
    pre_bits.push(0b010, 3)  # 3
    pre_bits.push(0b000, 3)  # 13
    pre_bits.push(0b011, 3)  # 2
    pre_bits.push(0b000, 3)  # 14
    pre_bits.push(0b010, 3)  # 1

    # Now we need a literal/lengths Huffman tree with
    #   0 -> 285
    #  10 -> 0
    # 110 -> 256
    # 111 -> excess_code

    # Have to reverse the bit codes here. 'Huffman codes are packed starting
    # with the most-significant bit of the code.'
    pre_bits.push_rev(0b111, 3)  # 0: code length 2
    pre_bits.push_rev(0b10, 2)  # code length 18, 7 extra bits to follow
    pre_bits.push(0b1111111, 7)  # 1-138: code length 0 (127 + 11 entries)
    pre_bits.push_rev(0b10, 2)  # code length 18, 7 extra bits to follow
    pre_bits.push(0b1101010, 7)  # 139-255: code length 0 (106 + 11 entries)
    pre_bits.push_rev(0b01, 2)  # 256: code length 3
    if excess_code:
        if excess_code - 257 >= 11:
            pre_bits.push_rev(0b10, 2)  # code length 18, extra bits to follow
            pre_bits.push(excess_code - 257 - 11, 7)
        else:
            for i in range(257, excess_code):
                pre_bits.push_rev(0b110, 3)  # code length 0
        pre_bits.push_rev(0b01, 2)  # excess_code: code length 3
        if 285 - (excess_code + 1) >= 11:
            pre_bits.push_rev(0b10, 2)  # code length 18, extra bits to follow
            pre_bits.push(285 - (excess_code + 1) - 11, 7)  # 18
        else:
            for i in range(excess_code + 1, 285):
                pre_bits.push_rev(0b110, 3)  # code length 0
    else:
        pre_bits.push_rev(0b01, 2)  # Insert dummy at 257 to make the tree work.
        pre_bits.push_rev(0b10, 2)  # code length 18, 7 extra bits to follow
        pre_bits.push(285 - 258 - 11, 7)  # 18
    pre_bits.push_rev(0b00, 2)  # 285: code length 1

    # Now the distance Huffman tree. 'If only one distance code is used, it is
    # encoded using one bit, not zero bits; in this case there is a single code
    # length of one, with one unused code.' We store the unused code with a code
    # length of 1 instead of 0 because that's one bit shorter in the code
    # lengths tree we constructed.
    pre_bits.push_rev(0b00, 2)  # code length 1
    pre_bits.push_rev(0b00, 2)  # code length 1

    # Write a literal \x00.
    assert n >= 1
    pre_bits.push(0b01, 2)
    n -= 1

    # At this point our bit offset may be even or odd. Every block of 258 we
    # write takes two bits. We want to do byte operations for most of the data.
    # Push up to an even bit offset, but remember if we have to add an extra bit
    # later.
    #
    # This is the part that doesn't work for small blocks. We may push up to 7
    # bits here, representing up to 4 * 258 = 1032 zeroes.
    odd = False
    while not pre_bits.is_even():
        pre_bits.push(0b0, 1)
        if not odd:
            assert n >= 258
            n -= 258
        odd = not odd

    preamble = pre_bits.get_buf()

    num_zeroes = n // (258 * 4)
    n -= num_zeroes * (258 * 4)

    post_bits = BitBuffer()
    while n >= 258:
        post_bits.push(0b00, 2)
        n -= 258

    # Round out the odd bit if necessary.
    if odd:
        post_bits.push(0b0, 1)

    if num_excess < 3:
        for i in range(num_excess):
            post_bits.push(0b01, 2)  # literal \x00
    else:
        post_bits.push(0b111, 3)  # excess_code
        post_bits.push(excess_bits, excess_nbits)
        post_bits.push(0b0, 1)  # distance 1

    n -= num_excess
    assert n == 0

    post_bits.push(0b011, 3)  # 256, end of data

    # Round to a byte boundary with an uncompressed block.
    if not post_bits.is_even():
        post_bits.push(0b0, 1)
        post_bits.push(0b00, 2)
        while not post_bits.is_even():
            post_bits.push(0b0, 1)
        post_bits.push(0b00000000, 8)
        post_bits.push(0b00000000, 8)
        post_bits.push(0b11111111, 8)
        post_bits.push(0b11111111, 8)

    postamble = post_bits.get_buf()

    return preamble, num_zeroes, postamble


def adler32_zeroes(n, start=1):
    s1 = (start & 0xffff) % 65521
    s2 = ((start >> 16) & 0xffff) % 65521
    return (((s2 + n * s1) % 65521) << 16) + s1


class ZlibStream:
    ZEROES_1024 = b'\x00' * 1024

    def __init__(self):
        self.parts = []
        self.checksum = 1

    def push_zeroes(self, n):
        self.checksum = adler32_zeroes(n, self.checksum)
        self.parts.append(deflate_zeroes(n))

    def length(self):
        l = 2  # Header.
        for part in self.parts:
            l += len(part[0]) + part[1] + len(part[2])
        l += 5  # Final empty block.
        l += 4  # Checksum.
        return l

    def out(self):
        yield struct.pack('BB', 0b01111000, 0b11011010)
        for i in range(len(self.parts)):
            part = self.parts[i]
            if type(part) == bytes:
                yield part
            else:
                preamble, n, postamble = part
                yield preamble
                while n >= 1024:
                    yield self.ZEROES_1024
                    n -= 1024
                yield b'\x00' * n
                yield postamble
        # Final empty non-compressed block with BFINAL=1.
        yield b'\x01\x00\x00\xff\xff'
        yield struct.pack('>L', self.checksum & 0xffffffff)

# PNG code.


def chunk(type, data):
    crc = zlib.crc32(type + data) & 0xffffffff
    return struct.pack('>L', len(data)) + type + data + struct.pack('>L', crc)


def chunk_IHDR(width, height, depth, color):
    data = struct.pack('>LLBBBBB', width, height, depth, color, 0, 0, 0)
    return chunk(b'IHDR', data)


def chunk_PLTE(rgb):
    data = b''.join(struct.pack('BBB', *x) for x in rgb)
    return chunk(b'PLTE', data)


def max_deflate_png(width, height):
    # PNG header.
    out = b'\x89\x50\x4e\x47\x0d\x0a\x1a\x0a'
    out += chunk_IHDR(width, height, 1, 3)
    out += chunk_PLTE(((0x40, 0x00, 0xe0), (0xe0, 0x00, 0x40)))

    zlib_data = ZlibStream()

    data_width = int(math.ceil(width / 8))  # 8 pixels per byte, rounded up
    zlib_data.push_zeroes((1 + data_width) * height)

    out += struct.pack('>L', zlib_data.length())
    crc = zlib.crc32(b'IDAT')
    out += b'IDAT'
    for x in zlib_data.out():
        crc = zlib.crc32(x, crc)
        out += x
    out += struct.pack('>L', crc & 0xffffffff)

    out += chunk(b'IEND', b'')
    return out
