/*
 * src/unpack.c
 * https://gitlab.com/bztsrc/assetpack
 *
 * Copyright (C) 2024 bzt, MIT license
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to
 * deal in the Software without restriction, including without limitation the
 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
 * sell copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL ANY
 * DEVELOPER OR DISTRIBUTOR BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
 * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
 * IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 *
 * @brief Main source file for unpacking
 */

#include <png.h>
#define STB_IMAGE_IMPLEMENTATION
#define APCK_IMPLEMENTATION
#include <apck.h>
#include "main.h"
#include "image.h"
#include "aes.h"

/* we need this to save the sprites */
uint8_t *apck_pack_png(uint32_t *buf, int w, int h, int p, uint64_t *size);

static char *dirnames[] = { "fonts", "bgm", "sfx", "videos", "models", "images", "atlases", "sprites", "maps", "reserved41",
    "reserved42", "reserved43", "reserved44", "reserved45", "objects", "scripts", "data01", "data02", "data03", "data04",
    "data05", "data06", "data07", "data08", "data09", "data10", "data11", "data12", "data13", "data14", "data15", "data16" };

static apck_t ctx;

/**
 * Dump Asset Pack archive
 */
void dump(void)
{
    apck_hdr_t hdr;
    uint8_t *ptr, *end, *tmp;
    uint64_t i, n, offs, size, orig, name;
    char *str, fn[32];
    int j;

    if(!ctx.archives[0].f) return;
    apck_fileseek(ctx.archives[0].f, 0);
    if(!apck_fileread(ctx.archives[0].f, &hdr, sizeof(hdr))) return;

    printf("Asset Pack Header\n");
    printf("  Magic:            %c%c%c%c\n", hdr.magic[0], hdr.magic[1], hdr.magic[2], hdr.magic[3]);
    printf("  Header size:      %u\n", hdr.hdr_size);
    printf("  Engine:           ");
    for(i = 0; i < 12 && hdr.engine[i]; i++) printf("%c", hdr.engine[i]);
    printf("\n  Engine version:   %u.%u.%u.%u\n",
        (hdr.version >> 24) & 0xff, (hdr.version >> 16) & 0xff, (hdr.version >> 8) & 0xff, hdr.version & 0xff);
    printf("  Asset Dir size:   %u\n", hdr.size);
    printf("  Checksum:         %08x\n", hdr.chksum);
    printf("  Encryption key:  ");
    for(i = 0; i < 32; i++) printf(" %02X", hdr.enckey[i]);
    printf("\n");

    str = (char*)ctx.archives[0].str;
    ptr = ctx.archives[0].str + *((uint32_t*)ctx.archives[0].str);
    end = ctx.archives[0].str + hdr.size;

    for(j = 0; j < 64 && ptr < end; j++) {
        ptr = apck_uleb128(ptr, &n);
        if(n) printf("\nAsset Directory #%u (%s, %u entries)\n"
            "  Offset     Size        Compressed  Name\n"
            "  ---------- ----------- ----------- ----------------------\n",
            j, j < APCK_FONT ? "locales" : dirnames[j - APCK_FONT], (uint32_t)n);
        for(i = 0; i < n && ptr < end; i++) {
            ptr = apck_uleb128(ptr, &name);
            ptr = apck_uleb128(ptr, &offs);
            ptr = apck_uleb128(ptr, &orig);
            ptr = apck_uleb128(ptr, &size);
            sprintf(fn, "%04u", (uint32_t)i + 1);
            printf("  %02x%08x%12u%12u %s", (uint32_t)((offs + hdr.hdr_size) >> 32), (uint32_t)(offs + hdr.hdr_size),
                (uint32_t)orig, (uint32_t)size, name > 4 ? str + name : fn);
            if(j < APCK_FONT && i && (tmp = apck_readbuf(&ctx, 0, offs + hdr.hdr_size, 4))) {
                printf(" (%s)", !memcmp(tmp, "OggS", 4) ? "voice" : (!memcmp(tmp, "\x89PNG", 4) ? "image" : "???"));
                free(tmp);
            }
            printf("\n");
        }
    }
    printf("\n");
}

/**
 * Check if buffer contains valid UTF-8 string
 */
int validutf8(uint8_t *buf, uint64_t size)
{
    uint8_t *s = buf, *e = buf + size;
    while(s < e) {
        if(s[0] < ' ' && s[0] != '\t' && s[0] != '\r' && s[0] != '\n') return 0;
        if(!(s[0] & 0x80)) s++; else
        if((s[0] & 0xE0) == 0xC0 && (s[1] & 0xC0) == 0x80) s += 2; else
        if((s[0] & 0xF0) == 0xE0 && (s[1] & 0xC0) == 0x80 && (s[2] & 0xC0) == 0x80) s += 3; else
        if((s[0] & 0xF8) == 0xF0 && (s[1] & 0xC0) == 0x80 && (s[2] & 0xC0) == 0x80 && (s[3] & 0xC0) == 0x80) s += 4;
        else return 0;
    }
    return 1;
}

/**
 * The main unpacker routine
 */
int unpack(char *input, char *output)
{
    apck_map_t *map;
    apck_sprite_t *sprite;
    apck_file_t *file = NULL;
    uint8_t *buf, *comp, *po;
    uint64_t size;
    char tmp[PATH_MAX + FILENAME_MAX + 1], fn[32];
    uint32_t pathlen, i, j, k;
    int s, d, n, w, h;

    if(!input) return 0;

    apck_init(&ctx, NULL, 0, aes_init, aes_read);
    if(apck_load(&ctx, input) == APCK_OK) {
        if(!output) dump();
        else {
            /*
             * NOTE: this code is a bit messy, because it accesses ctx directly for iterations. Normally you never
             * want *all* the assets at once, so the APCK API, which hides all is ugliness from you, is focused on
             * retriving one asset a time. That's nice, but not very efficient here where we really need all assets.
             */
            /* create destination directory */
            strcpy(tmp, output);
            pathlen = strlen(tmp);
            if(pathlen && (tmp[pathlen - 1] == '/' || tmp[pathlen - 1] == '\\')) tmp[--pathlen] = 0;
            apck_mkdir(tmp);

            /* application icon */
            if((buf = apck_asset(&ctx, APCK_IMAGE, "icon", &size))) {
                strcpy(tmp + pathlen, "/icon.png");
                if(!putfile(tmp, buf, size))
                    fprintf(stderr, "unable to write '%s'\n", tmp);
            }

            /* locales */
            for(j = 0; j < APCK_FONT; j++)
                if(ctx.files[j] && ctx.numfiles[j]) {
                    file = ctx.files[j];
                    strcpy(tmp + pathlen, "/locales");
                    apck_mkdir(tmp);
                    strcat(tmp + pathlen + 1, "/");
                    strcat(tmp + pathlen + 1, file->name);
                    k = strlen(tmp);

                    /* localized strings are already loaded and uncompressed (because that's needed to load the msgstr index) */
                    if(file->buf) {
                        /* save as a .po file */
                        if((po = po_encode(file->buf, file->size, &size))) {
                            strcat(tmp + pathlen + 1, ".po");
                            if(!putfile(tmp, po, size))
                                fprintf(stderr, "unable to write '%s'\n", tmp);
                            free(po);
                        } else
                            fprintf(stderr, "unable to encode '%s.po'\n", tmp);
                        free(file->buf); file->buf = NULL;
                        tmp[k] = 0;
                    }

                    /* then create a subdirectory and save voices and images with texts for this locale */
                    if(ctx.numfiles[j] > 1) {
                        apck_mkdir(tmp);
                        for(i = 1, file = &ctx.files[j][1]; i < ctx.numfiles[j]; i++, file++) {
                            /* read in, unencrypt and save data */
                            if(file->size && (buf = apck_readbuf(&ctx, file->archive, file->offs, file->size))) {
                                strcpy(tmp + k, "/");
                                if(file->name && *file->name) strcpy(tmp + k + 1, file->name);
                                else sprintf(tmp + k + 1, "%04u", i + 1);
                                if(!memcmp(buf, "OggS", 4)) strcat(tmp, ".ogg"); else
                                if(!memcmp(buf, "\x89PNG", 4)) strcat(tmp, ".png"); else
                                    strcat(tmp, ".bin");
                                if(!putfile(tmp, buf, file->size))
                                    fprintf(stderr, "unable to write '%s'\n", tmp);
                                free(buf);
                            } else { fprintf(stderr, "unable to read asset '%s'\n", file->name); }
                        }
                    }
                }

            /* all the other types (except sprites) */
            for(; j < APCK_NUMTYPES; j++)
                if(j != APCK_SPRITE && ctx.files[j] && ctx.numfiles[j]) {
                    for(i = 0, file = ctx.files[j]; i < ctx.numfiles[j]; i++, file++) {
                        /* we have already saved the application icon */
                        if(j == APCK_IMAGE && !strcmp(file->name, "icon")) continue;
                        /* read in, unencrypt and uncompress data */
                        if(file->size && (comp = apck_readbuf(&ctx, file->archive, file->offs, file->comp ? file->comp : file->size))) {
                            if(file->comp) {
                                if((buf = (uint8_t*)malloc(file->size))) {
                                    if(stbi_zlib_decode_buffer((char*)buf, (int)file->size, (char*)comp, (int)file->comp) < 1) {
                                        fprintf(stderr, "unable to uncompress asset '%s'\n", file->name);
                                        free(buf);
                                    } else
                                        file->buf = buf;
                                } else { fprintf(stderr, "unable to allocate asset '%s' (%u bytes)\n", file->name, (uint32_t)file->size); }
                                free(comp);
                            } else file->buf = comp;
                        } else { fprintf(stderr, "unable to read asset '%s'\n", file->name); }

                        /* if we have retrieved the raw asset */
                        if(file->buf) {
                            /* create subdirectory */
                            strcpy(tmp + pathlen, "/"); strcpy(tmp + pathlen + 1, dirnames[j - APCK_FONT]);
                            apck_mkdir(tmp);
                            strcat(tmp + pathlen + 1, "/");
                            sprintf(fn, "%04u", i + 1);
                            strcat(tmp + pathlen + 1, file->name && *file->name ? file->name : fn);
                            /* convert asset if necessary */
                            switch(j) {
                                case APCK_MAP:
                                    if((map = apck_map_decode(ctx.archives[file->archive].str, file->buf, file->size))) {
                                        free(file->buf);
                                        file->buf = tmx_encode(map, &file->size);
                                        strcat(tmp, ".tmx");
                                    } else {
                                        fprintf(stderr, "unable to decode map '%s'\n", file->name);
                                        strcat(tmp, ".bin");
                                    }
                                break;
                                case APCK_FONT: strcat(tmp, ".sfn"); break;
                                case APCK_MODEL: strcat(tmp, ".m3d"); break;
                                case APCK_VIDEO: strcat(tmp, ".ogv"); break;
                                default:
                                    if(!memcmp(file->buf, "OggS", 4)) strcat(tmp, ".ogg"); else
                                    if(!memcmp(file->buf, "\x89PNG", 4)) strcat(tmp, ".png"); else
                                    /* this is rudimentary, not that the actual extension matters to Asset Pack */
                                    if(file->buf[0] == '<') strcat(tmp, ".xml"); else
                                    if(file->buf[0] == '{' || file->buf[0] == '[') strcat(tmp, ".json"); else
                                    if(validutf8(file->buf, file->size)) strcat(tmp, ".txt"); else
                                        strcat(tmp, ".bin");
                                break;
                            }
                            /* write out converted asset */
                            if(file->buf) {
                                if(!putfile(tmp, file->buf, file->size))
                                    fprintf(stderr, "unable to write '%s'\n", tmp);
                                if(j == APCK_ATLAS && (buf = (uint8_t*)stbi_load_from_memory(file->buf, file->size, &file->w, &file->h, &s, 4))) {
                                    free(file->buf); file->buf = buf;
                                } else {
                                    free(file->buf); file->buf = NULL;
                                }
                            }
                        }
                    }
                }

            /* extract sprites from atlases */
            for(i = 0, j = APCK_SPRITE, file = ctx.files[j]; i < ctx.numfiles[j]; i++, file++) {
                sprite = NULL;
                /* read in, unencrypt and uncompress data */
                if(file->comp && (comp = apck_readbuf(&ctx, file->archive, file->offs, file->comp))) {
                    if((buf = (uint8_t*)malloc(file->size))) {
                        if(stbi_zlib_decode_buffer((char*)buf, (int)file->size, (char*)comp, (int)file->comp) < 1)
                            fprintf(stderr, "unable to uncompress asset '%s'\n", file->name);
                        else
                            sprite = apck_sprite_decode(ctx.archives[file->archive].aidx, ctx.archives[file->archive].str, buf, file->size);
                        free(buf);
                    } else { fprintf(stderr, "unable to allocate asset '%s' (%u bytes)\n", file->name, (uint32_t)file->size); }
                    free(comp);
                } else { fprintf(stderr, "unable to read asset '%s'\n", file->name); }

                /* if we have retrieved the raw asset */
                if(sprite) {
                    strcpy(tmp + pathlen, "/"); strcpy(tmp + pathlen + 1, dirnames[j - APCK_FONT]);
                    apck_mkdir(tmp);
                    strcat(tmp + pathlen + 1, "/");
                    sprintf(fn, "%04u", i + 1);
                    strcat(tmp + pathlen + 1, file->name && *file->name ? file->name : fn);
                    if(sprite->type != APCK_ANIM_LOOP)
                        strcat(tmp + pathlen + 1, loopext[sprite->type % 5]);
                    k = strlen(tmp);
                    /* foreach direction */
                    for(d = 0; d < APCK_NUMDIR; d++)
                        if(sprite->frames[d] && sprite->nf[d]) {
                            /* calculate pixel buffer size */
                            for(n = 1; n * n < sprite->nf[d]; n++);
                            w = sprite->w * n;
                            h = sprite->h * ((sprite->nf[d] + n - 1) / n);
                            if((buf = (uint8_t*)malloc(w * h * 4))) {
                                memset(buf, 0, w * h * 4);
                                /* place each frame on the pixel buffer */
                                for(s = 0; s < sprite->nf[d]; s++)
                                    apck_sprite_pose(&ctx, sprite, d, s, (uint32_t *)buf, (s % n) * sprite->w, (s / n) * sprite->h,
                                        w, h, w * 4);
                                /* compress to png and save */
                                if((comp = apck_pack_png((uint32_t*)buf, w, h, w * 4, &size))) {
                                    strcpy(tmp + k, dirext[d]);
                                    strcat(tmp + k, ".png");
                                    if(!putfile(tmp, comp, size))
                                        fprintf(stderr, "unable to write '%s'\n", tmp);
                                    free(comp);
                                } else
                                    fprintf(stderr, "unable to allocate encode png for '%s'\n", file->name);
                                free(buf);
                            } else
                                fprintf(stderr, "unable to allocate pixel buffer for '%s' (%u bytes)\n", file->name, w * h * 4);
                        }
                }
            }
        }
    } else
        fprintf(stderr, "unable to load archive\n");
    apck_free(&ctx);
    return 1;
}

