#include "dmlib.h"
#include "dmzlib.h"
#include "dmfile.h"
#include "dmargs.h"
#include <zlib.h>


#define SET_TMPBUF_SIZE (1024 * 1024)

char 	*optInFilename = NULL,
        *optOutFilename = NULL;

unsigned int optSkip = 0;
unsigned int optCompressLevel = Z_BEST_COMPRESSION;
BOOL    optCompress = FALSE,
        optUseZLIB = FALSE;


static const DMOptArg optList[] =
{
    {  0, '?', "help",        "Show this help", OPT_NONE },
    {  1, 'q', "quiet",       "Decrease verbosity", OPT_NONE },
    {  2, 'Z', "zlib",        "Use ZLIB instead of dmzlib", OPT_NONE },
    {  3, 'c', "compress",    "Compress instead of decompressing (ZLIB only)", OPT_NONE },
    {  4, 'l', "level",       "Set zlib compression level 1-9", OPT_ARGREQ },
    {  5, 's', "skip",        "Skip bytes from input start", OPT_ARGREQ },
};

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


void argShowHelp()
{
    dmPrintBanner(stdout, dmProgName, "[options] [input filename] [output filename]");
    dmArgsPrintHelp(stdout, optList, optListN, 0);
}


BOOL argHandleOpt(const int optN, char *optArg, char *currArg)
{
    (void) currArg;

    switch (optN)
    {
        case 0:
            argShowHelp();
            exit(0);
            break;

        case 1:
            dmVerbosity = 0;
            break;

        case 2:
            optUseZLIB = TRUE;
            break;

        case 3:
            optUseZLIB = TRUE;
            optCompress = TRUE;
            break;

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

        case 5:
            if (!dmGetIntVal(optArg, &optSkip, NULL))
            {
                dmErrorMsg("Invalid skip value '%s'.\n", optArg);
                return FALSE;
            }
            break;

        default:
            return FALSE;
    }

    return TRUE;
}


BOOL argHandleFile(char *currArg)
{
    if (optInFilename == NULL)
        optInFilename = currArg;
    else
    if (optOutFilename == NULL)
        optOutFilename = currArg;
    else
    {
        dmErrorMsg("Excess filenames specified.\n");
        return FALSE;
    }
    return TRUE;
}


int dmTestDMZlib(FILE *inFile, FILE *outFile, size_t *inSize, size_t *outSize, BOOL compress, int level)
{
    Uint8 *inBuffer = NULL, *outBuffer = NULL;
    DMZLibContext ctx;
    int ret;

    (void) level;

    if (compress)
    {
        ret = dmError(DMERR_NOT_SUPPORTED,
            "Compression is not supported with dmzlib.\n");
        goto out;
    }

    if ((ret = dmReadDataFile(inFile, "-", &inBuffer, inSize)) != DMERR_OK)
    {
        dmErrorMsg("Failed to read file.\n");
        goto out;
    }

    if ((outBuffer = dmMalloc(SET_TMPBUF_SIZE)) == NULL)
    {
        ret = dmError(DMERR_MALLOC,
            "Malloc failed.\n");
        goto out;
    }

    // Initialize decompression structures
    if ((ret = dmZLibInitInflate(&ctx)) != DMERR_OK)
        goto out;

    ctx.inBuffer     = ctx.inBufferStart = inBuffer;
    ctx.inBufferEnd  = inBuffer + *inSize;

    ctx.outBuffer    = ctx.outBufferStart = outBuffer;
    ctx.outBufferEnd = outBuffer + SET_TMPBUF_SIZE;
    ctx.expandable   = TRUE;

    // Attempt decompression
    if ((ret = dmZLibParseHeader(&ctx, TRUE)) != DMERR_OK)
    {
        dmErrorMsg("Error parsing ZLIB header: %d, %s\n", ret, dmErrorStr(ret));
        goto out;
    }

    if ((ret = dmZLibInflate(&ctx)) != DMERR_OK)
    {
        dmErrorMsg("Error in ZLIB decompress: %d, %s\n", ret, dmErrorStr(ret));
        goto out;
    }

    outBuffer = ctx.outBufferStart;
    *outSize  = ctx.outBuffer - ctx.outBufferStart;

    if (fwrite(outBuffer, sizeof(Uint8), *outSize, outFile) != *outSize)
    {
        ret = dmError(DMERR_FWRITE, "File write error.\n");
        goto out;
    }

out:
    dmZLibCloseInflate(&ctx);
    dmFree(inBuffer);
    dmFree(outBuffer);
    return ret;
}


int dmTestZlib(FILE *inFile, FILE *outFile, size_t *inSize, size_t *outSize, BOOL compress, int level)
{
    Uint8 *inBuffer = NULL, *outBuffer = NULL;
    int zret, zinit = FALSE;
    int ret = DMERR_OK;
    z_stream zstr;

    dmMsg(0, "Operating mode: %s\n",
        compress ? "compress" : "decompress");

    dmMemset(&zstr, 0, sizeof(zstr));

    if ((inBuffer = malloc(SET_TMPBUF_SIZE)) == NULL ||
        (outBuffer = malloc(SET_TMPBUF_SIZE)) == NULL)
    {
        ret = dmError(DMERR_MALLOC, "Malloc failed.\n");
        goto out;
    }

    if (compress)
        zret = deflateInit(&zstr, level);
    else
        zret = inflateInit(&zstr);

    if (zret != Z_OK)
    {
        ret = dmError(DMERR_INIT_FAIL, "Zlib init fail.\n");
        goto out;
    }
    zinit = TRUE;

    // Initialize compression streams
    zret = Z_OK;
    *outSize = 0;
    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;

        if (compress)
            zret = deflate(&zstr, Z_FULL_FLUSH);
        else
            zret = inflate(&zstr, Z_FULL_FLUSH);

        if (zstr.total_out > 0)
        {
            *outSize += zstr.total_out;
            if (fwrite(outBuffer, sizeof(Uint8), zstr.total_out, outFile) != zstr.total_out)
            {
                ret = dmError(DMERR_FWRITE,
                    "File write error.\n");
                goto out;
            }
        }
    }

    switch (zret)
    {
        case Z_OK:
            break;

        case Z_ERRNO:
            dmErrorMsg("Error: errno\n");
            break;

        case Z_STREAM_END:
            dmErrorMsg("Stream end.\n");
            break;

        default:
            ret = dmError(DMERR_COMPRESSION,
                "Error: %d, %s\n", zret, zError(zret));
    }

out:
    *inSize = zstr.total_in;

    if (zinit)
    {
        if (compress)
            deflateEnd(&zstr);
        else
            inflateEnd(&zstr);
    }

    dmFree(inBuffer);
    dmFree(outBuffer);

    return ret;
}


int main(int argc, char *argv[])
{
    FILE *inFile = NULL, *outFile = NULL;
    size_t inSize = 0, outSize = 0;
    int ret;

    dmInitProg("testdmzlib", "ZLIB/dmzlib tester", NULL, NULL, NULL);
    dmVerbosity = 1;

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

    dmZLibInit();

    // Input and output files
    if (optInFilename == NULL)
        inFile = stdin;
    else
    if ((inFile = fopen(optInFilename, "rb")) == NULL)
    {
        int res = dmGetErrno();
        dmErrorMsg("Failed to open input file '%s': %s\n",
            optInFilename, dmErrorStr(res));
        goto out;
    }

    if (optOutFilename == NULL)
        outFile = stdout;
    else
    if ((outFile = fopen(optOutFilename, "wb")) == NULL)
    {
        int res = dmGetErrno();
        dmErrorMsg("Failed to open output file '%s': %s\n",
            optOutFilename, dmErrorStr(res));
        goto out;
    }

    if (optSkip > 0 && fseeko(inFile, optSkip, SEEK_CUR) != 0)
    {
        int res = dmGetErrno();
        dmErrorMsg("Failed to seek in input stream: %s\n",
            dmErrorStr(res));
        goto out;
    }

    if (optUseZLIB)
        ret = dmTestZlib(inFile, outFile, &inSize, &outSize, optCompress, optCompressLevel);
    else
        ret = dmTestDMZlib(inFile, outFile, &inSize, &outSize, optCompress, optCompressLevel);

    dmMsg(0, "[%d] In %d, out %d bytes.\n", ret, inSize, outSize);

out:
    // Cleanup
    if (outFile != NULL)
        fclose(outFile);

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

    dmZLibClose();
    return 0;
}