/*
 * src/pack.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 Packer routine and its callbacks
 */

#include <time.h>
#include <vorbis/vorbisenc.h>
#include <png.h>
#define DBG(a) do{(printf a);fflush(stdout);}while(0)
#define APCK_PACK_IMPLEMENTATION
#include <apck_pack.h>
#include "main.h"
#include "aes.h"
#include "image.h"

static apck_pack_t ctx;
int category;

/* it sucks so much that we can't include stb_image in ANSI C */
char *stbi_zlib_decode_malloc_guesssize_headerflag(const char *buffer, int len, int initial_size, int *outlen, int parse_header);

/**
 * Callback to add audio files
 */
int audiocb(char *path)
{
    uint8_t *buf, *data;
    uint64_t size;
    char *ext;

    if((buf = getfile(path, &size))) {
        if((data = audio_decode(buf, size, 0, &size))) {
            if((ext = strrchr(path + pathlen, '.'))) *ext = 0;
            apck_pack_asset(&ctx, category, path + pathlen, data, size, 0, 0);
            free(data);
        } else
        if(!quiet) fprintf(stderr, "not an audio '%s'\n", path);
        free(buf);
    }
    return 1;
}

/**
 * Callback to add images (and voices too for locales, a little hack)
 */
int imagecb(char *path)
{
    uint8_t *buf, *data;
    uint32_t *pix;
    uint64_t size;
    int w, h;
    char *ext;

    if((buf = getfile(path, &size))) {
        if((pix = image_decode(buf, size, &w, &h))) {
            if((ext = strrchr(path + pathlen, '.'))) *ext = 0;
            apck_pack_asset(&ctx, category, path + pathlen, pix, w * h * 4, w, h);
            free(pix);
        } else
        if(category < APCK_FONT && (data = audio_decode(buf, size, 0, &size))) {
            if((ext = strrchr(path + pathlen, '.'))) *ext = 0;
            apck_pack_asset(&ctx, category, path + pathlen, data, size, 0, 0);
            free(data);
        } else
        if(!quiet) fprintf(stderr, "not an image '%s'\n", path);
        free(buf);
    }
    return 1;
}

/**
 * Callback to add sprites (will be added to atlas)
 */
int spritecb(char *path)
{
    uint8_t *buf;
    uint32_t *pix;
    uint64_t size;
    int w, h;
    char *ext;

    if((buf = getfile(path, &size))) {
        if((pix = image_decode(buf, size, &w, &h))) {
            if((ext = strrchr(path + pathlen, '.'))) *ext = 0;
            if(image_add(path + pathlen, pix, w, h))
                fprintf(stderr, "unable to allocate memory for sprite '%s'\n", path + pathlen);
        } else
        if(!quiet) fprintf(stderr, "not an image '%s'\n", path);
        free(buf);
    }
    return 1;
}

/**
 * Callback to add translated texts
 */
int localecb(char *path)
{
    uint8_t *buf, *data;
    uint64_t size;
    int pl;
    char *ext;

    if(category >= APCK_FONT) return 0;
    if((buf = getfile(path, &size))) {
        if((data = po_decode(buf, size, &size))) {
            if((ext = strrchr(path + pathlen, '.'))) *ext = 0;
            if(!apck_pack_asset(&ctx, category, path + pathlen, data, size, 0, 0)) {
                pl = pathlen;
                find(path, 1, imagecb);
                pathlen = pl;
                category++;
            }
        } else
        if(!quiet) fprintf(stderr, "not a po '%s'\n", path);
        free(buf);
    }
    return 1;
}

/**
 * Callback to add maps
 */
int mapcb(char *path)
{
    apck_map_t *map;
    uint8_t *buf;
    uint64_t size;
    char *ext;

    if((buf = getfile(path, &size))) {
        if((map = tmx_decode(buf, size))) {
            if((ext = strrchr(path + pathlen, '.'))) *ext = 0;
            apck_pack_asset(&ctx, APCK_MAP, path + pathlen, map, sizeof(apck_map_t), 0, 0);
            apck_map_free(map);
        } else
        if(!quiet) fprintf(stderr, "not a tmx '%s'\n", path);
        free(buf);
    }
    return 1;
}

/**
 * Callback for font files
 */
int fontcb(char *path)
{
    uint8_t *buf, *data, *ptr;
    uint64_t size;
    uint32_t len;
    char *ext;
    int r, c;

    if((buf = getfile(path, &size))) {
        /* if it's gzip compressed */
        if(buf[0] == 0x1f && buf[1] == 0x8b) {
            memcpy(&len, buf + size - 4, 4);
            ptr = buf + 3;
            c = *ptr++; ptr += 6;
            if(c & 4) { r = *ptr++; r += (*ptr++ << 8); ptr += r; }
            if(c & 8) { while(*ptr++ != 0); }
            if(c & 16) { while(*ptr++ != 0); }
            r = size - (ptr - buf);
            if((data = (unsigned char*)stbi_zlib_decode_malloc_guesssize_headerflag((const char*)ptr, r, 4096, &r, 0))) {
                free(buf);
                buf = data;
                size = r;
            } else {
                free(buf);
                if(!quiet) fprintf(stderr, "unable to uncompress '%s'\n", path + pathlen);
                return 1;
            }
        }
        if((data = font_decode(path + pathlen, buf, size, &len))) {
            if((ext = strrchr(path + pathlen, '.'))) *ext = 0;
            if(apck_pack_asset(&ctx, APCK_FONT, path + pathlen, data, len, 0, 0))
                free(data);
        } else
        if(!quiet) fprintf(stderr, "not a font '%s'\n", path);
        free(buf);
    }
    return 1;
}

/**
 * Callback for model files (might have inlined textures)
 */
int modelcb(char *path)
{
    uint8_t *buf, *data = NULL, *strs = NULL, *model, *ptr, *dst;
    uint64_t size;
    uint32_t len, s, si_s;
    char *ext;

    if((buf = getfile(path, &size))) {
        /* check file */
        if(memcmp(buf, "3DMO", 4)) goto err;
        ptr = buf + 8; size -= 8;
        /* skip preview */
        if(!memcmp(ptr, "PRVW", 4)) { len = *((uint32_t*)(ptr + 4)); ptr += len; size -= len; }
        /* check if its compressed */
        if(memcmp(ptr, "HEAD", 4)) {
            data = (unsigned char *)stbi_zlib_decode_malloc_guesssize_headerflag((const char*)ptr, size, 4096, (int*)&len, 1);
            if(!data || len < 20 || memcmp(data, "HEAD", 4)) {
err:            if(!quiet) fprintf(stderr, "not a model '%s'\n", path);
                if(data) free(data);
                free(buf);
                return 1;
            }
            ptr = data;
            size = len;
        }
        /* get string index size and string table */
        si_s = 1 << ((ptr[12] >> 4) & 3);
        strs = ptr + 16;
        /* copy chunks, but not inlined image assets, we add those as real image assets */
        len = 12 + size;
        if((model = (uint8_t*)malloc(len))) {
            memset(model, 0, len);
            memcpy(model, "3DMO", 4);
            dst = model + 8;
            while(size > 4 && memcmp(ptr, "OMD3", 4)) {
                len = *((uint32_t*)(ptr + 4));
                if(si_s < 8 && len > 12 + si_s && !memcmp(ptr, "ASET", 4) && !memcmp(ptr + 8 + si_s, "\x89PNG", 4)) {
                    s = 0; memcpy(&s, ptr + 8, si_s);
                    apck_pack_add(&ctx, APCK_IMAGE, (char*)strs + s, ptr + 8 + si_s, len - 8 - si_s);
                } else {
                    memcpy(dst, ptr, len);
                    dst += len;
                }
                ptr += len;
                size -= len;
            }
            memcpy(dst, "OMD3", 4);
            len = (uintptr_t)dst - (uintptr_t)model + 4;
            *((uint32_t*)(model + 4)) = len;
        }
        if(data) free(data);
        free(buf);
        if(!model) {
            fprintf(stderr, "unable to allocate memory for compressed model '%s' (%u bytes)\n", path, len);
            return 1;
        }
        if((ext = strrchr(path + pathlen, '.'))) *ext = 0;
        if(apck_pack_asset(&ctx, category, path + pathlen, model, len, 0, 0))
            free(model);
    }
    return 1;
}

/**
 * Callback for all the other asset kinds
 */
int assetcb(char *path)
{
    uint8_t *buf, *data, *ptr;
    uint64_t size;
    uint32_t len;
    char *ext;
    int r, c;

    if((buf = getfile(path, &size))) {
        if(category == APCK_VIDEO && memcmp(buf, "OggS", 4)) {
            if(!quiet) fprintf(stderr, "not an OGV video '%s'\n", path);
            free(buf);
            return 1;
        }
        /* if it's gzip compressed */
        if(buf[0] == 0x1f && buf[1] == 0x8b) {
            memcpy(&len, buf + size - 4, 4);
            ptr = buf + 3;
            c = *ptr++; ptr += 6;
            if(c & 4) { r = *ptr++; r += (*ptr++ << 8); ptr += r; }
            if(c & 8) { while(*ptr++ != 0); }
            if(c & 16) { while(*ptr++ != 0); }
            r = size - (ptr - buf);
            if((data = (unsigned char*)stbi_zlib_decode_malloc_guesssize_headerflag((const char*)ptr, r, 4096, &r, 0))) {
                free(buf);
                buf = data;
                size = r;
            } else {
                free(buf);
                if(!quiet) fprintf(stderr, "unable to uncompress '%s'\n", path + pathlen);
                return 1;
            }
        }
        if((ext = strrchr(path + pathlen, '.'))) *ext = 0;
        if(apck_pack_asset(&ctx, category, path + pathlen, buf, size, 0, 0))
            free(buf);
    }
    return 1;
}

/**
 * The main packer routine
 */
int pack(char *input, void *f, char *engine, uint32_t version, int encrypt, int maxsize)
{
    char tmp[PATH_MAX + FILENAME_MAX + 1], *fn;
    uint8_t enckey[32];
    int i, j;

    /* initialize the context */
    apck_pack_init(&ctx, f, engine, version);
    strcpy(tmp, input);
    pathlen = strlen(tmp);
    fn = tmp + pathlen++;

    /* if requested, set up encryption */
    if(encrypt) {
        srand(time(NULL));
        for(i = 0; i < 32; i += 4)
            *((uint32_t*)(enckey + i)) = rand();
        apck_pack_encrypt(&ctx, 16, enckey, aes_init, aes_write);
    }

    /* add files to the archive */
    strcpy(fn, "/icon.png");
    category = APCK_IMAGE;
    imagecb(tmp);

    strcpy(fn, "/locales");
    category = APCK_LOCALE;
    find(tmp, 0, localecb);

    strcpy(fn, "/fonts");
    category = APCK_FONT;
    find(tmp, 1, fontcb);

    strcpy(fn, "/bgm");
    category = APCK_BGM;
    find(tmp, 1, audiocb);

    strcpy(fn, "/sfx");
    category = APCK_SFX;
    find(tmp, 1, audiocb);

    strcpy(fn, "/videos");
    category = APCK_VIDEO;
    find(tmp, 1, assetcb);

    strcpy(fn, "/models");
    category = APCK_MODEL;
    find(tmp, 1, modelcb);

    strcpy(fn, "/images");
    category = APCK_IMAGE;
    find(tmp, 1, imagecb);

    strcpy(fn, "/sprites");
    category = APCK_ATLAS;
    find(tmp, 1, spritecb);

    strcpy(fn, "/maps");
    category = APCK_MAP;
    find(tmp, 1, mapcb);

    strcpy(fn, "/objects");
    category = APCK_OBJECT;
    find(tmp, 1, assetcb);

    strcpy(fn, "/scripts");
    category = APCK_SCRIPT;
    find(tmp, 1, assetcb);

    for(i = 0; i < 16; i++) {
        sprintf(fn, "/data%02u", i + 1);
        category = APCK_DATA01 + i;
        find(tmp, 1, assetcb);
    }

    if(image_build_atlas(maxsize, 1) == APCK_OK) {
        if(atlas.imgs) {
            for(i = 0; i < atlas.numimg; i++) {
                apck_pack_asset(&ctx, APCK_ATLAS, NULL, atlas.imgs[i], atlas.size * atlas.size * 4, atlas.size, atlas.size);
                free(atlas.imgs[i]);
            }
            free(atlas.imgs);
        }
        if(atlas.spr) free(atlas.spr);
        if(atlas.slots) free(atlas.slots);

        if(images) {
            for(i = 0; i < numimages; i++) {
                apck_pack_asset(&ctx, APCK_SPRITE, images[i].name, &images[i].sprite, sizeof(apck_sprite_t), 0, 0);
                for(j = 0; j < APCK_NUMDIR; j++) {
                    if(images[i].buf[j].pix) free(images[i].buf[j].pix);
                    if(images[i].sprite.frames[j]) free(images[i].sprite.frames[j]);
                }
            }
            free(images);
        }
    }
    /* actually create the archive */
    return apck_pack_flush(&ctx);
}
