// oodle-cli.cpp : This file contains the 'main' function. Program execution begins and ends there.
//

#include <iostream>
#include <fstream>
#include <string>
#include <vector>
#include <sstream>
#include <iterator>
#include <locale>
#include <codecvt>
#include <filesystem>
#include <stdlib.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_Decompress)* decompress_func;
typedef decltype(OodleCore_Plugins_SetPrintf)* setprintf_func;

void usage() {
    std::cerr <<
        "Usage: oodle-cli [OPTIONS] [--] <in_file> <out_file> <uncompressed_size>\n"
        "Decompress a <in_file> to <out_file> using the Oodle algorithm.\n"
        "The size of the uncompressed data must be known.\n"
        "\n"
        "Options:\n"
        "  -v  Increase Oodle's verbosity. May be specified up to three times.\n"
        "  -c  Only check if the library can be found and used.\n"
        "\n"
        "The following environmental variables are recognized:\n"
        "  DLL_SEARCH_PATH: A color (':') separated list of directories to search for the Oodle library. May be relative to the current 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(stdout, fmt, args);
}

int main(int argc, char* argv[])
{
    int verbosity = 0;
    bool check_lib = FALSE;
    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 {
            std::cerr << "WARN: Unknown option '" << arg << "'!\n\n";
        }
    }

    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;
    }

    auto decompress = reinterpret_cast<decompress_func>(GetProcAddress((HMODULE)hDLL, "OodleLZ_Decompress"));
    if (decompress == NULL) {
        std::cerr << "ERROR: The library is incompatible!\n";
        return 1;
    }

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

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

    std::string in_name = argv[i];
    std::string out_name = argv[i + 1];
    size_t raw_buffer_size = std::stoi(argv[i + 2]);

    std::cout << "INFO: Attempting to decompress from '" << in_name << "' to '" << out_name << "' (" << raw_buffer_size << " bytes).\n";

    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!\n";
            return 1;
        }
        reinterpret_cast<void*>(fn(printf_callback));
    }

    std::ifstream in_file(in_name, std::ios::binary | std::ios::ate);
    if (!in_file) {
        std::cerr << "ERROR: Failed to open compressed file!\n";
        return 1;
    }

    std::streamsize compressed_buffer_size = in_file.tellg();
    in_file.seekg(0, std::ios::beg);

    std::vector<char> compressed_buffer(compressed_buffer_size);
    in_file.read(compressed_buffer.data(), compressed_buffer_size);

    if (in_file.fail()) {
        std::cerr << "ERROR: Failed to read compressed file!\n";
        return 1;
    }

    std::vector<char> raw_buffer(raw_buffer_size);

    int res = decompress(
        compressed_buffer.data(),
        compressed_buffer_size,
        raw_buffer.data(),
        raw_buffer_size,
        OodleLZ_FuzzSafe_Yes,
        OodleLZ_CheckCRC_No,
        (OodleLZ_Verbosity)verbosity,
        nullptr,
        0,
        nullptr,
        nullptr,
        nullptr,
        0,
        OodleLZ_Decode_Unthreaded
    );

    out_file.write(raw_buffer.data(), res);

    if (out_file.fail()) {
        std::cerr << "Failed to write output file!\n";
        return 1;
    }

    if (res != raw_buffer_size) {
        std::cerr << "Failed to decompress. Expected " << raw_buffer_size << " bytes, got " << res << "!\n";
        return 1;
    }

    std::cout << "INFO: Done!\n";
}