/*
 * 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,
};


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

char * optPackFilename = NULL;
int    optCompressLevel = Z_BEST_COMPRESSION;
int    optCommand = CMD_NONE;


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 , 'x', "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,   0, "license"         , "Print out this program's license agreement", OPT_NONE },
    {  2, 'v', "verbose"         , "Be more verbose", OPT_NONE },

    {  4, 'C', "level"           , "Set zlib compression level 0-9", OPT_ARGREQ },
    {  6, 'x', "exclude"         , "Exclude name/glob pattern", OPT_ARGREQ },
};

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


void argShowHelp()
{
    FILE *fh = stdout;

    dmPrintBanner(fh, dmProgName, "<command> [options] <packfile> [<filename(s)>]");

    fprintf(fh, "\n");

    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(fh, "  %-15s   %s\n", tmpStr, cmd->desc);
    }

    fprintf(fh, "\n");

    dmArgsPrintHelp(fh, optList, optListN, 0, 80 - 2);

    fprintf(fh,
        "\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)
{
    switch (optN)
    {
        case 0:
            argShowHelp();
            exit(0);
            break;

        case 1:
            dmPrintLicense(stdout);
            exit(0);
            break;

        case 2:
            dmVerbosity++;
            break;

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

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


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

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

    entry->size     = src->size;
    entry->offset   = src->offset;
    entry->csize    = src->csize;
    entry->flags    = src->flags;

    return entry;
}


int dmPackWrite(DMPackFile * pack)
{
    DMPackEntry *entry;
    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 (entry = pack->entries; entry != NULL; entry = entry->next)
    {
        hdr.dirEntries++;
        hdr.dirOffset += entry->csize;
    }

    dmMsg(1, "%d entries in PACK, dir at offset 0x%08" DM_PRIx64 ".\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 (entry = pack->entries; entry != NULL; entry = entry->next)
    {
        // Write one entry
        if (!dm_fwrite_str(pack->file, entry->filename, DMRES_NAME_LEN) ||
            !dm_fwrite_le64(pack->file, entry->offset) ||
            !dm_fwrite_le32(pack->file, entry->csize) ||
            !dm_fwrite_le32(pack->file, entry->size) ||
            !dm_fwrite_le32(pack->file, entry->flags))
            return DMERR_FWRITE;
    }

    return DMERR_OK;
}


int dmPackCreate(const char *filename, DMPackFile ** pack)
{
    // Allocate packfile-structure
    if ((*pack = (DMPackFile *) dmMalloc0(sizeof(DMPackFile))) == 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 int level, DMPackEntry ** ppEntry)
{
    Uint64 startOffs, outSize;
    Uint8 *inBuffer = NULL, *outBuffer = NULL;
    DMPackEntry rentry, *pentry;
    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 (pentry = pack->entries; pentry != NULL; pentry = pentry->next)
    {
        startOffs += pentry->csize;
    }

    // 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 (level > 0)
    {
        int zret;

        // Initialize compression streams
        memset(&zstr, 0, sizeof(zstr));
        if (deflateInit(&zstr, level) != Z_OK)
        {
            ret = DMERR_COMPRESSION;
            goto out;
        }
        zinit = TRUE;

        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(rentry.filename, filename, DMRES_NAME_LEN);
    rentry.filename[DMRES_NAME_LEN] = 0;
    rentry.offset = startOffs;
    rentry.size   = zstr.total_in;
    rentry.csize  = outSize;
    rentry.flags  = level > 0 ? DMF_COMPRESSED : 0;

    // Add directory entry
    if ((*ppEntry = dmPackEntryCopy(&rentry)) == 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;
}


int dmPackExtractFile(DMPackFile *pack, DMPackEntry * entry, const 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))
    {
        // Initialize compression streams
        memset(&zstr, 0, sizeof(zstr));
        if (inflateInit(&zstr) != Z_OK)
        {
            ret = DMERR_COMPRESSION;
            goto out;
        }
        zinit = TRUE;
    }

    remaining = entry->csize;
    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 (zinit)
        {
            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;
                }
            }
        }
        else
        {
            if (fwrite(inBuffer, sizeof(Uint8), zstr.avail_in, outFile) != zstr.avail_in)
            {
                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 dmMatchList(char **list, const size_t nlist, const char *str)
{
    for (size_t i = 0; i < nlist; i++)
    {
        if (dmStrMatch(str, list[i]))
            return TRUE;
    }
    return FALSE;
}


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

    // Parse arguments
    dmInitProg("packed", "TNSP 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 || argc < 2)
    {
        argShowHelp();
        goto out;
    }

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

    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;
        }

        if (res != DMERR_OK)
        {
            dmErrorMsg("Could not open PACK file: %s\n",
                dmErrorStr(res));
            goto out;
        }

        // Add files into PACK
        dmMsg(1, "Adding files...\n");

        for (size_t i = 0; i < nsrcFilenames; i++)
        if (!dmMatchList(excFilenames, nexcFilenames, srcFilenames[i]))
        {
            res = dmPackAddFile(pack, srcFilenames[i], optCompressLevel, &entry);

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

        dmMsg(2, "w=%d\n", dmPackWrite(pack));
        break;

    case CMD_LIST:
        // List files in PACK
        if ((res = dmPackOpen(optPackFilename, &pack, TRUE)) != DMERR_OK)
        {
            dmErrorMsg("Could not open PACK file: %s\n",
                dmErrorStr(res));
            goto out;
        }
        else
        {
            size_t ntotal = 0, nshown = 0,
                totalSize = 0, totalCSize = 0;

            // Print the list header
            int hlen = printf(
                "%-32s | %10s | %10s | %10s | %s\n",
                "Name", "Size", "CSize", "Offset", "Flags");

            for (int n = 0; n < hlen; n++)
                fputc('-', stdout);

            fputs("\n", stdout);

            // Print all the matching entries
            for (entry = pack->entries; entry != NULL; entry = entry->next)
            {
                BOOL match;

                if (nsrcFilenames > 0)
                    // Check for matches in specified filenames/patterns
                    match = dmMatchList(srcFilenames, nsrcFilenames, entry->filename);
                else
                    // No filenames specified, everything shown
                    match = TRUE;

                if (match && !dmMatchList(excFilenames, nexcFilenames, entry->filename))
                {
                    printf("%-32s | %10d | %10d | %010" DM_PRIx64 " | %04x\n",
                        entry->filename, entry->size, entry->csize,
                        entry->offset, entry->flags);

                    nshown++;
                    totalSize += entry->size;
                    totalCSize += entry->csize;
                }

                ntotal++;
            }

            for (int n = 0; n < hlen; n++)
                fputc('-', stdout);

            printf(
                "\n"
                "%-32s | %10" DM_PRIu_SIZE_T " | %10" DM_PRIu_SIZE_T " |\n"
                "Listed %" DM_PRIu_SIZE_T " of %" DM_PRIu_SIZE_T " files total.\n",
                "Totals",
                totalSize, totalCSize,
                nshown, ntotal);
        }
        break;

    case CMD_EXTRACT:
        // Extract files from PACK
        if ((res = dmPackOpen(optPackFilename, &pack, TRUE)) != DMERR_OK)
        {
            dmErrorMsg("Could not open PACK file: %s\n",
                dmErrorStr(res));
            goto out;
        }

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

            // Check for matches
            if (nsrcFilenames > 0)
            {
                match = !dmMatchList(excFilenames, nexcFilenames, entry->filename) &&
                    dmMatchList(srcFilenames, nsrcFilenames, entry->filename);
            }
            else
            {
                match = !dmMatchList(excFilenames, nexcFilenames, entry->filename);
            }

            if (match && (entry->privFlags & PACK_EXTRACTED) == 0)
            {
                // Print one entry
                dmPrint(0, "Extracting: %-32s [siz=%d, cmp=%d, offs=0x%08" DM_PRIx64 ", flags=0x%04x]\n",
                        entry->filename, entry->size, entry->csize,
                        entry->offset, entry->flags);

                if ((res = dmPackExtractFile(pack, entry, optCompressLevel > 0)) == DMERR_OK)
                    entry->privFlags |= PACK_EXTRACTED;
            }
        }
        break;
    }

out:
    dmPackClose(pack);

    return 0;
}