#include <stdint.h>
#include <unistd.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <zlib.h>
#include <getopt.h>
#include <stdbool.h>

#include "espfsformat.h"
#include "heatshrink_encoder.h"
#include "parsing.h"

#define DEFAULT_GZIP_EXTS "html,css,js,svg,png,jpg,gif"

struct InputFileLinkedListEntry;

struct InputFileLinkedListEntry {
  char *name;
  struct InputFileLinkedListEntry *next;
};

/// two ends of a linked list with input files
static struct InputFileLinkedListEntry *s_inputFiles = NULL;
static struct InputFileLinkedListEntry *s_lastInputFile = NULL;

/// Output file FD
static int s_outFd = 1;

/// Array of gzipped extensions, ends with a NULL pointer
static char **s_gzipExtensions = NULL;
/// Gzip all files
static bool s_gzipAll = false;

// impls to satisfy defs in the config header
void *httpdPlatMalloc(size_t len)
{
    return malloc(len);
}

void httpdPlatFree(void *ptr)
{
    free(ptr);
}

/**
 * Compress a file using Heatshrink
 *
 * @param[in] in - pointer to the uncompressed input
 * @param insize - len of the uncompressed input
 * @param[out] out - destination buffer for the compressed data
 * @param outcap - capacity of the output buffer
 * @param level - compression level, 1-9; -1 for default.
 * @return actual length of the compressed data
 */
size_t compressHeatshrink(const uint8_t *in, size_t insize, uint8_t *out, size_t outcap, int level)
{
    const uint8_t *inp = in;
    uint8_t *outp = out;
    size_t len;
    uint8_t ws[] = {5, 6, 8, 11, 13};
    uint8_t ls[] = {3, 3, 4, 4, 4};
    HSE_poll_res pres = 0;
    HSE_sink_res sres = 0;
    size_t r;
    if (level == -1) { level = 8; }
    level = (level - 1) / 2; //level is now 0, 1, 2, 3, 4
    heatshrink_encoder *enc = heatshrink_encoder_alloc(ws[level], ls[level]);
    if (enc == NULL) {
        perror("allocating mem for heatshrink");
        exit(1);
    }
    //Save encoder parms as first byte
    *outp = (uint8_t) ((ws[level] << 4) | ls[level]);
    outp++;
    outcap--;

    r = 1;
    do {
        if (insize > 0) {
            sres = heatshrink_encoder_sink(enc, inp, insize, &len);
            if (sres != HSER_SINK_OK) { break; }
            inp += len;
            insize -= len;
            if (insize == 0) { heatshrink_encoder_finish(enc); }
        }
        do {
            pres = heatshrink_encoder_poll(enc, outp, outcap, &len);
            if (pres != HSER_POLL_MORE && pres != HSER_POLL_EMPTY) { break; }
            outp += len;
            outcap -= len;
            r += len;
        } while (pres == HSER_POLL_MORE);
    } while (insize != 0);

    if (insize != 0) {
        fprintf(stderr, "Heatshrink: Bug? insize is still %d. sres=%d pres=%d\n", (int) insize, sres, pres);
        exit(1);
    }

    heatshrink_encoder_free(enc);
    return r;
}

/**
 * Compress a file using Gzip
 *
 * @param[in] in - pointer to the uncompressed input
 * @param insize - len of the uncompressed input
 * @param[out] out - destination buffer for the compressed data
 * @param outcap - capacity of the output buffer
 * @param level - compression level, 1-9; -1 for default.
 * @return actual length of the compressed data
 */
size_t compressGzip(const uint8_t *in, size_t insize, uint8_t *out, size_t outcap, int level)
{
    z_stream stream;
    int zresult;

    stream.zalloc = Z_NULL;
    stream.zfree = Z_NULL;
    stream.opaque = Z_NULL;
    stream.next_in = in;
    stream.avail_in = (uInt) insize;
    stream.next_out = out;
    stream.avail_out = (uInt) outcap;
    // 31 -> 15 window bits + 16 for gzip
    zresult = deflateInit2(&stream, level, Z_DEFLATED, 31, 8, Z_DEFAULT_STRATEGY);
    if (zresult != Z_OK) {
        fprintf(stderr, "DeflateInit2 failed with code %d\n", zresult);
        exit(1);
    }

    zresult = deflate(&stream, Z_FINISH);
    if (zresult != Z_STREAM_END) {
        fprintf(stderr, "Deflate failed with code %d\n", zresult);
        exit(1);
    }

    zresult = deflateEnd(&stream);
    if (zresult != Z_OK) {
        fprintf(stderr, "DeflateEnd failed with code %d\n", zresult);
        exit(1);
    }

    return stream.total_out;
}

/**
 * Check if a file name should be compressed by gzip
 *
 * @param name - file name
 * @return true if should compress
 */
bool shouldCompressGzip(const char *name)
{
    if (!s_gzipExtensions) { return false; }
    if (s_gzipAll) { return true; }

    const char *ext = name + strlen(name);
    while (*ext != '.') {
        ext--;
        if (ext < name) {
            // no dot in file name -> no extension -> nothing to match against
            return false;
        }
    }
    ext++;

    int i = 0;
    while (s_gzipExtensions[i] != NULL) {
        if (strcasecmp(ext, s_gzipExtensions[i]) == 0) {
            return true;
        }
        i++;
    }

    return false;
}

/**
 * Parse a list of gzipped extensions
 *
 * @param input - list of comma-separated extensions, e.g. "jpg,png"
 * @return
 */
void parseGzipExtensions(char *input)
{
    char *token;
    char *extList = input;
    size_t count = 2; // one for first element, second for terminator

    // count elements
    while (*extList != 0) {
        if (*extList == ',') { count++; }
        extList++;
    }

    // split string
    extList = input;
    s_gzipExtensions = malloc(count * sizeof(char *));
    count = 0;
    token = strtok(extList, ",");
    while (token) {
        s_gzipExtensions[count++] = token;
        token = strtok(NULL, ",");
    }
    // terminate list
    s_gzipExtensions[count] = NULL;
}

/**
 * Process a file.
 *
 * @param fd - filedes
 * @param[in] name - filename to embed in the archive
 * @param compression_mode - compression mode
 * @param level - compression level for heatshrink, 1-9
 * @param[out] compName - the used compression is output here (for debug print)
 * @return - size of the output, in percent (100% = no compression)
 */
int handleFile(int fd, const char *name, int compression_mode, int level, const char **compName)
{
    uint8_t *fdat = NULL, *cdat = NULL, *cdatbuf = NULL;
    size_t size, csize;
    EspFsHeader h;
    uint16_t realNameLen;
    uint8_t flags = 0;
    size = (size_t) lseek(fd, 0, SEEK_END);
    fdat = malloc(size);
    lseek(fd, 0, SEEK_SET);
    read(fd, fdat, size);

    if (shouldCompressGzip(name)) {
        csize = size * 3;
        if (csize < 100) { // gzip has some headers that do not fit when trying to compress small files
            csize = 100;
        } // enlarge buffer if this is the case
        cdat = cdatbuf = malloc(csize);
        csize = compressGzip(fdat, size, cdat, csize, level);
        compression_mode = COMPRESS_NONE; // don't use heatshrink if gzip was already used - it would only make it bigger
        flags = FLAG_GZIP;
    } else if (compression_mode == COMPRESS_NONE) {
        csize = size;
        cdat = fdat;
    } else if (compression_mode == COMPRESS_HEATSHRINK) {
        cdat = cdatbuf = malloc(size * 2);
        csize = compressHeatshrink(fdat, size, cdat, size * 2, level);
    } else {
        fprintf(stderr, "Unknown compression - %d\n", compression_mode);
        exit(1);
    }

    if (csize > size) {
        fprintf(stderr, "! Compression enbiggened %s, embed as plain\n", name);
        //Compressing enbiggened this file. Revert to uncompressed store.
        compression_mode = COMPRESS_NONE;
        csize = size;
        cdat = fdat;
        flags = 0;
    }

    //Fill header data
    h.magic = htole32(ESPFS_MAGIC); // ('E' << 0) + ('S' << 8) + ('f' << 16) + ('s' << 24);
    h.flags = flags;
    h.compression = (uint8_t) compression_mode;
    h.nameLen = realNameLen = (uint16_t) strlen(name) + 1; // zero terminator
    uint16_t padbytes = 0;
    if (h.nameLen & 3) {
        //Round to next 32bit boundary
        padbytes = 4 - (h.nameLen & 3);
        h.nameLen += padbytes; // include the bytes in "name" to make parsing easier - these will be zeroed out, so the c-string remains the same.
    }
    h.nameLen = htole16(h.nameLen);
    h.fileLenComp = htole32((uint32_t) csize);
    h.fileLenDecomp = htole32((uint32_t) size);

    write(s_outFd, &h, sizeof(EspFsHeader));
    write(s_outFd, name, realNameLen);
    if (padbytes) {
        write(s_outFd, "\0\0\0", padbytes); // these zeros are included in h.nameLen
    }
    write(s_outFd, cdat, csize);
    //Pad out to 32bit boundary - the parser does this automatically when walking over the archive.
    if (csize & 3) {
        padbytes = 4 - (csize & 3);
        write(s_outFd, "\0\0\0", padbytes);
        csize += padbytes;
    }

    free(fdat);
    if (cdatbuf) {
        // free the buffer allocated for compression output
        free(cdatbuf);
    }

    // debug outputs ...

    if (compName != NULL) {
        if (h.compression == COMPRESS_HEATSHRINK) {
            *compName = "heatshrink";
        } else if (h.compression == COMPRESS_NONE) {
            if (h.flags & FLAG_GZIP) {
                *compName = "gzip";
            } else {
                *compName = "none";
            }
        } else {
            *compName = "unknown";
        }
    }

    // get compression % (lower is better)
    return size ? (int) ((csize * 100) / size) : 100;
}

/**
 * Write final dummy header with FLAG_LASTFILE set.
 */
void finishArchive(void)
{
    EspFsHeader h;
    h.magic = htole32(ESPFS_MAGIC);
    h.flags = FLAG_LASTFILE;
    h.compression = COMPRESS_NONE;
    h.nameLen = 0;
    h.fileLenComp = 0;
    h.fileLenDecomp = 0;
    write(s_outFd, &h, sizeof(EspFsHeader));
}

/**
 * Queue a file for adding to the archive.
 * Appends it to the `s_inputFiles` linked list.
 *
 * @param name - file name to add
 */
void queueInputFile(const char *name)
{
    fprintf(stderr, "INFILE: %s\n", name);

    struct InputFileLinkedListEntry *tmp = malloc(sizeof(struct InputFileLinkedListEntry));
    tmp->name = strdup(name);
    tmp->next = NULL;

    if (s_lastInputFile == NULL) {
        s_inputFiles = tmp;
        s_lastInputFile = tmp;
    } else {
        s_lastInputFile->next = tmp;
        s_lastInputFile = tmp;
    }
}

int main(int argc, char **argv)
{
    int f;
    char inputFileName[1024];
    struct stat statBuf;
    int serr;
    int rate;
    int err = 0;
    int compType;  //default compression type - heatshrink
    int compLvl = -1;
    bool use_gzip = false;

    compType = COMPRESS_HEATSHRINK;

    int c;
    char *outfile = NULL;
    char *c_outfile = NULL;
    char *c_varname = NULL;
    char *parseFile = NULL;
    char *stripPath = NULL;
    char *extractFile = NULL;

    while (1) {
        int option_index = 0;
        static struct option long_options[] = {
                {"parse",      required_argument, 0, 'p'},
                {"extract",    required_argument, 0, 'e'},
                {"compress",   required_argument, 0, 'c'},
                {"gzip",       no_argument,       0, 'z'},
                {"gzip-all",   no_argument,       0, 'G'},
                {"level",      required_argument, 0, 'l'},
                {"gzip-exts",  required_argument, 0, 'g'},
                {"input",      required_argument, 0, 'i'},
                {"output",     required_argument, 0, 'o'},
                {"c-output",   required_argument, 0, 'C'},
                {"c-varname",  required_argument, 0, 'N'},
                {"strip-path", required_argument, 0, 'S'},
                {"help",       no_argument,       0, 'h'},
                { /* end marker */ }
        };

        c = getopt_long(argc, argv, "c:l:g:zGS:e:hp:C:i:o:0123456789",
                        long_options, &option_index);
        if (c == -1) {
            break;
        }

        switch (c) {
            case 'h':
                goto show_help;

            case 'z':
                use_gzip = true;
                break;

            case '0' ... '9':
                compLvl = c - '0';
                break;

            case 'p':
                parseFile = strdup(optarg);
                break;

            case 'e':
                extractFile = strdup(optarg);
                break;

            case 'C':
                c_outfile = strdup(optarg);
                break;

            case 'N':
                c_varname = strdup(optarg);
                break;

            case 'S':
                stripPath = strdup(optarg);
                break;

            case 'c':
                compType = atoi(optarg);
                break;

            case 'G':
                use_gzip = true;
                s_gzipAll = true;
                break;

            case 'g':
                use_gzip = true;
                parseGzipExtensions(optarg);
                break;

            case 'l':
                compLvl = atoi(optarg);
                if (compLvl < 1 || compLvl > 9) {
                    fprintf(stderr, "Bad compression level: %d\n", compLvl);
                    err = 1;
                    goto show_help;
                }
                break;

            case 'i':
                queueInputFile(optarg);
                break;

            case 'o':
                outfile = strdup(optarg);
                break;

            case '?':
                goto show_help;

            default:
                fprintf(stderr, "Unknown option: %c\n", c);
                err = 1;
                goto show_help;
        }
    }

    FILE *outfp = NULL;
    if (outfile && 0 != strcmp("-", outfile)) {
        fprintf(stderr, "Writing to %s\n", outfile);
        outfp = fopen(outfile, "w+");
        if (!outfp) {
            perror(outfile);
            return 1;
        }
        s_outFd = fileno(outfp);
        ftruncate(s_outFd, 0);
    } else {
        if (outfile) {
            free(outfile);
            outfile = NULL;
        }
        fprintf(stderr, "Writing to stdout\n\n");
    }

    if (parseFile) {
        parseEspfsImage(parseFile, extractFile, s_outFd);
        exit(0);
    }

    if (s_gzipExtensions == NULL && use_gzip) {
        parseGzipExtensions(strdup(DEFAULT_GZIP_EXTS));
    }

    if (optind < argc) {
        while (optind < argc) {
            queueInputFile(argv[optind++]);
        }
    }

    if (!s_inputFiles) {
        fprintf(stderr, "Reading input file names from stdin\n");
        while (fgets(inputFileName, sizeof(inputFileName), stdin)) {
            //Kill off '\n' at the end
            inputFileName[strlen(inputFileName) - 1] = 0;
            queueInputFile(inputFileName);
        }
    }

    struct InputFileLinkedListEntry *entry = s_inputFiles;
    while (entry) {
        char *name = entry->name;
        //Only include files
        serr = stat(name, &statBuf);
        if ((serr == 0) && S_ISREG(statBuf.st_mode)) {
            //Strip off './' or '/' madness.
            char *embeddedName = name;
            f = open(name, O_RDONLY);

            // relative path starting with ./, remove that
            if (embeddedName[0] == '.' && embeddedName[1] == '/') {
                embeddedName += 2;
            }
            // remove prefix
            if (stripPath && 0 == strncmp(embeddedName, stripPath, strlen(stripPath))) {
                embeddedName += strlen(stripPath);
            }
            // remove leading slash, if any
            if (embeddedName[0] == '/') { embeddedName++; }

            if (f > 0) {
                const char *compName = "unknown";
                rate = handleFile(f, embeddedName, compType, compLvl, &compName);
                fprintf(stderr, "%s (%d%%, %s)\n", embeddedName, rate, compName);
                close(f);
            } else {
                perror(name);
            }
        } else if (serr != 0) {
            perror(name);
        }

        entry = entry->next;
    }

    finishArchive();
    fsync(s_outFd);

    if (outfp) {
        if (outfile && c_outfile) {
            const size_t imagelen = (size_t) lseek(s_outFd, 0, SEEK_END);
            lseek(s_outFd, 0, SEEK_SET);

            FILE *coutf = fopen(c_outfile, "w+");
            if (!coutf) {
                perror(c_outfile);
                return 1;
            }
            int cfd = fileno(coutf);
            ftruncate(cfd, 0);

            if (!c_varname) {
                c_varname = strdup("espfs_image");
            }

            fprintf(coutf, "unsigned char %s[%lu] = {", c_varname, imagelen);
            for (size_t i = 0; i < imagelen; i++) {
                fprintf(coutf, (i % 16) ? " " : "\n  ");
                uint8_t u;
                read(s_outFd, &u, 1);
                fprintf(coutf, "0x%02x,", u);
            }
            fprintf(coutf, "\n};\nunsigned int %s_len = %lu;\n", c_varname, imagelen);
            fsync(cfd);
            fclose(coutf);
        }

        fclose(outfp);
    }

    return 0;


    show_help:
    //               ##########**********##########**********##########----------$$$$$$$$$$----------$$$$$$$$$$----------| 80 chars
    fprintf(stderr, "%s - Program to create espfs images\n", argv[0]);
    fprintf(stderr, "Options:\n");
    fprintf(stderr, "[-p|--parse FILE]\n"
                    "    Parse an espfs file and show a list of its contents. No other options apply in this mode.\n");
    fprintf(stderr, "[-e|--extract FILE]\n"
                    "    Extract a file with the given name from the parsed file (-p)\n");
    fprintf(stderr, "[-S|--strip-path PATH]\n"
                    "    Remove the given path from input file names before packing\n");
    fprintf(stderr, "[-c|--compress COMPRESSOR]\n"
                    "    0 - None, 1 - Heatshrink (default)\n");
    fprintf(stderr, "[-l|--level LEVEL] or [-0 through -9]\n"
                    "    compression level 1-9, higher is better but uses more RAM\n");
    fprintf(stderr, "[-z|--gzip]\n"
                    "    use gzip for files with extensions matching "DEFAULT_GZIP_EXTS"\n");
    fprintf(stderr, "[-G|--gzip-all]\n"
                    "    use gzip for all files\n");
    fprintf(stderr, "[-g|--gzip-exts GZIPPED_EXTENSIONS]\n"
                    "    use gzip for files with custom extensions, comma-separated\n");
    fprintf(stderr, "[-i|--input FILE]\n"
                    "    Input file, can be multiple. Files can also be passed at the end without -i, or as lines on stdin\n"
                    "    if not specified by args\n");
    fprintf(stderr, "[-o|--output FILE]\n"
                    "    Output file name; if not specified or equal to \"-\", outputs to stdout\n");
    fprintf(stderr, "[-C|--c-output FILE]\n"
                    "    C embed output file name (an uint8_t array for embedding); can be used only if a binary output\n"
                    "    file is specified as well (using -o), as this reads the output file and converts it as a second step.\n");
    fprintf(stderr, "[-h|--help\n"
                    "    Show help.\n\n");
    exit(err);
}