import logging
import argparse
import struct
import os.path

logger = None


def main():
    global logger
    _ch = logging.StreamHandler()
    _ch.setLevel("DEBUG")
    _formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    _ch.setFormatter(_formatter)
    logger = logging.getLogger(__name__)
    logger.addHandler(_ch)
    logger.setLevel("DEBUG")
    parser = argparse.ArgumentParser()
    parser.add_argument("gzp_file_directory", help="gzp file or directory to read")
    parser.add_argument("output_dir", help="output directory to send all extracted files")
    parser.add_argument("--extensions", help="extensions to extract, separated by a comma (ex: 'tga,wav' will extract only .tga and .wav files)")
    args = parser.parse_args()
    gzp_file = args.gzp_file_directory
    output_dir = args.output_dir
    extensions = args.extensions

    if extensions:
        extensions = extensions.split(",")

    if os.path.isdir(gzp_file):
        logger.info("Reading .gzp files in "+gzp_file)
        files = os.listdir(gzp_file)
        for file in files:
            if file.endswith(".gzp"):
                extract_gzp(os.path.join(gzp_file, file), extensions, output_dir)
    else:
        extract_gzp(gzp_file, extensions, output_dir)


def extract_gzp(gzp_file, extensions, output_dir):
    logger.info("Opening " + gzp_file)
    with open(gzp_file, "rb") as gzp_fp:
        checksum = read_int(gzp_fp)

        if checksum != 0x6608F101:
            pass
            #raise Exception("Invalid GZP checksum")

        logger.info("Checksum OK")

        meta_info_offset = read_int(gzp_fp)

        gzp_fp.seek(meta_info_offset)
        unk = read_int(gzp_fp)
        entries_count = read_int(gzp_fp)

        if entries_count == 0:
            logger.info("No entries found, skipping")
            return

        logger.info(str(entries_count) + " entries in GZP")

        for index in range(entries_count):
            logger.info("Reading index " + str(index))
            compressed_size = read_int(gzp_fp)
            original_size = read_int(gzp_fp)
            file_time = read_int(gzp_fp)
            content_offset = read_int(gzp_fp) + 16
            compression = read_byte(gzp_fp)  # compression: 1 if compressed else 0
            name_length = read_byte(gzp_fp)
            name = read_bytes(gzp_fp, name_length).decode("utf8").strip('\x00')
            logger.info(name + " compression: " + str(compression) + ", filesize: " + sizeof_fmt(
                original_size) + ", start: " + str(content_offset))

            if extensions:
                extract_file = False
                for extension in extensions:
                    if name.endswith(extension):
                        extract_file = True
                        break
                if not extract_file:
                    logger.info("File " + name + " does not match any of wanted extension. Skipping...")
                    continue

            curr_pos = gzp_fp.tell()
            gzp_fp.seek(content_offset)
            buffer = gzp_fp.read(compressed_size)
            gzp_fp.seek(curr_pos)

            if compression == 1:
                logger.info("File is compressed, decompressing it")
                buffer = decompress(buffer, original_size)

            logger.info("Writing file " + name)
            if not os.path.exists(output_dir):
                os.mkdir(output_dir)
            with open(os.path.join(output_dir, name), "wb") as entry_fp:
                entry_fp.write(buffer)


def sizeof_fmt(num, suffix='B'):
    for unit in ['', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi']:
        if abs(num) < 1024.0:
            return "%3.1f%s%s" % (num, unit, suffix)
        num /= 1024.0
    return "%.1f%s%s" % (num, 'Yi', suffix)


def read_byte(fp):
    return struct.unpack("<B", fp.read(1))[0]


def read_bytes(fp, num):
    return bytes(struct.unpack('<'+str(num)+'B', fp.read(num)))


def read_int(fp):
    return struct.unpack("<L", fp.read(4))[0]


def decompress(buffer, finalsize):
    i = 0
    j = 0
    dec_byte = 0
    dec_bits = 8
    buff_start = 0xFEE
    res = bytearray(finalsize)

    if finalsize == 0:
        return res

    while j < finalsize:
        if dec_bits == 8:
            dec_byte = buffer[i]
            i += 1
            dec_bits = 0

        if (dec_byte >> dec_bits & 1) == 0:
            dec_pos = ((buffer[i] + ((buffer[i + 1] & 0xF0) << 4) - buff_start - j) & 0xFFF) - 0x1000 + j
            dec_len = (buffer[i + 1] & 0xF) + 3
            i += 2

            while dec_len > 0:
                if dec_pos >= 0:
                    res[j] = res[dec_pos]
                else:
                    res[j] = 32

                j += 1
                dec_pos += 1
                dec_len -= 1
        else:
            res[j] = buffer[i]
            i += 1
            j += 1

        dec_bits += 1

    return res


if __name__ == '__main__':
    main()