import zlib
import hashlib
from pathlib import Path
import click

def print_carray(f, payload):
    while len(payload) > 0:
        f.write('\n    ')
        f.write(', '.join('0x{:02x}'.format(x) for x in payload[0:16]))
        f.write(',')
        payload = payload[16:]
    f.write('\n')


@click.command()
@click.argument('dpath', type=click.Path(exists=True), default='.')
@click.option('--sd', is_flag=True, default=False, help='Generate header for using with SDCard')
def main(dpath, sd):
    """
    This script takes a directory path (default is current directory) and
    create compressed .bin.gz file along 'esp_binaries.h' containing metadata in esp32_zipfile_t struct.
    Which can be used to program ESP32 with TestBed or Brain.
    If '--sd' option is specified, it will skip generating C array data for gz file.
    """
    output_header = 'esp_binaries.h'
    with open(output_header, 'w') as fc:
        # print typedef struct
        fc.write('// Generated by tools/esp_compress.py\n')
        file_list = sorted(Path(dpath).glob('*.bin'))
        fc.write(f'#define ESP_BINARIES_COUNT ({len(file_list)})\n\n')

        # print list of struct first
        for fname in file_list:
            var = fname.stem
            var = var.replace('.', '_')
            var = var.replace('-', '_')
            fc.write(f'// const esp32_zipfile_t {var}\n')
        fc.write('\n')

        for fname in file_list:
            with open(fname, 'rb') as fi:
                image = fi.read()
                zimage = zlib.compress(image, 9)
                fzname = fname.with_suffix('.bin.gz')
                md5 = hashlib.md5(image)

                # write .gz file
                with open(fzname, 'wb') as fz:
                    fz.write(zimage)

                # write to c header file
                var = fname.stem
                var = var.replace('.', '_')
                var = var.replace('-', '_')

                # bin gz contents
                if not sd:
                    fc.write(f'const uint8_t _{var}_gz[{len(zimage)}] = {{')
                    print_carray(fc, zimage)
                    fc.write('};\n\n')

                fc.write(f'const esp32_zipfile_t {var} = {{\n')
                fc.write(f'  .name = "{fzname}",\n')
                if sd:
                    fc.write(f'  .data = NULL,\n')
                else:
                    fc.write(f'  .data = _{var}_gz,\n')
                fc.write(f'  .compressed_len = {len(zimage)},\n')
                fc.write(f'  .uncompressed_len = {len(image)},\n')
                fc.write('  .md5 = {')
                print_carray(fc, md5.digest())
                fc.write('  },\n')
                fc.write('};\n')
                fc.write('\n')

                deflation = 100*(1-len(zimage)/len(image))
                print(f"Compressed {var}: {len(image)} to {len(zimage)} bytes, deflation: {deflation:.2f}%")


if __name__ == '__main__':
    main()