#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <stdbool.h>
#include <unistd.h>
#include <getopt.h>
#include <sys/stat.h>
#include "lzma_c/LZMAEncoderApi.h"

#define METHOD_LZMA 1
#define METHOD_FLAG_CALLTRICK 0
#define METHOD_FLAG_CALLTRICK_1 0x40
#define METHOD_FLAG_CALLTRICK_2 0x80

static const char* str_usage = "Usage: %s [--nologo / -n] [--kernel / -k] [--unpack / -u] <infile> [<outfile>]\n";

static void error_load_infile() {
    printf("Error: cannot load input file\n");
    exit(1);
}

static void error_malloc_failed() {
    printf("Error: memory allocation failed\n");
    exit(1);
}

unsigned byteswap(unsigned val)
{
    val = ((val << 8) & 0xFF00FF00 ) | ((val >> 8) & 0xFF00FF ); 
    return (val << 16) | (val >> 16);
}

unsigned pack_lzma(unsigned char *infile, unsigned char *outfile, unsigned insize, unsigned char *workmem, int is_kerpack)
{
    unsigned outsize = lzma_compress(infile, outfile + 11, insize, workmem, is_kerpack);
    unsigned tmp = 0;
    memcpy(&tmp, outfile + 12, 4);
    tmp = byteswap(tmp);
    memcpy(outfile + 12, &tmp, 4);
    return outsize - 1;
}

unsigned preprocess_calltrick(unsigned char *infile, unsigned insize, unsigned char ct1[256], unsigned *ctn, unsigned char *cti)
{
    unsigned char *buf = malloc(2*insize);
    if (buf ==  NULL) {
        goto pack_calltrick_fail;
    }
    memset(ct1, 0, 256);

    unsigned char *infile_ptr = infile;
    unsigned count = insize;
    unsigned char *buf_ptr = buf;
    do {
        unsigned al = *infile_ptr++;
        al -= 0x0E8;
        if (al > 1) {
            continue;
        }
        if (count < 5) {
            break;
        }
        unsigned num;
        memcpy(&num, infile_ptr, 4);
        infile_ptr += 4;
        num += (infile_ptr - infile);
        if (num >= insize || num >= 0x1000000) { // xxx:
            infile_ptr -= 4;
            ct1[*infile_ptr] = 1;
        } else {
            count -= 4;
            num = byteswap(num);
            memcpy(infile_ptr - 4, &num, 4);
            (*ctn)++;
            memcpy(buf_ptr, &infile_ptr, 4);
            buf_ptr += 8;
        }
        
    } while(count > 0);

    unsigned idx = 0;
    while (idx < 256 && ct1[idx] != 0) {
        idx++;
    }
    if (idx >= 256) {
        goto pack_calltrick_fail;
    }
    *cti = ~(unsigned char)256;

    unsigned tmp;
    while (buf_ptr != buf) {
        buf_ptr -= 8;
        memcpy(&tmp, *buf_ptr, 4);
        memcpy(tmp - 4, cti, 1);
    }
    free(buf);
    return tmp;

pack_calltrick_fail:
    *ctn = 0;
    return 0;
}


int main(int argc, char *argv[])
{
    bool flag_nologo = false;
    bool flag_kernel = false;
    bool flag_unpack = false;
    char *infile_name = NULL;
    char *outfile_name = NULL;
    int opt;
    static struct option long_options[] = {
        {"nologo", no_argument, NULL, 'n'},
        {"kernel", no_argument, NULL, 'k'},
        {"unpack", no_argument, NULL, 'u'},
        {0, 0, 0, 0} // must be zeros here
    };
    while ((opt = getopt_long(argc, argv, "nku", long_options, NULL)) != -1) {
        switch (opt) {
            case 'n':
                flag_nologo = true;
                break;
            case 'k':
                flag_kernel = true;
                break;
            case 'u':
                flag_unpack = true;
                break;
            default: // '?'
                fprintf(stderr, str_usage, argv[0]);
                return 1;
        }
    }
    // check for required infile
    if (optind >= argc) {
        printf("Error: Input file is required.\n");
        printf(str_usage, argv[0]);
        return 1;
    }
    infile_name = argv[optind];

    // optional outfile
    if (optind + 1 < argc) {
        outfile_name = argv[optind + 1];
        optind++;
    }

    // printf("argc = %d, optind = %d\n", argc - 1, optind);
    if (argc - 1 > optind) {
        printf("Error: too many arguments\n");
        printf(str_usage, argv[0]);
        return 1;
    }

    // for debug: print parsed options and args
    printf("nologo: %s\n", flag_nologo ? "true" : "false");
    printf("kernel: %s\n", flag_kernel ? "true" : "false");
    printf("unpack: %s\n", flag_unpack ? "true" : "false");
    printf("infile: %s\n", infile_name);
    printf("outfile: %s\n", outfile_name ? outfile_name : "Not specified");
    printf("\n");

    if (!flag_nologo) {
        printf("KPackC - Kolibri Packer in C\nUses LZMA v4.32 compression library\n\n");
    }

    struct stat64 statbuf;
    int res_stat = stat64(infile_name, &statbuf);
    if (res_stat == -1) {
        error_load_infile();
    }
    // printf("statbuf.st_size = %lld\n", statbuf.st_size);

    if (statbuf.st_size == 0 || statbuf.st_size >= 4294967296LL) {
        printf("Error: input file must be non zero and less than 4G\n");
        return 1;
    }
    unsigned infile_size = (unsigned int)statbuf.st_size;
    printf("infile_size = %u\n", infile_size);

    FILE *infile_handle = fopen(infile_name, "rb");
    if (!infile_handle) {
        error_load_infile();
    }

    unsigned char *infile_buf = malloc(infile_size);
    if (infile_buf == NULL) {
        error_malloc_failed();
    }
    size_t bytes_read = fread(infile_buf, 1, infile_size, infile_handle);
    if (bytes_read != infile_size) {
        error_load_infile();
    }
    fclose(infile_handle);

    // TODO check option -k or not. Below goes packing for without -k option

    // calculate maximum size of the output:
    unsigned max_output_size = infile_size*9/8 + 0x400; // should be enough for header
    printf("max_output_size = 0x%x\n", max_output_size);
    unsigned char *buf = malloc(2*max_output_size); // allocate memory for two copies of maximum output
    if (buf == NULL) {
        error_malloc_failed();
    }

    unsigned char *outfile = buf, *outfile1 = buf, *outfilebest = buf, *outfile2 = buf + max_output_size;
    memcpy(buf, "KPCK", 4);
    memcpy(buf + 4, (void*)&infile_size, sizeof(infile_size));

    unsigned log_dict_size = (31 - __builtin_clz(infile_size - 1)) + 1;
    if (log_dict_size > 28) {
        log_dict_size = 28;
    }
    lzma_set_dict_size(log_dict_size);
    unsigned dict_size = (1 << log_dict_size);

    unsigned workmem_size = dict_size*19/2 + 0x509000;
    unsigned char *workmem = malloc(workmem_size);
    if (workmem == NULL) {
        error_malloc_failed();
    }
    
    printf("Compressing ... \n");

    outfile = outfile2;
    memcpy(outfile2, outfile1, 8);
    unsigned outsize = pack_lzma(infile_buf, outfile, infile_size, workmem, 0);
    printf("1. outsize = %u bytes\n", outsize);
    outfilebest = outfile;
    unsigned method = METHOD_LZMA;

    unsigned char ct1[256];
    unsigned ctn;
    unsigned char cti;

    // TODO ...

    return 0;
}