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

#include <boost/program_options.hpp>

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

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

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.1.0-dev");

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 usage()
{
    std::cout << BOLD << "Read Test Utility\n\n" << ENDC;
    std::cout << "Usage: read_test <--help | -h>\n";
    std::cout << "       read_test --version\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);    
//    BlackLynx::SearchLynx::CharacterUtilities::lowercaseInPlace(val);
    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 filename;
    uint32_t numThreads = 0;
    
    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>(&filename)->required(), 
             "input file")
            ;
     
        // 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:   " << filename << "\n";

    // Start time measurement
    auto start = std::chrono::system_clock::now();
    
    // Prime the config object then allocate buffers
    auto slConfig = bsl::SearchLynxConfiguration::getInstance();    
    auto bufAllocator = bsl::BufferAllocator::getInstance();
    bufAllocator->initAllocation();
    
    // 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(filename, numThreads);
            break;
            
        case ReadType::NORMAL:
            // Set default number of threads
            if (numThreads == 0)
            {
                numThreads = 22;
            }
            driver = new bsl::ExecDriver(filename, numThreads);
            break;
            
        case ReadType::URING:
            numThreads = 1;
            driver = new bsl::URingExecDriver(filename, 1);
            break;
    }
    std::cout << "Num threads:  " << numThreads << "\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;

    uint64_t fileSize = bsl::FileUtils::getFileSize(filename);
    
    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;
}