/*
 * PACKed - PACKfile EDitor
 * Programmed and designed by Matti 'ccr' Hamalainen
 * (C) Copyright 2011, 2018 Tecnic Software productions (TNSP)
 */
#include "dmtool.h"
#include "dmlib.h"
#include "dmargs.h"
#include "dmpack.h"
#include "dmfile.h"
#include "dmres.h"
#include <zlib.h>


#define	SET_MAX_FILES	  (4096)
#define	SET_TMPBUF_SIZE   (1024 * 1024)

enum
{
    CMD_NONE = 0,
    CMD_CREATE,
    CMD_ADD,
    CMD_LIST,
    CMD_EXTRACT
} DCOMMAND;

enum
{
    PACK_EXTRACTED = 0x0001,
};

int    nsrcFilenames = 0, nexcFilenames = 0;
char * srcFilenames[SET_MAX_FILES];
char * excFilenames[SET_MAX_FILES];

char * optPackFilename = NULL;
BOOL   optDoCompress = TRUE;
int    optCompressLevel = Z_BEST_COMPRESSION;
int    optCommand = CMD_NONE;
int    optDefResFlags = 0;


static const DMOptArg cmdList[] =
{
    { CMD_CREATE  , 'c', "create",      "Create and add files to PACK", OPT_NONE },
    { CMD_ADD     , 'a', "add",         "Add files to PACK", OPT_NONE },
    { CMD_LIST    , 'l', "list",        "List files in PACK", OPT_NONE },
    { CMD_EXTRACT , 'e', "extract",     "Extract files from PACK", OPT_NONE },
};

static const int cmdListN = sizeof(cmdList) / sizeof(cmdList[0]);


static const DMOptArg optList[] =
{
    {  0, '?', "help",        "Show this help", OPT_NONE },
    {  1, 'n', "nocompress",  "No compression / decompression", OPT_NONE },
    {  2, 'C', "level",       "Set zlib compression level 1-9", OPT_ARGREQ },
    {  3, 'v', "verbose",     "Increase verbosity", OPT_NONE },
    {  4, 'f', "resflags",    "Set default resource flags (-f 0xff)", OPT_ARGREQ },
    {  5, 'x', "exclude",     "Exclude name/glob pattern", OPT_ARGREQ },
};

static const int optListN = sizeof(optList) / sizeof(optList[0]);


void argShowHelp()
{
    dmPrintBanner(stdout, dmProgName, "<command> [options] <packfile> [<filename(s)>]");

    for (int i = 0; i < cmdListN; i++)
    {
        const DMOptArg *cmd = &cmdList[i];
        char tmpStr[128];
        snprintf(tmpStr, sizeof(tmpStr), "%c / %s",
            cmd->o_short, cmd->o_long);

        fprintf(stdout, "  %-15s   %s\n", tmpStr, cmd->desc);
    }

    fprintf(stdout, "\n");

    dmArgsPrintHelp(stdout, optList, optListN, 0);

    fprintf(stdout,
        "\n"
        "Examples:\n"
        "$ %s l test.pak             -- list files in test.pak\n"
        "$ %s c test.pak foobar.jpg  -- create test.pak and add foobar.jpg in it\n"
        "$ %s e test.pak foobar.jpg  -- extract foobar.jpg from test.pak\n",
        dmProgName, dmProgName, dmProgName);
}


BOOL argHandleOpt(const int optN, char *optArg, char *currArg)
{
    (void) currArg;
    switch (optN)
    {
        case 0:
            argShowHelp();
            exit(0);
            break;

        case 1:
            optDoCompress = FALSE;
            break;

        case 2:
            optCompressLevel = atoi(optArg);
            if (optCompressLevel < 1 || optCompressLevel > 9)
            {
                dmErrorMsg("Invalid compression level argument '%s', must be 1 .. 9.\n", optArg);
                return FALSE;
            }
            break;

        case 3:
            dmVerbosity++;
            break;

        case 4:
            {
                unsigned int tmpUI;
                if (!dmGetIntVal(optArg, &tmpUI, NULL))
                {
                    dmErrorMsg("Invalid flags value '%s'.\n", optArg);
                    return FALSE;
                }
                optDefResFlags = tmpUI;
            }
            break;

        case 5:
            if (nexcFilenames < SET_MAX_FILES)
            {
                excFilenames[nexcFilenames++] = optArg;
            }
            else
            {
                dmErrorMsg("Maximum number of exclusion patterns (%d) exceeded!\n",
                      SET_MAX_FILES);
                return FALSE;
            }
            break;

        default:
            dmErrorMsg("Unimplemented option argument '%s'.\n", currArg);
            return FALSE;
    }

    return TRUE;
}


BOOL argHandleNonOpt(char *currArg)
{
    if (optCommand == CMD_NONE)
    {
        for (int n = 0; n < cmdListN; n++)
        {
            const DMOptArg *cmd = &cmdList[n];
            if ((currArg[0] == cmd->o_short && currArg[1] == 0) ||
                strcmp(currArg, cmd->o_long) == 0)
            {
                optCommand = cmd->id;
                break;
            }
        }
        if (optCommand == CMD_NONE)
        {
            dmErrorMsg("Invalid command '%s'.\n", currArg);
            return FALSE;
        }
    }
    else
    if (optPackFilename == NULL)
    {
        optPackFilename = currArg;
    }
    else
    if (nsrcFilenames < SET_MAX_FILES)
    {
        srcFilenames[nsrcFilenames++] = currArg;
    }
    else
    {
        dmErrorMsg("Maximum number of input files (%d) exceeded!\n",
              SET_MAX_FILES);
        return FALSE;
    }
    return TRUE;
}


DMPackEntry *dmPackEntryCopy(const DMPackEntry *src)
{
    DMPackEntry *node = dmPackEntryNew();
    if (node == NULL)
        return NULL;

    strncpy(node->filename, src->filename, DMRES_NAME_LEN);
    node->filename[DMRES_NAME_LEN] = 0;

    node->size     = src->size;
    node->offset   = src->offset;
    node->length   = src->length;
    node->flags    = src->flags;

    return node;
}


int dmPackWrite(DMPackFile * pack)
{
    DMPackEntry *node;
    DMPackFileHeader hdr;

    if (pack == NULL)
        return DMERR_OK;

    if (pack->file == NULL)
        return DMERR_FOPEN;

    // Compute directory offset and number of entries
    memcpy(&hdr.ident, DPACK_IDENT, sizeof(hdr.ident));
    hdr.version = DPACK_VERSION;
    hdr.dirEntries = 0;
    hdr.dirOffset =
        sizeof(hdr.ident) + sizeof(hdr.version) +
        sizeof(hdr.dirEntries) + sizeof(hdr.dirOffset);

    for (node = pack->entries; node != NULL; node = node->next)
    {
        hdr.dirEntries++;
        hdr.dirOffset += node->length;
    }

    dmMsg(1, "%d entries in PACK, dir at offset 0x%08x.\n",
        hdr.dirEntries, hdr.dirOffset);

    // Write PACK header
    if (fseeko(pack->file, 0, SEEK_SET) != 0)
        return DMERR_FSEEK;

    if (!dm_fwrite_str(pack->file, (Uint8 *) & hdr.ident, sizeof(hdr.ident)) ||
        !dm_fwrite_le16(pack->file, hdr.version) ||
        !dm_fwrite_le32(pack->file, hdr.dirEntries) ||
        !dm_fwrite_le64(pack->file, hdr.dirOffset))
        return DMERR_FWRITE;

    // Write the directory
    if (fseeko(pack->file, hdr.dirOffset, SEEK_SET) != 0)
        return DMERR_FSEEK;

    for (node = pack->entries; node != NULL; node = node->next)
    {
        // Write one entry
        if (!dm_fwrite_str(pack->file, node->filename, DMRES_NAME_LEN) ||
            !dm_fwrite_le64(pack->file, node->offset) ||
            !dm_fwrite_le32(pack->file, node->length) ||
            !dm_fwrite_le32(pack->file, node->size) ||
            !dm_fwrite_le32(pack->file, node->flags))
            return DMERR_FWRITE;
    }

    return DMERR_OK;
}


int dmPackCreate(const char *filename, DMPackFile ** pack)
{
    // Allocate packfile-structure
    *pack = (DMPackFile *) dmMalloc0(sizeof(DMPackFile));
    if (*pack == NULL)
        return DMERR_MALLOC;

    // Open the file
    (*pack)->file = fopen(filename, "wb");
    if ((*pack)->file == NULL)
    {
        dmFree(*pack);
        return DMERR_FOPEN;
    }

    (*pack)->filename = dm_strdup(filename);

    // Set the result
    return DMERR_OK;
}


int dmPackAddFile(DMPackFile * pack, const char *filename,
    const Uint32 flags, const BOOL compress, int level,
    DMPackEntry ** ppEntry)
{
    Uint64 startOffs, outSize;
    Uint8 *inBuffer = NULL, *outBuffer = NULL;
    DMPackEntry entry, *node;
    FILE *inFile = NULL;
    int ret = DMERR_OK;
    BOOL zinit = FALSE;
    z_stream zstr;

    *ppEntry = NULL;

    if (pack == NULL)
        return DMERR_NULLPTR;

    if (pack->file == NULL)
        return DMERR_FOPEN;

    // Compute starting offset
    outSize = 0;
    startOffs = sizeof(DMPackFileHeader);
    for (node = pack->entries; node != NULL; node = node->next)
    {
        startOffs += node->length;
    }

    // Seek to the position
    if (fseeko(pack->file, startOffs, SEEK_SET) != 0)
        return DMERR_INVALID;

    // Read file data
    if ((inFile = fopen(filename, "rb")) == NULL)
        return DMERR_FOPEN;

    // Allocate temporary buffer
    if ((inBuffer = (Uint8 *) dmMalloc(SET_TMPBUF_SIZE)) == NULL ||
        (outBuffer = (Uint8 *) dmMalloc(SET_TMPBUF_SIZE)) == NULL)
    {
        ret = DMERR_MALLOC;
        goto out;
    }

    // Read (and possibly compress) the data
    if (compress)
    {
        int zret;
        dmMemset(&zstr, 0, sizeof(zstr));
        if (deflateInit(&zstr, level) != Z_OK)
        {
            ret = DMERR_COMPRESSION;
            goto out;
        }
        zinit = TRUE;

        // Initialize compression streams
        zret = Z_OK;
        while (!feof(inFile) && zret == Z_OK)
        {
            zstr.avail_in = fread(inBuffer, sizeof(Uint8), SET_TMPBUF_SIZE, inFile);

            zstr.next_in = inBuffer;
            zstr.next_out = outBuffer;
            zstr.avail_out = SET_TMPBUF_SIZE;
            zstr.total_out = 0;
            zret = deflate(&zstr, Z_FULL_FLUSH);

            if (zret == Z_OK && zstr.total_out > 0)
            {
                outSize += zstr.total_out;
                if (fwrite(outBuffer, sizeof(Uint8), zstr.total_out, pack->file) != zstr.total_out)
                {
                    ret = DMERR_FWRITE;
                    goto out;
                }
            }
        }
    }
    else
    {
        while (!feof(inFile))
        {
            size_t read = fread(inBuffer, sizeof(Uint8), SET_TMPBUF_SIZE, inFile);
            if (read > 0)
            {
                outSize += read;
                if (fwrite(inBuffer, sizeof(Uint8), read, pack->file) != read)
                {
                    ret = DMERR_FWRITE;
                    goto out;
                }
            }
        }
        zstr.total_in = outSize;
    }

    // Create directory entry
    strncpy(entry.filename, filename, DMRES_NAME_LEN);
    entry.filename[DMRES_NAME_LEN] = 0;
    entry.offset = startOffs;
    entry.size   = zstr.total_in;
    entry.length = outSize;
    entry.flags  = flags | (compress ? DMF_COMPRESSED : 0);

    // Add directory entry
    if ((*ppEntry = dmPackEntryCopy(&entry)) == NULL)
    {
        ret = DMERR_MALLOC;
        goto out;
    }

    dmPackEntryInsert(&pack->entries, *ppEntry);

out:
    // Cleanup
    if (zinit)
        deflateEnd(&zstr);

    dmFree(inBuffer);
    dmFree(outBuffer);
    if (inFile != NULL)
        fclose(inFile);

    return ret;
}


/*
 * EXTRACT a file from the PACK
 */
int dmPackExtractFile(DMPackFile *pack, DMPackEntry * entry, BOOL decompress)
{
    z_stream zstr;
    FILE *outFile = NULL;
    Uint8 *inBuffer = NULL, *outBuffer = NULL;
    size_t remaining;
    int zret, ret = DMERR_OK;
    BOOL zinit = FALSE;

    if (pack == NULL)
        return DMERR_NULLPTR;

    if (pack->file == NULL)
        return DMERR_FOPEN;

    // Seek to the position
    if (fseeko(pack->file, entry->offset, SEEK_SET) != 0)
        return DMERR_FSEEK;

    // Open destination file
    if ((outFile = fopen(entry->filename, "wb")) == NULL)
        return DMERR_FOPEN;

    // Allocate temporary buffer
    if ((inBuffer = (Uint8 *) dmMalloc(SET_TMPBUF_SIZE)) == NULL ||
        (outBuffer = (Uint8 *) dmMalloc(SET_TMPBUF_SIZE)) == NULL)
    {
        ret = DMERR_MALLOC;
        goto out;
    }

    // Read and uncompress the data, if needed
    if (decompress || (entry->flags & DMF_COMPRESSED) == 0)
    {
        dmMemset(&zstr, 0, sizeof(zstr));
        if (inflateInit(&zstr) != Z_OK)
        {
            ret = DMERR_COMPRESSION;
            goto out;
        }
        zinit = TRUE;
    }

    // Initialize compression streams
    remaining = entry->length;
    zret = Z_OK;
    while (remaining > 0 && zret == Z_OK)
    {
        size_t needed = remaining > SET_TMPBUF_SIZE ? SET_TMPBUF_SIZE : remaining;
        zstr.avail_in = fread(inBuffer, sizeof(Uint8), needed, pack->file);
        if (zstr.avail_in < needed)
        {
            ret = DMERR_FREAD;
            goto out;
        }

        remaining   -= zstr.avail_in;
        zstr.next_in = inBuffer;

        if (!decompress)
        {
            if (fwrite(inBuffer, sizeof(Uint8), zstr.avail_in, outFile) != zstr.avail_in)
            {
                ret = DMERR_FWRITE;
                goto out;
            }
        }
        else
        while (zstr.avail_in > 0 && zret == Z_OK)
        {
            zstr.next_out  = outBuffer;
            zstr.avail_out = SET_TMPBUF_SIZE;
            zstr.total_out = 0;
            zret = inflate(&zstr, Z_FULL_FLUSH);
            if (zstr.total_out > 0 &&
                fwrite(outBuffer, sizeof(Uint8), zstr.total_out, outFile) != zstr.total_out)
            {
                ret = DMERR_FWRITE;
                goto out;
            }
        }
    }

out:
    // Cleanup
    if (zinit)
        inflateEnd(&zstr);

    dmFree(inBuffer);
    dmFree(outBuffer);

    if (outFile != NULL)
        fclose(outFile);

    return ret;
}


/* Compare a string to a pattern. Case-SENSITIVE version.
 * The matching pattern can consist of any normal characters plus
 * wildcards ? and *. "?" matches any character and "*" matches
 * any number of characters.
 */
BOOL dmStrMatch(const char *str, const char *pattern)
{
    BOOL matched = TRUE, any = FALSE, end = FALSE;
    const char *tmp = NULL;

    // Check given pattern and string
    if (str == NULL || pattern == NULL)
        return FALSE;

    // Start comparision
    do {
        matched = FALSE;
        switch (*pattern)
        {
        case '?':
            // Any single character matches
            if (*str)
            {
                matched = TRUE;
                pattern++;
                str++;
            }
            break;

        case '*':
            matched = TRUE;
            pattern++;
            if (!*pattern)
                end = TRUE;
            any = TRUE;
            tmp = pattern;
            break;

        case 0:
            if (any)
            {
                if (*str)
                    str++;
                else
                    end = TRUE;
            }
            else
            {
                if (*str)
                {
                    if (tmp)
                    {
                        any = TRUE;
                        pattern = tmp;
                    }
                    else
                        matched = FALSE;
                }
                else
                    end = TRUE;
            }
            break;
        default:
            if (any)
            {
                if (*pattern == *str)
                {
                    any = FALSE;
                    matched = TRUE;
                }
                else
                {
                    if (*str)
                    {
                        matched = TRUE;
                        str++;
                    }
                }
            }
            else
            {
                if (*pattern == *str)
                {
                    matched = TRUE;
                    if (*pattern)
                        pattern++;
                    if (*str)
                        str++;
                }
                else
                {
                    if (tmp)
                    {
                        matched = TRUE;
                        any = TRUE;
                        pattern = tmp;
                    }
                }
            }

            if (!*str && !*pattern)
                end = TRUE;
            break;

        }        // switch

    } while (matched && !end);

    return matched;
}


BOOL dmCheckExcluded(const char *filename)
{
    for (int i = 0; i < nexcFilenames; i++)
    {
        if (dmStrMatch(filename, excFilenames[i]))
            return TRUE;
    }
    return FALSE;
}


int main(int argc, char *argv[])
{
    int res = 0;
    DMPackFile *pack = NULL;

    // Parse arguments
    dmInitProg("packed", "Pack File Editor", "0.7", NULL, NULL);
    dmVerbosity = 0;

    if (!dmArgsProcess(argc, argv, optList, optListN,
        argHandleOpt, argHandleNonOpt, OPTH_BAILOUT))
        exit(1);

    if (optCommand == CMD_NONE)
    {
        dmErrorMsg("Nothing to do. Perhaps try \"%s --help\"?\n",
            dmProgName);
        goto error;
    }

    if (optPackFilename == NULL)
    {
        dmErrorMsg("No PACK file specified.\n");
        goto error;
    }

    dmMsg(1, "Processing PACK '%s' ...\n", optPackFilename);

    // Execute command
    switch (optCommand)
    {
    case CMD_CREATE:
    case CMD_ADD:
        switch (optCommand)
        {
        case CMD_CREATE:
            dmMsg(1, "Creating new PACK\n");
            res = dmPackCreate(optPackFilename, &pack);
            break;

        case CMD_ADD:
            dmMsg(1, "Opening existing PACK\n");
            res = dmPackOpen(optPackFilename, &pack, FALSE);
            break;
        }

        // Add files into PACK
        if (res == DMERR_OK)
        {
            dmMsg(1, "Adding files...\n");

            for (int i = 0; i < nsrcFilenames; i++)
            if (!dmCheckExcluded(srcFilenames[i]))
            {
                DMPackEntry *node = NULL;
                int res = dmPackAddFile(pack, srcFilenames[i],
                    optDefResFlags, optDoCompress,
                    optCompressLevel, &node);

                if (res != DMERR_OK)
                {
                    dmPrint(1, "%-32s [ERROR:%d]\n",
                        srcFilenames[i], res);
                }
                else
                {
                    dmPrint(1, "%-32s ['%s', s=%d, c=%d, o=%ld, f=0x%04x]\n",
                        srcFilenames[i], node->filename,
                        node->size, node->length, node->offset, node->flags);
                }
            }

            dmMsg(2, "w=%d\n", dmPackWrite(pack));
            dmMsg(2, "c=%d\n", dmPackClose(pack));
        }
        else
        {
            dmErrorMsg("Could not open packfile, error #%d: %s\n", res,
                  dmErrorStr(res));
        }
        break;

    case CMD_LIST:
        // List files in PACK
        res = dmPackOpen(optPackFilename, &pack, TRUE);
        if (res == DMERR_OK)
        {
            DMPackEntry *node;
            int i;
            for (i = 0, node = pack->entries; node; i++)
                node = node->next;

            dmMsg(1, "%d files total\n", i);

            dmPrint(0, "%-32s | %8s | %8s | %8s | %s\n",
                "Name", "Size", "CSize", "Offset", "ResFlags");

            for (node = pack->entries; node != NULL; node = node->next)
            {
                BOOL match;

                // Check for matches in specified filenames/patterns
                if (nsrcFilenames > 0)
                {
                    match = FALSE;
                    for (int i = 0; i < nsrcFilenames && !match; i++)
                        match = dmStrMatch(node->filename, srcFilenames[i]);
                }
                else
                // No filenames specified, everything shown
                    match = TRUE;

                if (match)
                {
                    dmPrint(0, "%-32s | %8d | %8d | %08x | %04x\n",
                        node->filename, node->size, node->length,
                        node->offset, node->flags);
                }
            }

            dmMsg(2, "c=%d\n", dmPackClose(pack));
        }
        else
            dmErrorMsg("Could not open packfile, error #%d: %s\n", res,
                  dmErrorStr(res));
        break;

    case CMD_EXTRACT:
        // Extract files from PACK
        res = dmPackOpen(optPackFilename, &pack, TRUE);
        if (res == DMERR_OK)
        {
            DMPackEntry *node;
            FILE *resFile = fopen(DMRES_RES_FILE, "w");
            if (resFile == NULL)
            {
                int err = dmGetErrno();
                dmErrorMsg("Could not create resource output file '%s' #%d: %s\n",
                    DMRES_RES_FILE, err, dmErrorStr(err));
            }

            for (node = pack->entries; node != NULL; node = node->next)
            {
                BOOL match;

                // Check for matches
                if (nsrcFilenames > 0)
                {
                    match = FALSE;
                    for (int i = 0; i < nsrcFilenames && !match; i++)
                    {
                        if (dmStrMatch(node->filename, srcFilenames[i]) &&
                            !dmCheckExcluded(node->filename))
                            match = TRUE;
                    }
                }
                else
                {
                    match = !dmCheckExcluded(node->filename);
                }

                if (match && (node->privFlags & PACK_EXTRACTED) == 0)
                {
                    // Mark as done
                    node->privFlags |= PACK_EXTRACTED;

                    // Print one entry
                    dmPrint(0, "Extracting: %-32s [siz=%d, cmp=%d, offs=0x%08x, flags=0x%04x]\n",
                            node->filename, node->size, node->length,
                            node->offset, node->flags);

                    dmPackExtractFile(pack, node, optDoCompress);
                }
            }

            dmMsg(2, "c=%d\n", dmPackClose(pack));

            if (resFile != NULL)
                fclose(resFile);
        }
        else
            dmErrorMsg("Could not open packfile, error #%d: %s\n", res,
                  dmErrorStr(res));
        break;
    }

error:
    return 0;
}