#include <fcntl.h>
#include <htslib/sam.h>
#include <htslib/thread_pool.h>
#include <klib/kthread.h>
#include <math.h>
#include <pthread.h>
#include <spdlog/spdlog.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/mman.h>
#include <unistd.h>
#include <zlib.h>

#include <algorithm>

#include "sam_io.h"
#include "sort_args.h"
#include "util/profiling.h"

#define BAM_BLOCK_SIZE 16L * 1024 * 1024

namespace nsgv {

SortArg gSortArg; // 参数
HeaderBuf gInHdr; // 输入文件的header
bool gIsBigEndian;

};  // namespace nsgv

struct BamSortData {
    uint16_t flag;
    uint16_t qnameLen;
    uint16_t bamLen;
    int32_t tid;
    int64_t pos;
    uint8_t *uDataPos;
    char *qname;  // pointer to qname
};

// thread data for sorting
struct SortData {
    vector<vector<BamSortData>> bamArr; // for each thread
    vector<DataBuffer> uData;  // for each thread uncompressed data
    vector<uint8_t *> *blockAddrArr;
    int round = 0;
};


int sam_realloc_bam_data(bam1_t *b, size_t desired) {
    uint32_t new_m_data;
    uint8_t *new_data;
    new_m_data = desired;
    kroundup32(new_m_data);  // next power of 2
    new_m_data += 32;        // reduces malloc arena migrations?
    if (new_m_data < desired) {
        errno = ENOMEM;  // Not strictly true but we can't store the size
        return -1;
    }
#ifdef FUZZING_BUILD_MODE_UNSAFE_FOR_PRODUCTION
    if (new_m_data > FUZZ_ALLOC_LIMIT) {
        errno = ENOMEM;
        return -1;
    }
#endif
    if ((bam_get_mempolicy(b) & BAM_USER_OWNS_DATA) == 0) {
        new_data = (uint8_t*)realloc(b->data, new_m_data);
    } else {
        if ((new_data = (uint8_t *)malloc(new_m_data)) != NULL) {
            if (b->l_data > 0)
                memcpy(new_data, b->data, b->l_data < b->m_data ? b->l_data : b->m_data);
            bam_set_mempolicy(b, bam_get_mempolicy(b) & (~BAM_USER_OWNS_DATA));
        }
    }
    if (!new_data)
        return -1;
    b->data = new_data;
    b->m_data = new_m_data;
    return 0;
}

static inline int realloc_bam_data(bam1_t *b, size_t desired) {
    if (desired <= b->m_data)
        return 0;
    return sam_realloc_bam_data(b, desired);
}

// Fix bad records where qname is not terminated correctly.
static int fixup_missing_qname_nul(bam1_t *b) {
    bam1_core_t *c = &b->core;

    // Note this is called before c->l_extranul is added to c->l_qname
    if (c->l_extranul > 0) {
        b->data[c->l_qname++] = '\0';
        c->l_extranul--;
    } else {
        if (b->l_data > INT_MAX - 4)
            return -1;
        if (realloc_bam_data(b, b->l_data + 4) < 0)
            return -1;
        b->l_data += 4;
        b->data[c->l_qname++] = '\0';
        c->l_extranul = 3;
    }
    return 0;
}

static void bam_cigar2rqlens(int n_cigar, const uint32_t *cigar, hts_pos_t *rlen, hts_pos_t *qlen) {
    int k;
    *rlen = *qlen = 0;
    for (k = 0; k < n_cigar; ++k) {
        int type = bam_cigar_type(bam_cigar_op(cigar[k]));
        int len = bam_cigar_oplen(cigar[k]);
        if (type & 1)
            *qlen += len;
        if (type & 2)
            *rlen += len;
    }
}

int parseBam(uint8_t *data, bam1_t *b) {
    bam1_core_t *c = &b->core;
    int32_t block_len, ret, i;
    uint32_t new_l_data;
    uint8_t tmp[32], *x;

    b->l_data = 0;

    block_len = *(uint32_t *)data;
    spdlog::info("record_len: {}", block_len);
    data += 4;
    x = data;

    c->tid = le_to_u32(x);
    c->pos = le_to_i32(x + 4);
    uint32_t x2 = le_to_u32(x + 8);
    c->bin = x2 >> 16;
    c->qual = x2 >> 8 & 0xff;
    c->l_qname = x2 & 0xff;
    c->l_extranul = (c->l_qname % 4 != 0) ? (4 - c->l_qname % 4) : 0;
    uint32_t x3 = le_to_u32(x + 12);
    c->flag = x3 >> 16;
    c->n_cigar = x3 & 0xffff;
    c->l_qseq = le_to_u32(x + 16);
    c->mtid = le_to_u32(x + 20);
    c->mpos = le_to_i32(x + 24);
    c->isize = le_to_i32(x + 28);

    new_l_data = block_len - 32 + c->l_extranul;
    if (new_l_data > INT_MAX || c->l_qseq < 0 || c->l_qname < 1)
        return -4;
    if (((uint64_t)c->n_cigar << 2) + c->l_qname + c->l_extranul + (((uint64_t)c->l_qseq + 1) >> 1) + c->l_qseq >
        (uint64_t)new_l_data)
        return -4;
    if (realloc_bam_data(b, new_l_data) < 0)
        return -4;
    b->l_data = new_l_data;

    b->data = data;
    data += c->l_qname;

    if (b->data[c->l_qname - 1] != '\0') {  // try to fix missing nul termination
        if (fixup_missing_qname_nul(b) < 0)
            return -4;
    }
    for (i = 0; i < c->l_extranul; ++i) b->data[c->l_qname + i] = '\0';
    c->l_qname += c->l_extranul;
    data += b->l_data - c->l_qname;


    // TODO: consider making this conditional
    if (c->n_cigar > 0) {  // recompute "bin" and check CIGAR-qlen consistency
        hts_pos_t rlen, qlen;
        bam_cigar2rqlens(c->n_cigar, bam_get_cigar(b), &rlen, &qlen);
        if ((b->core.flag & BAM_FUNMAP) || rlen == 0)
            rlen = 1;
        b->core.bin = hts_reg2bin(b->core.pos, b->core.pos + rlen, 14, 5);
        // Sanity check for broken CIGAR alignments
        if (c->l_qseq > 0 && !(c->flag & BAM_FUNMAP) && qlen != c->l_qseq) {
            hts_log_error("CIGAR and query sequence lengths differ for %s", bam_get_qname(b));
            return -4;
        }
    }

    return 4 + block_len;
}

// multi-thread uncompress bam blocks
static void mtUncompressBlock(void *data, long idx, int tid) {
    auto &p = *(SortData*) data;
    auto &bamArr = p.bamArr[tid];
    auto &udata = p.uData[tid];

    // if (udata.maxLen < nsgv::gSortArg.MAX_MEM / nsgv::gSortArg.NUM_THREADS) {
    //     udata.allocMem(nsgv::gSortArg.MAX_MEM / nsgv::gSortArg.NUM_THREADS);
    // }
    // return;

    // spdlog::info("start: {} end: {} size: {} idx: {}, tid: {}", start, end, p.blockAddrArr->size(), idx, tid);

    long start = idx;
    int step = nsgv::gSortArg.NUM_THREADS;
    // long stop = p.blockAddrArr->size();
    long stop = idx + 1;

    // int chunkSize = (p.blockAddrArr->size() + nsgv::gSortArg.NUM_THREADS - 1) / nsgv::gSortArg.NUM_THREADS;
    // long start = idx * chunkSize;
    // int step = 1;
    // long stop = std::min((idx + 1) * chunkSize, (long)p.blockAddrArr->size());

    uint8_t buf[65535];
    // if (tid == 0 ) spdlog::info("round:{} tid:{} start:{} stop:{}", p.round, tid, start, stop);

    for (long i = start; i < stop; i += step) {
        uint8_t *block = (*p.blockAddrArr)[i];
        size_t dlen = 0x10000;  // 65535
        if (udata.maxLen - udata.curLen < dlen) {
            // spdlog::info("re alloc buf, maxLen: {} M", udata.maxLen / 1024 / 1024);
            udata.maxLen += INCREASE_SIZE;
            udata.data = (uint8_t *)realloc(udata.data, udata.maxLen);
        }
        int block_length = unpackInt16(&block[16]) + 1;
        uint32_t crc = le_to_u32(block + block_length - 8);
        int ret = bgzfUncompress(buf, &dlen, (Bytef *)block + 18, block_length - 18, crc);
        // spdlog::info("BamSortData size: {}, bam size: {}", sizeof(BamSortData), *(uint32_t *)(udata.data +
        // udata.curLen));
        // spdlog::info("i: {} slen: {} dlen: {}", i, block_length, dlen);
        // 不用对bam进行完全解析，反正只需要排序就行
        // int pos = 0;
        // uint8_t *ud = udata.data + udata.curLen;
        // bam1_t *bp = bam_init1();
        //while (pos < dlen) {
        //    bamArr.push_back(BamSortData());
        //    BamSortData &b = bamArr.back();
        //    b.bamLen = *(uint32_t *)(buf + pos);
        //    // // spdlog::info("bamLen: {}", b.bamLen);
        //    pos += 4 + b.bamLen;
        //    // if (b.bamLen > 1000)
        //    // spdlog::info("tid: {} record len: {}, dlen: {}", tid, b.bamLen, dlen);
        //    // exit(0);
        //    // int bamLen = parseBam(buf, bp);
        //    // exit(0);
        //    // spdlog::info("bamLen: {}", b.bamLen);
        //    //  pos += bamLen;
        //}
        // exit(0);
        // usleep(1);
        PROF_T_BEG(mem_copy);
        memcpy(udata.data + udata.curLen, buf, dlen);
        PROF_T_END(tid, mem_copy);
        udata.curLen += dlen * 1;
    }
}

// multi-thread sort bam blocks
static void mtSortBlocks(void *data, long idx, int tid) {
    auto &p = *(SortData *)data;
    auto &bamArr = p.bamArr[tid];
    auto &uData = p.uData[tid];
}

static void mergeSortedBlocks(const SortData &sortData) {

}

static void *nonBlockingUncompress(void *data) {
    PROF_G_BEG(uncompress);
    if (((SortData *)data)->round % 10 == 0)
        spdlog::info("round-{} block size: {}", ((SortData *)data)->round++, ((SortData *)data)->blockAddrArr->size());
    else
        ((SortData *)data)->round++;

    // spdlog::info("block arr size: {}", ((SortData *)data)->blockAddrArr->size());
    kt_for(nsgv::gSortArg.NUM_THREADS, mtUncompressBlock, data, ((SortData *)data)->blockAddrArr->size());
    // kt_for(nsgv::gSortArg.NUM_THREADS, mtUncompressBlock, data, nsgv::gSortArg.NUM_THREADS);
    PROF_G_END(uncompress);
    return nullptr;
}

static size_t uncommpressedDataSize(const SortData &sortData) {
    size_t curSize = 0;
    for (int i = 0; i < sortData.uData.size(); ++i) {
        curSize += sortData.uData[i].curLen;
    }
    return curSize;
}

template <typename T>
static void switchPointer(T *p1, T *p2) {
    const T tmp = *p1;
    *p1 = *p2;
    *p2 = tmp;
}

static size_t threadMemSize(SortData &sortData) {
    size_t memSize = 0;
    for (int i = 0; i < sortData.uData.size(); ++i) {
        memSize += sortData.uData[i].curLen;
    }
    return memSize;
}

static size_t threadMaxMemSize(SortData &sortData) {
    size_t memSize = 0;
    for (int i = 0; i < sortData.uData.size(); ++i) {
        memSize += sortData.uData[i].maxLen;
    }
    return memSize;
}

void *threadRead(void *data) {
    spdlog::info("test file start");
    FILE *fp = (FILE*) data;
    size_t bufSize = 64L * 1024 * 1024;
    uint8_t *testBuf = (uint8_t *)malloc(bufSize);
    size_t readState = 0;
    size_t fileSize = 0;
    while ((readState = fread(testBuf, 1, bufSize, fp)) > 0) {
        fileSize += readState;
    }
    spdlog::info("test file size: {}", fileSize);
    return NULL;
}

int doSort() {

#if 0
    /* 打开输入bam文件 */
    samFile *inBamFp = sam_open_format(nsgv::gSortArg.INPUT_FILE.c_str(), "r", nullptr);
    if (!inBamFp) {
        spdlog::error("[{}] load sam/bam file failed.\n", __func__);
        return -1;
    }
    hts_set_opt(inBamFp, HTS_OPT_BLOCK_SIZE, BAM_BLOCK_SIZE);
    sam_hdr_t *inBamHdr = sam_hdr_read(inBamFp);  // 读取header
    
    htsThreadPool htsPoolRead = {NULL, 0};        // 多线程读取，创建线程池
    htsPoolRead.pool = hts_tpool_init(nsgv::gSortArg.NUM_THREADS);
    if (!htsPoolRead.pool) {
        spdlog::error("[{}] failed to set up thread pool", __LINE__);
        sam_close(inBamFp);
        return -1;
    }
    hts_set_opt(inBamFp, HTS_OPT_THREAD_POOL, &htsPoolRead);

    // 测试读取bam的速度

    int max_bam_len = 0;
    bam1_t *bamp = bam_init1();
    while (sam_read1(inBamFp, inBamHdr, bamp) >= 0) {
        if (max_bam_len < bamp->l_data) max_bam_len = bamp->l_data;
        if (bamp->l_data > 1000) {
            spdlog::info("large record len: {}", bamp->l_data);
        }
    }
    sam_close(inBamFp);
    spdlog::info("max record len: {}", max_bam_len);
    return 0;

#endif
    // 初始化算法用到的数据结构
    SortData sortData;
    sortData.bamArr.resize(nsgv::gSortArg.NUM_THREADS);
    sortData.uData.resize(nsgv::gSortArg.NUM_THREADS);
    for (int i = 0; i < nsgv::gSortArg.NUM_THREADS; ++i) {
        sortData.uData[i].allocMem(nsgv::gSortArg.MAX_MEM / nsgv::gSortArg.NUM_THREADS);
    }
    nsgv::gIsBigEndian = ed_is_big();

    // 打开文件
    FILE *fpr = fopen(nsgv::gSortArg.INPUT_FILE.c_str(), "rb");

    // threadRead(fpr); exit(0);

    parseSamHeader(fpr, nsgv::gInHdr);

    //    FILE *fpw = fopen(nsgv::gSortArg.OUTPUT_FILE.c_str(), "rb");

    const size_t READ_BUFSIZE = 4L * 1024 * 1024 * nsgv::gSortArg.NUM_THREADS;
    const size_t MARGIN_SIZE = READ_BUFSIZE * 4;
    // const size_t READ_BUFSIZE = 8L * 1024 * 1024;
    const size_t SINGLE_BLOCK_SIZE = 65535;
    uint8_t *fbuf[5];
    fbuf[0] = (uint8_t *)malloc(READ_BUFSIZE);
    fbuf[1] = (uint8_t *)malloc(READ_BUFSIZE);
    fbuf[2] = (uint8_t *)malloc(SINGLE_BLOCK_SIZE);
    fbuf[3] = (uint8_t *)malloc(SINGLE_BLOCK_SIZE);
    fbuf[4] = (uint8_t *)malloc(READ_BUFSIZE);

    //setvbuf(fpr, (char*)fbuf[4], _IOFBF, 0);
    //setvbuf(fpr, (char *)fbuf[4], _IONBF, READ_BUFSIZE);

    uint8_t *curBuf = fbuf[0];
    uint8_t *preBuf = fbuf[1];
    uint8_t *preBlock = fbuf[2];
    uint8_t *curBlock = fbuf[3];

    size_t fsize = 0;
    size_t block_num = 0;
    int round = 0;
    size_t totalBlocks = 0;

    /*
    有几种情况
    第一种：上一个buf里，只剩17字节，那么不能解析下一个block的长度
    第二种：上一个buf里，完全不剩字节，即上一个buf读入完成时正好卡在下一个block开头
    第三种：上一个buf里，剩余18字节，只能解析下一个block的长度
    第四种：上一个buf里，剩余超过18字节，但是没有放下所有block的数据

    简而言之：1. 够解析下一个block的header，不够解析完整的block
    2. 不够解析下一个block的header。
    */
    size_t readState = 0;
    size_t curReadPos = 0;
    int blockLen = 0;
    vector<uint8_t *> startAddrArr[2];
    vector<uint8_t *> *curStartAddrArr = &startAddrArr[0];
    vector<uint8_t *> *preStartAddrArr = &startAddrArr[1];
    size_t curBLockPos = 0;
    size_t lastBufRemain = 0;
    uint8_t *switchBuf;
    vector<uint8_t *> *switchAddrArr;
    uint8_t *switchBlock;
    pthread_t uncompressTid = 0;
    pthread_t sortMergeTid = 0;

// 读取一个压缩的block
#define READ_BLOCKS                                            \
    curBLockPos = curReadPos;                                  \
    while (curReadPos + BLOCK_HEADER_LENGTH < readState) {     \
        blockLen = unpackInt16(&curBuf[curBLockPos + 16]) + 1; \
        curReadPos += blockLen;                                \
        if (curReadPos < readState) {                          \
            /* 这个block完整的包含在curBuf里 */                  \
            curStartAddrArr->push_back(&curBuf[curBLockPos]);  \
            curBLockPos = curReadPos;                          \
        }                                                      \
    }                                                          \
    lastBufRemain = readState - curBLockPos;                   \
    memcpy(curBlock, &curBuf[curBLockPos], lastBufRemain);  // 将不完整的block拷贝到curBlock

// 交换buffer指针
#define SWITCH_POINTER                 \
    switchPointer(&curBuf, &preBuf); \
    switchPointer(&curStartAddrArr, &preStartAddrArr); \
    switchPointer(&curBlock, &preBlock);

    PROF_G_BEG(mid_all);

    PROF_G_BEG(read);
    readState = fread(fbuf[4], 1, READ_BUFSIZE, fpr);
    // readState = fread(curBuf, 1, READ_BUFSIZE, fpr);
    PROF_G_BEG(mem_copy);
    memcpy(curBuf, fbuf[4], READ_BUFSIZE);
    PROF_G_END(mem_copy);
    PROF_G_END(read);

    while (readState > 0) {
        //spdlog::info("readState-0: {}", readState);
        //    while (readState > 0) {
        PROF_G_BEG(parse_block);
        curStartAddrArr->clear();
        fsize += readState;
        curReadPos = 0;
//        if (lastBufRemain < BLOCK_HEADER_LENGTH) { // 上次遗留的数据不够解析block 长度
//            memcpy(&preBlock[lastBufRemain], curBuf, BLOCK_HEADER_LENGTH);
//            curReadPos += BLOCK_HEADER_LENGTH;
//            lastBufRemain += BLOCK_HEADER_LENGTH;
//        }
//        blockLen = unpackInt16(&preBlock[16]) + 1;
//        const int lenToRead = blockLen - lastBufRemain;
//        memcpy(&preBlock[lastBufRemain], &curBuf[curReadPos], lenToRead);
//        curReadPos += lenToRead;
//        curStartAddrArr->push_back(&preBlock[0]);

        // curStartAddrArr->push_back(&curBuf[0]);
        // spdlog::info("first blocklen: {}, remain: {}, toread: {}", blockLen, lastBufRemain, lenToRead);
        // uint8_t buf[65535];
        // int block_length = unpackInt16(&preBlock[16]) + 1;
        // uint32_t crc = le_to_u32(preBlock + block_length - 8);
        // size_t dlen = 65535;
        // int ret = bgzf_uncompress(buf, &dlen, (Bytef *)preBlock + 18, block_length - 18, crc); // 开头是header，得跳过
        // int bamLen = *(uint32_t *)(buf);
        // spdlog::info("ret: {}, bamLen: {}, dlen: {}, block_length: {}", ret, bamLen, dlen, block_length);
        // exit(0);

        READ_BLOCKS;
        // spdlog::info("read block len: {}", blockLen);

        // spdlog::info("cur block size: {}", curStartAddrArr->size());

        totalBlocks += curStartAddrArr->size();
        PROF_G_END(parse_block);

#if 1
        // 并行解压
        PROF_G_BEG(sort);
        pthread_join(uncompressTid, NULL);
        {
            // 如果sortData的解压buf满了，那么进行排序并输出
            if (true) {
            //if (threadMemSize(sortData) + MARGIN_SIZE >= nsgv::gSortArg.MAX_MEM) {
                // spdlog::info("clear buf data, buf size: {} G, {} G", threadMemSize(sortData) / 1024.0 / 1024 / 1024, threadMaxMemSize(sortData) / 1024.0 / 1024 / 1024);
                for (int i = 0; i < sortData.uData.size(); ++i) {
                    sortData.uData[i].curLen = 0; // 清空数据
                    sortData.bamArr[i].clear();
                }
            }
        }
        PROF_G_END(sort);

        sortData.blockAddrArr = curStartAddrArr;
        pthread_create(&uncompressTid, NULL, nonBlockingUncompress, &sortData);
        //nonBlockingUncompress(&sortData);
#endif
        // kt_for(nsgv::gSortArg.NUM_THREADS, mtUncompressBlock, &sortData, curStartAddrArr->size() / BATCH_SIZE);
        // if (uncommpressedDataSize(sortData) >= nsgv::gSortArg.MAX_MEM) {
        //     kt_for(nsgv::gSortArg.NUM_THREADS, mtSortBlocks, &sortData, nsgv::gSortArg.NUM_THREADS);
        //     mergeSortedBlocks(sortData);
        // }

        SWITCH_POINTER;

        PROF_G_BEG(read);
        readState = fread(curBuf, 1, READ_BUFSIZE, fpr);
        //spdlog::info("readState: {}", readState);
        PROF_G_BEG(mem_copy);
        // memcpy(curBuf, fbuf[4], readState); readState = READ_BUFSIZE;
        memcpy(curBuf, fbuf[4], readState);
        PROF_G_END(mem_copy);
        PROF_G_END(read);
        // if (fsize >= 6245369164) break;
    }
    pthread_join(uncompressTid, NULL);
    PROF_G_END(mid_all);


    spdlog::info("total blocks: {}", totalBlocks);
    fprintf(stderr, "file size: %ld\n\n", fsize);
err:
    free(fbuf[0]);
    free(fbuf[1]);
    free(fbuf[2]);
    free(fbuf[3]);
    fclose(fpr);
    //    fclose(fpw);

    return 0;
}