#!/usr/bin/env python3

import sys

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)

@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()
    data = struct.pack('>I', len(data)) + data
    data = np.frombuffer(data, dtype=np.uint8)

    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()

    win = np.zeros([height*width*3], dtype=np.uint8)
    win[12*3:12*3+len(data_expanded)] = data_expanded
    win = win.reshape((height, width, 3))

    win[0, :8] = WINDOW_MAGIC
    win[0, 8] = tb_int_to_px(0)
    win[0, 9] = tb_int_to_px(0)
    win[0, 10] = tb_int_to_px(width)
    win[0, 11] = tb_int_to_px(height)
    
    print('out:', ' '.join(f'{x:02x}' for x in win.flatten()[:256]))

    Image.fromarray(win).save(output_file)

if __name__ == '__main__':
    tb_data_encode()