#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 <errno.h>

#include "espfsformat.h"
#include "heatshrink_encoder.h"
#include "parsing.h"
#include "httpd-logging.h"

void show_version(char **argv);
void show_help(int retval, char **argv);

#define DEFAULT_GZIP_EXTS "css,js,svg,png,jpg,jpeg,webm,ico,gif"
#define DEFAULT_C_VARNAME "espfs_image"

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 = STDOUT_FILENO;

/// 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) {
        espfs_error("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) {
        espfs_error("Heatshrink: Bug? insize is still %d. sres=%d pres=%d", (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
 * @return actual length of the compressed data
 */
size_t compressGzip(const uint8_t *in, size_t insize, uint8_t *out, size_t outcap)
{
    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, Z_BEST_COMPRESSION /* we want the smallest possible files */, Z_DEFLATED, 31, 8, Z_DEFAULT_STRATEGY);
    if (zresult != Z_OK) {
        espfs_error("DeflateInit2 failed with code %d", zresult);
        exit(1);
    }

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

    zresult = deflateEnd(&stream);
    if (zresult != Z_OK) {
        espfs_error("DeflateEnd failed with code %d", 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);
        compression_mode = EFS_COMPRESS_NONE; // don't use heatshrink if gzip was already used - it would only make it bigger
        flags = EFS_FLAG_GZIP;
    } else if (compression_mode == EFS_COMPRESS_NONE) {
        csize = size;
        cdat = fdat;
    } else if (compression_mode == EFS_COMPRESS_HEATSHRINK) {
        cdat = cdatbuf = malloc(size * 2);
        csize = compressHeatshrink(fdat, size, cdat, size * 2, level);
    } else {
        espfs_error("Unknown compression - %d", compression_mode);
        exit(1);
    }

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

    //Fill header data
    h.magic = htole32(EFS_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 == EFS_COMPRESS_HEATSHRINK) {
            *compName = "heatshrink";
        } else if (h.compression == EFS_COMPRESS_NONE) {
            if (h.flags & EFS_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(EFS_MAGIC);
    h.flags = EFS_FLAG_LASTFILE;
    h.compression = EFS_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)
{
    espfs_dbg("INFILE: %s", 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;
    }
}

enum OpMode {
  OM_INVALID = 0,
  OM_PACK = 'P',
  OM_LIST = 'L',
  OM_EXTRACT = 'X',
  OM_EMBED = 'M',
};

#define BUFLEN 1024
int main(const int argc, char **argv)
{
    int f;
    char inputFileName[BUFLEN];
    char tempbuf[BUFLEN];
    struct stat statBuf;
    int serr;
    int rate;
    int compType;  //default compression type - heatshrink
    int compLvl = -1;
    bool use_gzip = false;
    enum OpMode opmode = OM_INVALID;

    compType = EFS_COMPRESS_HEATSHRINK;

    int c;
    char *outFileName = NULL;
    char *c_varname = NULL;
    char *stripPath = NULL;
    char *extractFileName = NULL;

    size_t num_input_files = 0;
    bool read_from_stdin = false;

    while (1) {
        int option_index = 0;
        static struct option long_options[] = {
                // Help
                {"help",       no_argument,       0, 'h'},
                {"version",    no_argument,       0, 'V'},

                // Main operation (one at a time)
                {"pack",       no_argument,       0, 'P'},
                {"list",       no_argument,       0, 'L'},
                {"extract",    required_argument, 0, 'X'},
                {"embed",      no_argument,       0, 'M'},

                // Common options
                {"input",      required_argument, 0, 'i'}, // input file; "-" to read them as lines on stdin; can be repeated in case of --pack.
                // Input files can also be given as stand-alone arguments at the end, e.g. as a result of glob
                {"output",     required_argument, 0, 'o'}, // output file; "-" for stdout (default)

                // Options for --pack
                {"compress",   required_argument, 0, 'c'}, // 0 = no, 1 = heatshrink (def)
                {"level",      required_argument, 0, 'l'}, // Heatshrink compression level 1-9, -1 = default
                {"gzip",       optional_argument, 0, 'z'}, // Gzipped extensions; no arg = default, "*" = all, comma-separated list = custom exts
                {"strip",      required_argument, 0, 's'}, // path removed from all input files, e.g. when they are in a subfolder

                // Options for --embed
                {"varname",    required_argument, 0, 'n'}, // name of the array; the length variable is {varname}_len

                { /* end marker */ }
        };

        c = getopt_long(argc, argv, "h?VPLX:Mi:o:c:l:z::s:n:",
                        long_options, &option_index);
        if (c == -1) {
            break;
        }

        switch (c) {
            case 'h':
            case '?':
                show_help(0, argv);

            case 'V':
                show_version(argv);

            case 'P':
                opmode = OM_PACK;
                break;

            case 'L':
                opmode = OM_LIST;
                break;

            case 'X':
                opmode = OM_EXTRACT;
                if (extractFileName) {
                    espfs_error("can extract only one file at a time!");
                    exit(1);
                }
                extractFileName = strdup(optarg);
                break;

            case 'M':
                opmode = OM_EMBED;
                break;

            case 'i':
                if (0 == strcmp(optarg, "-")) {
                    read_from_stdin = true;
                } else {
                    queueInputFile(optarg);
                }
                num_input_files++;
                break;

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

            case 'c':
                errno = 0;
                compType = (int) strtol(optarg, NULL, 10);
                if (errno != 0 || compType < 0 || compType > 1) {
                    espfs_error("Bad compression mode: %s", optarg);
                    exit(1);
                }
                break;

            case 'l':
                errno = 0;
                compLvl = (int) strtol(optarg, NULL, 10);
                if (errno != 0 || compLvl < 1 || compLvl > 9) {
                    espfs_error("Bad compression level: %s", optarg);
                    exit(1);
                }
                break;

            case 'z':
                use_gzip = true;
                if (optarg) {
                    if (0 == strcmp("*", optarg)) {
                        s_gzipAll = true;
                    } else {
                        parseGzipExtensions(optarg);
                    }
                } else {
                    parseGzipExtensions(strdup(DEFAULT_GZIP_EXTS)); // memory leak! ehh
                }
                break;

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

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

            default:
                espfs_error("Unknown option: %c", c);
                exit(1);
        }
    }

    if (!use_gzip) {
        s_gzipExtensions = NULL;
    }

    bool want_output;
    bool allows_multiple_inputs;

    switch (opmode) {
        case OM_PACK:
            want_output = true;
            allows_multiple_inputs = true;
            break;

        case OM_LIST:
            want_output = false;
            allows_multiple_inputs = false;
            break;

        case OM_EXTRACT:
        case OM_EMBED:
            want_output = true;
            allows_multiple_inputs = false;
            break;

        default:
            espfs_error("Specify one of the operation modes: -P, -L, -X, -M");
            exit(1);
    }


    /* Input */

    while (optind < argc) {
        char *s = argv[optind++];
        if (0 == strcmp(s, "-")) {
            read_from_stdin = true;
        } else {
            queueInputFile(s);
            num_input_files++;
        }
    }

    if (num_input_files == 0 && read_from_stdin && opmode == OM_PACK) {
        read_from_stdin = false;
        espfs_dbg("Reading input file names from stdin");
        while (fgets(inputFileName, sizeof(inputFileName), stdin)) {
            //Kill off '\n' at the end
            inputFileName[strlen(inputFileName) - 1] = 0;
            queueInputFile(inputFileName);
            num_input_files++;
        }
    }

    if (!read_from_stdin) {
        if (num_input_files == 0) {
            if (allows_multiple_inputs) {
                espfs_error("Specify input file(s)!");
            } else {
                espfs_error("Specify input file!");
            }
            exit(1);
        } else if (!allows_multiple_inputs && num_input_files > 1) {
            espfs_error("Mode %c requires exactly one input file!", opmode);
            exit(1);
        }
    }

    char *inFileName = read_from_stdin ? NULL : s_inputFiles->name;


    /* Output */

    if (!want_output && outFileName) {
        espfs_error("Output file is not allowed in %c mode!", opmode);
        exit(1);
    }

    FILE *outFile = NULL;
    bool write_to_stdout = outFileName && (0 == strcmp("-", outFileName));
    if (outFileName && !write_to_stdout) {
        espfs_dbg("Writing to %s", outFileName);
        outFile = fopen(outFileName, "w+");
        if (!outFile) {
            perror(outFileName);
            return 1;
        }
        s_outFd = fileno(outFile);
        ftruncate(s_outFd, 0);
    } else {
        if (outFileName) {
            free(outFileName);
            outFileName = NULL;
        }
        if (want_output && !write_to_stdout) {
            espfs_error("Specify output file! Use -o - for stdout");
            exit(1);
        }
    }


    /* Do it! */

    switch (opmode) {
        case OM_PACK: {
            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);
                    if (f > 0) {
                        // 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++; }

                        const char *compName = "unknown";
                        rate = handleFile(f, embeddedName, compType, compLvl, &compName);
                        (void) rate;
                        espfs_dbg("%s (%d%%, %s)", embeddedName, rate, compName);
                        close(f);
                    } else {
                        snprintf(tempbuf, BUFLEN, "Open %s for reading: %s", name, strerror(errno));
                        espfs_error("%s", tempbuf);
                        exit(1);
                    }
                } else if (serr != 0) {
                    snprintf(tempbuf, BUFLEN, "Stat %s: %s", name, strerror(errno));
                    espfs_error("%s", tempbuf);
                    exit(1);
                }

                entry = entry->next;
            }

            finishArchive();
            fsync(s_outFd);
            if (outFile) {
                fclose(outFile);
            }
        } break;

        case OM_LIST: {
            if (!inFileName) {
                espfs_error("Input file required!");
                exit(1);
            }
            parseEspfsImage(inFileName, NULL, s_outFd);
        } break;

        case OM_EXTRACT: {
            if (!inFileName) {
                espfs_error("Input file required!");
                exit(1);
            }
            parseEspfsImage(inFileName, extractFileName, s_outFd);
        } break;

        case OM_EMBED: {
            FILE *inFile = NULL;
            int inFD = STDIN_FILENO;

            if (!read_from_stdin) {
                inFile = fopen(inFileName, "r");
                if (!inFile) {
                    snprintf(tempbuf, BUFLEN, "Open %s for reading: %s", inFileName, strerror(errno));
                    espfs_error("%s", tempbuf);
                    exit(1);
                }
                inFD = fileno(inFile);
            }

            if (!c_varname) {
                c_varname = strdup(DEFAULT_C_VARNAME);
            }

            size_t len;
            len = (size_t) snprintf(tempbuf, BUFLEN, "unsigned char %s[] = {", c_varname);
            write(s_outFd, tempbuf, len);
            uint8_t u;
            size_t imageLen = 0;
            while (1 == read(inFD, &u, 1)) {
                len = (size_t) snprintf(tempbuf, BUFLEN, "%s0x%02x,", ((imageLen % 16) ? " " : "\n  "), u);
                write(s_outFd, tempbuf, len);
                imageLen++;
            }
            len = (size_t) snprintf(tempbuf, BUFLEN, "\n};\nunsigned int %s_len = %lu;\n", c_varname, imageLen);
            write(s_outFd, tempbuf, len);

            fsync(s_outFd);
            if (outFile) {
                fclose(outFile);
            }
            if (inFile) {
                fclose(inFile);
            }
        } break;

        default:
            __builtin_unreachable();
    }

    return 0;
}

__attribute__((noreturn))
void show_version(char **argv)
{
    // to stdout
    printf("%s #%s\n", argv[0], GIT_HASH);
    exit(0);
}

__attribute__((noreturn))
void show_help(int retval, char **argv)
{
    // to stderr, this can happen as a response to bad args

    //               ##########**********##########**********##########----------$$$$$$$$$$----------$$$$$$$$$$----------| 80 chars
    fprintf(stderr, "%s #%s - Program to create and parse espfs images\n", argv[0], GIT_HASH);
    fprintf(stderr, "\n"
                    "One main operation mode must be selected:\n"
                    "  -P, --pack              Create a binary fs image\n"
                    "  -L, --list              Read a binary fs image and show the contents\n"
                    "  -X, --extract=NAME      Read a binary fs image and extract a file with the given name\n"
                    "  -M, --embed             Read a binary file (typically the binary fs image produced by -P) and\n"
                    "                          convert it to C syntax byte array and a constant with its length.\n"
                    "  -h, -?, --help          Show this help (has precedence over other options)\n"
                    "\n"
                    "Additional arguments specify the input, output and parameters:\n"
                    "  -i, --input=FILE        Input file name. Can be used multiple times. For convenience with globs,\n"
                    "                          input files can be given at the end of the option string without -i.\n"
                    "                          The \"-\" filename means \"read from stdin\":\n"
                    "                           - pack: reads file names to pack from stdin (e.g. piped from `find`)\n"
                    "                           - embed: read the binary file from stdin (e.g. piped from the --pack\n"
                    "                             option, avoiding the creation of a temporary binary file)\n"
                    "                          Stdin reading is not supported for options --extract and --list.\n"
                    "\n"
                    "  -o, --output=FILE       Output file, \"-\" for stdout.\n"
                    "                          Output can't be changed in --list mode.\n"
                    "\n"
                    "Pack options:\n"
                    "  -c, --compress=MODE     Compression mode, 0=none, 1=heatshrink (default). Not used if the file\n"
                    "                          is gzipped - additional compression wouldn't be useful.\n"
                    "  -l, --level=LEVEL       Compression level, 1-9, -1=default. 1=worst, 9=best compression, but uses\n"
                    "                          more RAM to unpack.\n"
                    "  -s, --strip=PATH        Strip a path prefix from all packed files (e.g. a subfolder name)\n"
                    "  -z, --gzip[=EXTS]       Enable gzip compression for some file types (filename extensions).\n"
                    "                          By default, enabled for "DEFAULT_GZIP_EXTS".\n"
                    "                          Set EXTS to * to enable gzip for all files. Custom extensions are set as\n"
                    "                          a comma-separated list.\n"
                    "\n"
                    "Embed options:\n"
                    "  -n, --varname=VARNAME   Set custom array name to use in the generated C code. The length constant\n"
                    "                          is called {VARNAME}_len. The default VARNAME is \""DEFAULT_C_VARNAME"\"\n"
    );

    exit(retval);
}