#include <iostream>
#include <fstream>
#include <string>
#include <vector>
#include <sstream>
#include <iterator>
#include <locale>
#include <codecvt>
#include <filesystem>
#include <stdlib.h>
#include <io.h>
#include <fcntl.h>
#include <share.h>
#include <sys\stat.h>
#include <string.h>

#include <windows.h>
#include <libloaderapi.h>

#include "oodle2.h"

const LPCWSTR LIB_NAME = L"oo2core_8_win64";
// Printing a `LPCWSTR` is a pain, so it's much easier to just have this second variable.
// It's not like this is going to change often.
const char* LIB_FILE = "oo2core_8_win64.dll";

typedef decltype(OodleLZ_Compress)* compress_func;
typedef decltype(OodleLZ_Decompress)* decompress_func;
typedef decltype(OodleCore_Plugins_SetPrintf)* setprintf_func;

typedef std::uint8_t u8;
typedef std::uint32_t u32;
typedef std::uint64_t u64;

const u32 CHUNK_RAW_SIZE = 512 * 1024;


void usage() {
    std::cerr <<
        "Usage: oodle-cli [OPTIONS] [--] <operation>\n"
        "Perform Oodle 2 compression or decompression for Darktide data.\n"
        "Deflated data is a stream of chunks with 512k uncompressed size.\n"
        "Inflated data is an arbitrary byte stream.\n"
        "Data is read from stdin and written to stdout\n"
        "\n"
        "<operation> may be either 'compress' or 'decompress'.\n"
        "\n"
        "Options:\n"
        "  -v            Increase Oodle's verbosity.\n"
        "                May be specified up to three times.\n"
        "  -c            Only check if the library can be found and used.\n"
        "  --fuzz-safe   Use fuzz safe decompression.\n"
        "  --check-crc   Check CRC during decompression.\n"
        "  --padding     The amount of padding to assume at the start of the data.\n"
        "                Set this to `(x % 16)`, where `x` is the position at which\n"
        "                the data passed to oodle - cli starts in the source file.\n"
        "  --help        Show this help message.\n"
        "\n"
        "The following environmental variables are recognized:\n"
        "  DLL_SEARCH_PATH: A color (':') separated list of directories to search\n"
        "                   for the Oodle library. May be relative to the current\n"
        "                   working directory.\n";
}


void printf_callback(int verboseLevel, const char* file, int line, const char* fmt, ...) {
    std::cout << "[OODLE] ";

    // For some reason, this doesn't actually map to the config value.
    switch (verboseLevel) {
    case 0:
        std::cout << "EXTRAORDINARILY IMPORTANT: ";
        break;
    case 1:
        std::cout << "EXTERMELY IMPORTANT: ";
        break;
    case 2:
        std::cout << "VERY IMPORTANT: ";
        break;
    }

    va_list args;
    va_start(args, fmt);
    vfprintf(stderr, fmt, args);
}


u32 read_u32(FILE* stream) {
    u32 buf;
    if (fread_s(&buf, 4, sizeof(u8), 4, stream) < 4) {
        throw "ERROR: Failed to read u32 from stream!";
    }

    return _byteswap_ulong(buf);
}


size_t write_u32(FILE* stream, u32 val) {
    val = _byteswap_ulong(val);

    if (fwrite(&val, 4, sizeof(u8), stream) < 4) {
        throw "ERROR: Failed to read u32 from stream!";
    }

    return 4;
}

size_t write_padding(FILE* stream, size_t offset) {
    size_t pos = ftell(stream) + offset;
    if (errno != 0) {
        return 0;
    }

    size_t padding_size = 16 - ((pos) % 16);
    if (padding_size < 16) {
        u64 padding = 0x0;

        _set_errno(0);
        fwrite(&padding, sizeof(u8), padding_size, stream);

        return padding_size;
    }

    return 0;
}


int do_compress(HMODULE hDLL, FILE* in_file, FILE* out_file, size_t padding_start) {
    auto fn = reinterpret_cast<compress_func>(GetProcAddress(hDLL, "OodleLZ_Compress"));
    if (fn == NULL) {
        std::cerr << "ERROR: The library is incompatible, no 'OodleLZ_Compress'!\n";
        return 1;
    }

    // I'm not sure how Fatshark does this efficiently.
    // Within each chunk, there is a padding whos size depends on the position where the padding is written.
    // And that position is affected by both the size of previous chunks and the total number of chunks, since
    // list of chunk sizes is written before all chunks.
    // Because of that, all chunks have to be compressed and buffered before any of them can be written to the output.
    std::vector<u32> chunk_sizes;
    std::vector<u8> final_buffer(CHUNK_RAW_SIZE);

    u8* write_address = final_buffer.data();

    // This tracks the total amount of data written to `final_buffer`.
    // It also doubles as the current position in the output stream, required
    // to calculate padding sizes.
    size_t final_size = 0;

    // Re-use the same buffers for each chunk.
    // Allocating the full CHUNK_RAW_SIZE for compressed buffer is pretty
    // much guranteed to be too much, but only slightly, and more performant than
    // re-allocating anyways.
    u8* raw_buffer = new u8[CHUNK_RAW_SIZE];

    size_t chunk_index = 0;
    boolean is_eof = false;

    do {
        chunk_index++;

        _set_errno(0);
        size_t read = fread_s(raw_buffer, CHUNK_RAW_SIZE, sizeof(u8), CHUNK_RAW_SIZE, in_file);
        if (errno != 0) {
            char msg[80];
            _strerror_s(msg, 80, NULL);
            std::cerr << "ERROR: Failed to read chunk " << chunk_index << " from input file : " << msg << std::endl;
            return 1;
        }

        if (read < CHUNK_RAW_SIZE) {
            is_eof = true;
        }

        size_t remaining = CHUNK_RAW_SIZE - read;

        // Bitsquid always writes full chunks. If the last chunk can't be filled completely
        // with data the remainder is filled with `0x0`.
        memset(raw_buffer + read, 0x0, remaining);

        // Bitsquid also adds an end marker consisting of four `0x66`.
        if (remaining >= 4) {
            memset(raw_buffer + read, 0x66, 4);
        }
        else if (CHUNK_RAW_SIZE - read >= 0) {
            std::cerr << "ERROR: Not enough space left in chunk to add end marker. Don't know how to proceed.\n";
            return 1;
        }

        intptr_t res = fn(
            OodleLZ_Compressor_Kraken,
            raw_buffer,
            CHUNK_RAW_SIZE,
            write_address,
            OodleLZ_CompressionLevel_Optimal2,
            nullptr,
            0,
            nullptr,
            nullptr,
            0
        );

        if (res <= 0) {
            std::cerr << "ERROR: Failed to compress chunk " << chunk_index << "!\n";
            return 1;
        }

        chunk_sizes.push_back(res);
        write_address += res;
    } while (!is_eof);

    _set_errno(0);
    write_u32(out_file, chunk_sizes.size());
    if (errno != 0) {
        char msg[80];
        _strerror_s(msg, 80, NULL);
        std::cerr << "ERROR: Failed to write number of chunks to output file: " << msg << std::endl;
        return 1;
    }

    // `padding_start` is the point in the final file, passed from CLI.
    // The `4` is added for the u32 at the start of our output, which is the
    // number of chunks.
    size_t pos = ftell(out_file) + 4;
    _set_errno(0);
    write_padding(out_file, pos);
    if (errno != 0) {
        char msg[80];
        _strerror_s(msg, 80, NULL);
        std::cerr << "ERROR: Failed to write chunk size padding to output file: " << msg << std::endl;
        return 1;
    }

    chunk_index = 0;
    u8* read_address = final_buffer.data();

    for (auto size : chunk_sizes) {
        chunk_index++;

        _set_errno(0);
        write_u32(out_file, size);
        if (errno != 0) {
            char msg[80];
            _strerror_s(msg, 80, NULL);
            std::cerr << "ERROR: Failed to write size for chunk " << chunk_index << " to output file : " << msg << std::endl;
            return 1;
        }

        size_t pos = ftell(out_file) + padding_start;
        _set_errno(0);
        write_padding(out_file, pos);
        if (errno != 0) {
            char msg[80];
            _strerror_s(msg, 80, NULL);
            std::cerr << "ERROR: Failed to write chunk size padding to output file: " << msg << std::endl;
            return 1;
        }

        _set_errno(0);
        fwrite(&read_address, sizeof(u8), size, out_file);
        if (errno != 0) {
            char msg[80];
            _strerror_s(msg, 80, NULL);
            std::cerr << "ERROR: Failed to write data for chunk " << chunk_index << " to output file : " << msg << std::endl;
            return 1;
        }

        read_address += size;
    }


    delete[] raw_buffer;

    return 0;
}


int do_decompress(
    HMODULE hDLL,
    FILE* in_file,
    FILE* out_file,
    OodleLZ_FuzzSafe fuzz_safe,
    OodleLZ_CheckCRC check_crc,
    OodleLZ_Verbosity verbosity,
    size_t padding_start,
    size_t num_chunks
) {
    auto fn = reinterpret_cast<decompress_func>(GetProcAddress(hDLL, "OodleLZ_Decompress"));
    if (fn == NULL) {
        std::cerr << "ERROR: The library is incompatible, no 'OodleLZ_Decompress'!\n";
        return 1;
    }

    // Re-use the same buffers for each chunk.
    // Allocating the full CHUNK_RAW_SIZE for the compressed buffer is pretty
    // much guranteed to be too much, but should be more performant than the
    // exact size allocating each time. And who cares about <512k RAM wasted?
    u8* raw_buffer = new u8[CHUNK_RAW_SIZE];
    u8* compressed_buffer = new u8[CHUNK_RAW_SIZE];
    if (verbosity >= OodleLZ_Verbosity_Some) {
        std::cerr << "DEBUG: " << num_chunks << " chunks\n";
    }

    for (size_t chunk_index = 0; chunk_index < num_chunks; chunk_index++) {
        u32 chunk_size = read_u32(in_file);

        _set_errno(0);
        size_t pos = ftell(in_file) + padding_start;
        if (errno != 0) {
            std::cerr << "ERROR: Failed to get position in input file: " << strerror(errno) << std::endl;
            return 1;
        }

        size_t padding_size = 16 - (pos % 16);
        if (padding_size < 16) {
            _set_errno(0);
            if (fseek(in_file, (long)padding_size, SEEK_CUR) != 0) {
                std::cerr << "ERROR: Failed to seek input file for chunk " << chunk_index <<  ": " << strerror(errno) << std::endl;
                return 1;
            }
        }

        if (verbosity >= OodleLZ_Verbosity_Some) {
            std::cerr << "DEBUG: Chunk " << chunk_index + 1 << ": " << chunk_size << "\n";
        }

        _set_errno(0);
        if (fread_s(compressed_buffer, chunk_size, sizeof(u8), chunk_size, in_file) < chunk_size) {
            std::cerr << "ERROR: Failed to read compressed data for chunk " << chunk_index << " from input file: " << strerror(errno) << "\n";
            return 1;
        }

        intptr_t res = fn(
            compressed_buffer,
            chunk_size,
            raw_buffer,
            CHUNK_RAW_SIZE,
            fuzz_safe,
            check_crc,
            verbosity,
            nullptr,
            0,
            nullptr,
            nullptr,
            nullptr,
            0,
            OodleLZ_Decode_Unthreaded
        );
        
        if (res != CHUNK_RAW_SIZE) {
            std::cerr << "ERROR: Failed to decompress chunk " << chunk_index << "!\n";
            return 1;
        }

        _set_errno(0);
        if (fwrite(raw_buffer, sizeof(u8), CHUNK_RAW_SIZE, out_file) < CHUNK_RAW_SIZE) {
            std::cerr << "ERROR: Failed to write decompressed chunk " << chunk_index << " to output file" << strerror(errno) << "\n";
            return 1;
        }
    }

    delete[] raw_buffer;
    delete[] compressed_buffer;

    return 0;
}


int main(int argc, char* argv[])
{
    int verbosity = 0;
    bool check_lib = FALSE;
    OodleLZ_FuzzSafe fuzz_safe = OodleLZ_FuzzSafe_No;
    OodleLZ_CheckCRC check_crc = OodleLZ_CheckCRC_No;
    size_t padding_start = 0;
    size_t num_chunks = 0;
    int i = 1;

    for (; i < argc; i++) {
        char* arg = argv[i];

        if (strcmp(arg, "--") == 0 || arg[0] != '-') {
            break;
        }

        if (strcmp(arg, "--help") == 0) {
            usage();
            return 0;
        }
        else if (strcmp(arg, "-c") == 0) {
            check_lib = TRUE;
        }
        else if (strcmp(arg, "-v") == 0) {
            verbosity++;
        }
        else if (strcmp(arg, "--fuzz-safe") == 0) {
            fuzz_safe = OodleLZ_FuzzSafe_Yes;
        }
        else if (strcmp(arg, "--check-crc") == 0) {
            check_crc = OodleLZ_CheckCRC_Yes;
        }
        else if (strcmp(arg, "--padding") == 0) {
            i++;
            padding_start = std::stoi(argv[i]);
        }
        else if (strcmp(arg, "--chunks") == 0) {
            i++;
            num_chunks = std::stoi(argv[i]);
        }
        else {
            std::cerr << "ERROR: Unknown option '" << arg << "'!\n\n";
            usage();
            return 1;
        }
    }

    char* var;
    size_t len;
    DWORD load_flags = 0;
    if (_dupenv_s(&var, &len, "DLL_SEARCH_PATH") == 0) {
        if (var) {
            std::string search_path(var, len);
            std::istringstream ss(search_path);
            std::string token;

            while (std::getline(ss, token, ':')) {
                auto path = std::filesystem::path(token);
                if (!path.is_absolute()) {
                    path = std::filesystem::absolute(path);
                }

                if (!path.is_absolute() || !AddDllDirectory(path.c_str())) {
                    std::cerr << "WARN: Failed to add DLL search path: '" << token << "'!\n";
                }
                else {
                    std::cout << "INFO: Added DLL search path: '" << path << "'.\n";
                }
            }

            load_flags = LOAD_LIBRARY_SEARCH_DEFAULT_DIRS;
            free(var);
        }
    }
    else {
        std::cerr << "ERROR: Failed to read environment variable 'DLL_SEARCH_PATH'. Skipping.\n";
    }

    HINSTANCE hDLL = LoadLibraryEx(LIB_NAME, NULL, load_flags);
    if (hDLL == NULL) {
        std::cerr << "ERROR: Couldn't find library file '" << LIB_FILE << "'!\n";
        return 1;
    }

    if (check_lib) {
        std::cout << "INFO: '" << LIB_FILE << "' found and loaded.\n";
        return 0;
    }

    if (num_chunks == 0) {
        std::cerr << "ERROR: Number of chunks not specified!\n";
        usage();
        return 1;
    }

    if (argc - i < 1) {
        std::cerr << "ERROR: Arguments missing!\n\n";
        usage();
        return 1;
    }

    if (verbosity > 0) {
        auto fn = reinterpret_cast<setprintf_func>(GetProcAddress((HMODULE)hDLL, "OodleCore_Plugins_SetPrintf"));
        if (fn == NULL) {
            std::cerr << "ERROR: The library is incompatible, no 'OodleCore_Plugins_SetPrintf'!\n";
            return 1;
        }
        reinterpret_cast<void*>(fn(printf_callback));
    }

    std::string operation = argv[i];
    FILE* in_file = nullptr;
    FILE* out_file = nullptr;

    if (argc - i >= 2) {
        errno_t err = fopen_s(&in_file, argv[i + 1], "rbS");

        if (err != 0) {
            std::cerr << "ERROR: Failed to open input file: " << strerror(errno) << "\n";
            return 1;
        }
    }
    else {
        if (_setmode(_fileno(stdin), _O_BINARY) < 0) {
            std::cerr << "ERROR: Failed to prepare stdin for binary reading!\n";
            return 1;
        }
        in_file = stdin;
    }

    if (argc - i >= 3) {
        errno_t err = fopen_s(&out_file, argv[i + 2], "wbS");

        if (err != 0) {
            std::cerr << "ERROR: Failed to open output file: " << strerror(errno) << "\n";
            return 1;
        }
    }
    else {
        if (_setmode(_fileno(stdout), _O_BINARY) < 0) {
            std::cerr << "ERROR: Failed to prepare stdout for binary writing!\n";
            return 1;
        }
        out_file = stdout;
    }

    int res;

    if (operation == "compress") {
        res = do_compress(hDLL, in_file, out_file, padding_start);
    }
    else {
        res = do_decompress(hDLL, in_file, out_file, fuzz_safe, check_crc, (OodleLZ_Verbosity)verbosity, padding_start, num_chunks);
    }

    if (in_file != stdin) {
        fclose(in_file);
    }

    if (out_file != stdout) {
        fclose(out_file);
    }

    return res;
}