#!/usr/bin/env python3

import sys
import math

import numpy as np
import click
import struct
from PIL import Image

WINDOW_MAGIC = np.array([
    (0x00, 0x10, 0x20),
    (0x30, 0x40, 0x50),
    (0x60, 0x70, 0x80),
    (0x90, 0xa0, 0xb0),
    (0xc0, 0xd0, 0xe0),
    (0xf0, 0x40, 0x20),
    (0x50, 0x40, 0x30),
    (0x20, 0x10, 0x00),
])

def tb_int_to_px(val):
    return np.array([(val//0x10000)&0xff, (val//0x100)&0xff, val&0xff], dtype=np.uint8)

def data_encode(data, width, height):
    HEADER_W = 12
    assert width > HEADER_W
    data = struct.pack('>I', len(data)) + data
    data = np.frombuffer(data, dtype=np.uint8)
    for i in range(len(data) - 4):
        print(f'{data[i+4]:02x} ', end='')
        if i%16 == 15 and i > 0:
            print()

    data_expanded = np.repeat(data, 2).reshape((-1, 2))
    data_expanded[:, 0] &= 0xf0
    data_expanded[:, 1] = (data_expanded[:, 1] & 0x0f) << 4
    data_expanded = data_expanded.flatten()

    payload_w = width-HEADER_W
    line_capacity = payload_w*3
    print(payload_w, line_capacity, line_capacity / 2)
    block_height = math.ceil(len(data_expanded)/line_capacity)
    data_expanded = np.hstack([data_expanded, np.zeros(line_capacity*block_height - len(data_expanded), dtype=np.uint8)])
    block = np.zeros((block_height, payload_w, 3), dtype=np.uint8)
    block[:,:,:] = data_expanded.reshape((block_height, payload_w, 3))

    payload_area = np.vstack(
            [block] * (height // block_height) +
            [np.zeros((height % block_height, payload_w, 3), dtype=np.uint8)])

    header = np.zeros((height, HEADER_W, 3), dtype=np.uint8)
    for y in range(height):
        header[y, :8] = WINDOW_MAGIC
        if y % block_height == 0 and y + block_height <= height:
            header[y, 8] = tb_int_to_px(0x88)
        else:
            header[y, 8] = tb_int_to_px(0)
        header[y, 9] = tb_int_to_px(y)
        header[y, 10] = tb_int_to_px(width)
        header[y, 11] = tb_int_to_px(height)
    
    return Image.fromarray(np.hstack([header, payload_area]))

@click.command()
@click.option('-w', '--width', type=int, default=400)
@click.option('-h', '--height', type=int, default=300)
@click.argument('input_file', type=click.File('rb'))
@click.argument('output_file', type=click.Path(dir_okay=False, writable=True))
def tb_data_encode(width, height, input_file, output_file):
    data = input_file.read()
    img = data_encode(data, width, height)
    img.save(output_file)

if __name__ == '__main__':
    tb_data_encode()