#include <algorithm>
#include <chrono>
#include <cstdlib>
#include <filesystem>
#include <iomanip>
#include <iostream>
#include <string>
#include <vector>

#include <sys/utsname.h>

#include <boost/program_options.hpp>

#include "AIO_ExecDriver.h"
#include "ExecDriver.h"
#include "URingExecDriver.h"
#include "Utilities/BufferAllocator.h"
#include "Utilities/FileUtils.h"

// Alias the namespace for readability
namespace bpo = boost::program_options;
namespace bsl = BlackLynx::SearchLynx;
namespace sfs = std::filesystem;

const std::string BOLD("\033[1m");
const std::string ENDC("\033[0m");
const std::string RED("\033[31m");
const std::string YELLOW("\033[33m");
const std::string UP_ONE = "\033[1A";

const std::string VERSION("ReadTest version 1.2.0-dev");

const std::vector<std::string> UNITS{"b/sec", "KB/s", "MB/s", "GB/s", "TB/s"};

double printRate(double rawRate, uint32_t& unitIdx)
{
    uint32_t idx = 0;
    double rate = rawRate;
    while (rate > 1024)
    {
        rate /= 1024;
        ++idx;
    }
    
    unitIdx = idx;
    return rate;
}

void getKernelVersion(uint32_t& major, uint32_t& minor, uint32_t& bugfix)
{
    // Read "uname" info
    struct utsname unameBuf;
    uname(&unameBuf);
    
    // Release contains kernel version; e.g. 5.15.0-33-generic
    char* startPtr = unameBuf.release;
    char* endPtr = nullptr;
    major = std::strtoul(startPtr, &endPtr, 10);
    
    startPtr = endPtr + 1;
    minor = std::strtoul(startPtr, &endPtr, 10);
    
    startPtr = endPtr + 1;
    bugfix = std::strtoul(startPtr, &endPtr, 10);
    
#if 0
    std::cout << "MAJOR:  " << major << std::endl;
    std::cout << "MINOR:  " << minor << std::endl;
    std::cout << "BUGFIX: " << bugfix << std::endl;

    std::cout << "sysname:  " << unameBuf.sysname << std::endl;
    std::cout << "nodename: " << unameBuf.nodename << std::endl;
    std::cout << "release:  " << unameBuf.release << std::endl;
    std::cout << "version:  " << unameBuf.version << std::endl;
    std::cout << "machine:  " << unameBuf.machine << std::endl;
#endif
}

bool validateKernelVersion(uint32_t major, uint32_t minor, uint32_t bugfix)
{
    uint32_t majorVer, minorVer, bugfixVer;
    getKernelVersion(majorVer, minorVer, bugfixVer);

    if (majorVer >= major)
    {
        if (minorVer >= minor)
        {
            if (bugfixVer >= bugfix)
            {
                return true;
            }
        }
    }
    
    return false;
}

void usage()
{
    std::cout << BOLD << "Read Test Utility\n\n" << ENDC
              << "Usage: read_test <--help | -h>\n"
              << "       read_test --version\n\n"
              << "Options:\n"
              << "  --input-file    Path to the file to read.\n"
              << "  --read-type     The type of read to perform: auto, normal, aio, or uring\n"
              << "  --num-threads   The number of reader threads to use.\n"
              << "  --chunk-size    Chunk size in kilobytes.\n\n";
}

template <typename T>
void lowercase(std::basic_string<T>& str)
{    
    std::transform(str.begin(), str.end(), str.begin(), tolower);
}

// Type used to set "read-type" value
enum class ReadType
{
    DEFAULT_INVALID,
    AUTO,
    NORMAL,
    AIO,
    URING
};

struct ReadTypeValue 
{
    ReadTypeValue()
    {
    }
    
    ReadTypeValue(std::string const& val)
        : value(val)
    { 
    }
    
    ReadType getValue()
    {
        if (value == "auto")
        {
            return ReadType::AUTO;
        }
        else if (value == "normal")
        {
            return ReadType::NORMAL;
        }
        else if (value == "aio")
        {
            return ReadType::AIO;
        }
        else if (value == "uring")
        {
            return ReadType::URING;
        }
        return ReadType::DEFAULT_INVALID;
    }
    
    std::string value;
};

void validate(
    boost::any& v, 
    std::vector<std::string> const& values,
    ReadTypeValue*,
    int)
{
    // Make sure no previous assignment to 'v' was made.
    bpo::validators::check_first_occurrence(v);
    // Extract the first string from 'values'. If there is more than one string,
    // it's an error, and exception will be thrown.
    std::string val = bpo::validators::get_single_string(values);    
    lowercase(val);
    // Check for valid values
    if (val == "auto" || val == "normal" || val == "aio" || val == "uring")
    {
        v = boost::any(ReadTypeValue(val));
    } 
    else 
    {
        throw bpo::validation_error(bpo::validation_error::invalid_option_value);
    }
}

int main(int argc, char** argv) 
{
    bpo::variables_map optionsMap;
    ReadTypeValue readType("default");
    std::string filePath;
    uint32_t numThreads = 0;
    uint32_t chunkSize = 65536; // in Kb
    
    try
    {       
        // General options
        bpo::options_description generalOpts("General options");
        generalOpts.add_options()
            ("help,h", "display this help message")
                
            ("version", 
             "output version number and exit")
                
            ("read-type", bpo::value<ReadTypeValue>(&readType),
             "read type")
        
            ("num-threads", bpo::value<uint32_t>(&numThreads),
             "number reader threads")
        
            ("input-file", bpo::value<std::string>(&filePath)->required(), 
             "input file")
        
            ("chunk-size", bpo::value<uint32_t>(&chunkSize), "chunk size in Kb")
            ;
     
        // Positional options
        bpo::positional_options_description posOpts;
        posOpts.add("input-file", 1);
        
        bpo::parsed_options parsed = bpo::command_line_parser(argc, argv).
                                        options(generalOpts).positional(posOpts)
                                        .run();
        
        bpo::store(parsed, optionsMap);
        
        // Display help if the option is listed or if no arguments are provided
        if (optionsMap.count("help"))// || (argc == 1)) 
        {
            usage();
            return 1;
        }
        
        // Display version and exist if option listed
        if (optionsMap.count("version"))
        {
            std::cout << VERSION << std::endl;
            return 2;
        }
        
        // Notify any errors
        bpo::notify(optionsMap);
    }
    catch (std::exception& e) 
    {
        std::cerr << BOLD << RED << "\nERROR: " << ENDC << e.what() << "\n"
                  << std::endl;
        return -1;
    }
    
    std::cout << "Read Test Utility\n" << std::endl;
    std::cout << "Read type:    " << readType.value << "\n";
    std::cout << "Input file:   " << filePath << "\n";

    // Validate kernel version if user selected uring read type
    if (readType.getValue() == ReadType::URING && 
        (! validateKernelVersion(5, 6, 0)))
    {
        uint32_t majorVer, minorVer, bugfixVer;
        getKernelVersion(majorVer, minorVer, bugfixVer);
        std::cerr << BOLD << RED << "\nERROR: " << ENDC << "The uring read type "
                  << "requires kernel version >= 5.6; the current kernel version"
                  << " " << majorVer << "." << minorVer << " is too old.\n"
                  << std::endl;
        return -1;
    }
    
    // Start time measurement
    auto start = std::chrono::system_clock::now();
    
    uint64_t fileSize = sfs::file_size(filePath);
    
    // Prime the config object then allocate buffers
    auto bufAllocator = bsl::BufferAllocator::getInstance();
    bufAllocator->initAllocation(chunkSize, fileSize);
    
    // Use the read type argument to decide which ExecDriver instance to create
    bsl::ExecDriver* driver = nullptr;
    switch (readType.getValue())
    {
        case ReadType::DEFAULT_INVALID:
        case ReadType::AUTO:
        case ReadType::AIO:
            // Set default number of threads
            if (numThreads == 0)
            {
                numThreads = 3;
            }
            driver = new bsl::AIO_ExecDriver(filePath, numThreads, chunkSize);
            break;
            
        case ReadType::NORMAL:
            // Set default number of threads
            if (numThreads == 0)
            {
                numThreads = 22;
            }
            driver = new bsl::ExecDriver(filePath, numThreads, chunkSize);
            break;
            
        case ReadType::URING:
//            numThreads = 1;
            driver = new bsl::URingExecDriver(filePath, numThreads, chunkSize);
            break;
    }
    std::cout << "Num threads:  " << numThreads << "\n";
    std::cout << "Chunk size:   " << chunkSize << "\n";
    
    driver->start();
    driver->waitUntilComplete();
    
    // Deallocate buffers
    bufAllocator->deallocateBuffers();
    
    // End time measurement
    auto end = std::chrono::system_clock::now();
    auto diff = end - start;
    double duration = std::chrono::duration <double>(diff).count() ;
    std::cout << "Duration:     "  << std::fixed << std::setprecision(3) 
              << duration << " ms" << std::endl;

    double rawRate = fileSize / duration;
    std::cout << "Raw rate:     "  << rawRate << " b/sec" << std::endl;
    
    uint32_t idx = 0;
    double rate = printRate(rawRate, idx);
    std::cout << "Rate:         "  << std::fixed << std::setprecision(3) << rate 
              << " " << UNITS[idx] << std::endl;
    
    // Clean up
    delete driver;
    driver = nullptr;
    
    return 0;
}