#!/usr/bin/env python3

import argparse
import sys
import os.path
import subprocess
import xml.etree.ElementTree as xml

def train(files, output, zstd='zstd', maxdict=512):
    cmd = [ zstd, '--train', '--stdout', '--maxdict=' + str(maxdict), '-o', output ]

    for f in files:
        cmd.append(f.name)

    subprocess.run(cmd)

def compress(file, zstd='zstd', level=19, dictionary=None):
    cmd = [ zstd, '--compress', '--stdout', '-' + str(level) ]

    if dictionary is not None:
        cmd.append('-D')
        cmd.append(dictionary.name)
    
    cmd.append(file.name)
    return subprocess.run(cmd, capture_output=True).stdout

def hexdump(array_name, input_file, out_h):
    array_len = 0
    
    print("constexpr unsigned char {}[] = {{".format(array_name), file=out_h)

    for byte in input_file[0:len(input_file)]:
        array_len+=1
        if array_len%16 == 0:
            print(" 0x{:02X},".format(byte), file=out_h)
        else:
            print(" 0x{:02X},".format(byte), file=out_h, end='')


    print("};", file=out_h)
    print("constexpr size_t {}_len = {};\n".format(array_name, array_len), file=out_h)

def name(path):
    name = path.replace('/', '_')
    if name.endswith('.zstd'):
        name = name[:-5]
    name = name.replace('-', '_')
    name = name.replace('.', '_')
    return name

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description='Resource Compiler for C++',
        epilog='If using compression, make sure the required dependencies are provided.',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )

    parser.add_argument('input', type=argparse.FileType('rt'), help='input file (.xrc)')
    parser.add_argument('-o', '--output', type=argparse.FileType('wt'), metavar='OUT', default=sys.stdout, help='output header file')

    parser.add_argument('--compress', choices=[ 'None', 'Zstd' ], help='compress input files using algorightm')
    parser.add_argument('--dict', type=argparse.FileType('rb'), help='[zstd] use specified dictionary, recommended for many similar small files')
    parser.add_argument('--train', action='store_true', help='[zstd] train dictionary')

    args=parser.parse_args()

    entries_list = ""

    if args.compress=='Zstd' and args.train:
        train(args.input, args.dict.name)

    # write header
    print("// Autogenerated binary file hexdump", file=args.output)
    print("// This file may get overwritten by the build system\n", file=args.output)
    print("#include <embed.h>\n", file=args.output)

    # write file data
    for child in xml.parse(args.input).getroot():
        if child.tag == 'qresource':
            prefix = child.attrib['prefix']
            for i in child:
                vname = name(i.text)
                with open(i.text) as f:
                    hexdump(vname, compress(f, dictionary=args.dict), args.output)
                entries_list += "  {{ \"{}/{}\", std::span({}, {}_len) }},\n".format(prefix, i.attrib['alias'], vname, vname)

    # write dictionary
    if args.dict is not None:
        hexdump('dictionary', args.dict.read(), args.output)
    
    # write entries
    print("constexpr auto entries = frozen::make_unordered_map<frozen::string, std::span<const unsigned char>>({", file=args.output)
    print(entries_list, file=args.output)
    print("});\n", file=args.output)
    
    # write metadata struct
    print("constexpr embed::ResourceData metadata = {", file=args.output)
    print("    .compression = embed::{},".format(args.compress), file=args.output)
    if args.dict is not None:
        print("    .dictionary = std::span(dictionary, dictionary_len),", file=args.output)
    print("};", file=args.output)