#include <spdlog/cfg/env.h>
#include <spdlog/sinks/stdout_color_sinks.h>
#include <spdlog/spdlog.h>
#include <time.h>

#include <argparse/argparse.hpp>
#include <iostream>
#include <map>
#include <set>
#include <string>

#include "fastsort_version.h"
#include "sort/sort_args.h"
#include "util/profiling.h"

namespace nsgv {
extern SortArg gSortArg;
};

extern int doSort(); // sort 函数声明

// 获取当前时间戳
string getCurrentTimeStr() {
    time_t time_val;
    struct tm *at;
    char now[80];
    time(&time_val);
    at = localtime(&time_val);
    strftime(now, 79, "%B %d, %Y at %I:%M:%S %p %Z", at);
    return string(now);
}

// 入口函数
int main(int argc, char *argv[]) {
    // init log
    spdlog::set_default_logger(spdlog::stderr_color_st("fastsort"));
    spdlog::cfg::load_env_levels();

    // init arg parser
    argparse::ArgumentParser program("FastSort", FASTSORT_VERSION,
                                     argparse::default_arguments::none);
    program.add_description(
        "Sort SAM/BAM files in user specific order such as coordinate, query-name etc.");

    program.add_argument("-i", "--input")
        .help("Input file. One SAM or BAM file.")
        .metavar("<INPUT>")
        .required();

    program.add_argument("-o", "--output")
        .help("Output file. SAM or BAM file to write sorted SAM/BAM to.")
        .metavar("<OUTPUT>")
        .required();

    program.add_argument("--num-threads")
        .help("Number of threads to use.")
        .scan<'i', int>()
        .default_value(1)
        .nargs(1)
        .metavar("<NUM_THREADS>");

    program.add_argument("--none-duplex-io")
        .help("Do not use writing-while-reading mode.")
        .default_value(false)
        .implicit_value(true);

    program.add_argument("--create-index")
        .help("Whether to create an index when writing sorted BAM output.(only work when sort type is coordinate)")
        .default_value(false)
        .implicit_value(true);

    program.add_argument("--index-format")
        .help("Format for bam index file. Possible values: {BAI, CSI}")
        .default_value(std::string("BAI"))
        .choices("BAI", "CSI")
        .nargs(1)
        .metavar("<IndexFormat>");

    program.add_argument("-l", "--compress-level")
        .help("Set compression level, from 0 (uncompressed) to 9 (best).")
        .nargs(1);

    program.add_argument("-m", "--mem")
        .help("approximate total memory limit for all threads (by default 4GB)")
        .default_value(std::string("4G"))
        .nargs(1)
        .metavar("<MEM>");

    program.add_argument("--sort-type")
        .help("\nUse coordinate sort or query name sort. Possible values: {COORDINATE, NAME}")
        .default_value(std::string("COORDINATE"))
        .choices("COORDINATE", "NAME")
        .nargs(1)
        .metavar("<SORT_TYPE>");

    program.add_argument("--query-name")
        .help("\nChoose the query name sort type when sorting by quey name. Possible values: {PICARD, NATURAL, ASCII}")
        .default_value(std::string("PICARD"))
        .choices("PICARD", "NATURAL", "ASCII")
        .nargs(1)
        .metavar("<QUERY_NAME>");

    // add help and version args
    program.add_argument("-h", "--help")
        .action([&](const auto & /*unused*/) {
            std::cout << program.help().str();
            std::exit(0);
        })
        .default_value(false)
        .help("shows help message and exits")
        .implicit_value(true)
        .nargs(0);

    program.add_argument("-v", "--version")
        .action([&](const auto & /*unused*/) {
            std::cout << FASTSORT_VERSION << std::endl;
            std::exit(0);
        })
        .default_value(false)
        .help("prints version information and exits")
        .implicit_value(true)
        .nargs(0);

    // std::cout << program << std::endl;

    nsgv::gSortArg.START_TIME = getCurrentTimeStr();
    nsgv::gSortArg.CLI_STR = argv[0];
    for (int i = 1; i < argc; ++i) {
        nsgv::gSortArg.CLI_STR += " " + std::string(argv[i]);
    }

    try {
        program.parse_args(argc, argv);
        nsgv::gSortArg.INPUT_FILE = program.get("--input");
        nsgv::gSortArg.OUTPUT_FILE = program.get("--output");
        nsgv::gSortArg.NUM_THREADS = program.get<int>("--num-threads");
        nsgv::gSortArg.DUPLEX_IO = !program.get<bool>("--none-duplex-io");
        nsgv::gSortArg.CREATE_INDEX = program.get<bool>("--create-index");
        nsgv::gSortArg.INDEX_FORMAT =
            program.get("--index-format") == "BAI" ? nsmd::IndexFormat::BAI : nsmd::IndexFormat::CSI;
        nsgv::gSortArg.SORT_COORIDINATE = program.get("--sort-type") == "COORDINATE";

        // 三种排序方式
        std::map<std::string, nsmd::QueryNameType> query_name_args = {
            {"NATURAL,", nsmd::QueryNameType::NATURAL},
            {"ASCII", nsmd::QueryNameType::ASCII},
            {"PICARD", nsmd::QueryNameType::PICARD}};

        nsgv::gSortArg.QUERY_NAME_TYPE = query_name_args[program.get("--query-name")];

        // 设置缓存大小
        string maxMemStr = program.get("--mem");
        {
            char* q;
            uint64_t maxMem = strtol(maxMemStr.c_str(), &q, 0);
            if (*q == 'k' || *q == 'K')
                maxMem <<= 10;
            else if (*q == 'm' || *q == 'M')
                maxMem <<= 20;
            else if (*q == 'g' || *q == 'G')
                maxMem <<= 30;
            nsgv::gSortArg.MAX_MEM = maxMem;
        }

    } catch (const std::exception &err) {
        spdlog::error(err.what());
        std::cerr << std::endl << program;
        return 1;
    }

/**
 * sort 流程
 * 第一阶段（phase-1）：
 * 1. 读入bam文件header（要注意如果是sam或者管道输入则不需要解压，如果是bam，需要考虑samtools和gatk的压缩方式不同，gatk不会block对齐）
 * 2. phase-1-stage-1: 读入待解压的数据块，直到填满phase-1-stage-1的buffer，读入数据期间分割压缩的block
 * 3. phase-1-stage-2: 多线程解压数据块，每个线程解压后放在自己的buf里，乱序解压，线程buf满后，进行线程内的排序
 * 4. phase-1-stage-3: 线程间进行归并排序，然后压缩写入中间文件，清空线程buf，继续读入下一批数据块，直到文件读完
 * 第二阶段（phase-2）：
 * 1. 打开所有中间文件，并关联上解压线程池
 * 2. phase-2-stage-1：解压中间文件，并进行归并排序，放入buf
 * 3. phase-2-stage-2：多线程进行压缩，放入buf
 * 4. phase-2-stage-3：将buf写入最终的bam文件
 * 5. 要注意计算index并写入index文件
 */

    spdlog::info("fast bam sort start");
    
    doSort();
    
    spdlog::info("fast bam sort end");

    displayProfiling(nsgv::gSortArg.NUM_THREADS);

    return 0;
}