/*
     Description: 第一阶段的排序

     Copyright : All right reserved by ICT

     Author : Zhang Zhonghai
     Date : 2026/06/02
*/

#include "phase_1_sort.h"

#include <klib/kthread.h>
#include <spdlog/spdlog.h>

#include <algorithm>
#include <string>

#include "common_data.h"
#include "const_val.h"
#include "phase_1.h"
#include "phase_1_compress.h"
#include "phase_1_write.h"
#include "sort.h"
#include "util/profiling.h"

uint64_t sort_bam_id = 0;
uint64_t block_nums = 0;

/* bam 排序堆 */
struct Phase1BamArrIdIdx {
    size_t idx = 0;  // 下一个待读入数据的idx
    size_t endIdx = 0;
    const OneBam* bam = nullptr;
};

struct Phase1BamPosGreaterThan {
    bool operator()(const Phase1BamArrIdIdx& a, const Phase1BamArrIdIdx& b) const { return a.bam->pos > b.bam->pos; }
};

static uint8_t* gUncompressDataBuf = nullptr;  // 全局的解压数据buf，排序的时候用来比较bam的坐标或者名字信息
struct Phase1BamNameGreaterThan {
    bool operator()(const Phase1BamArrIdIdx& a, const Phase1BamArrIdIdx& b) const {
        int cmp = strncmp((char*)(gUncompressDataBuf + a.bam->offset + OneBam::QnameOffset), (char*)(gUncompressDataBuf + b.bam->offset + OneBam::QnameOffset),
                          std::min(a.bam->qnameLen, b.bam->qnameLen));
        if (cmp == 0)
            return a.bam->qnameLen > b.bam->qnameLen;
        return cmp > 0;
    }
};

/* 用来排序 bam*/
template <class GreaterThan>
struct Phase1BamHeap {
    Phase1PipelineArg* p;
    priority_queue<Phase1BamArrIdIdx, vector<Phase1BamArrIdIdx>, GreaterThan> minHeap;
    size_t popNum = 0;

    int Init(Phase1PipelineArg* _p) {
        p = _p;
        if (p == nullptr) {
            return -1;
        }
        // for (int i = 0; i < 4; ++i) {
        for (int i = 0; i < p->numThread; ++i) {
            size_t startIdx = START_IDX(i, p->numThread, p->allBams.Size());
            size_t stopIdx = STOP_IDX(i, p->numThread, p->allBams.Size());
            if (startIdx < stopIdx) {
                minHeap.push({startIdx, stopIdx, &p->allBams.arr[startIdx]});
            }
        }
        return 0;
    }

    const OneBam* Pop() {
        const OneBam* ret = nullptr;
        if (!minHeap.empty()) {
            auto minVal = minHeap.top();
            minHeap.pop();
            ++popNum;
            ret = minVal.bam;
            if (minVal.idx + 1 < minVal.endIdx) {
                minVal.idx++;
                minVal.bam = &p->allBams.arr[minVal.idx];
                minHeap.push(minVal);
            }
        }
        return ret;
    }

    const OneBam* Top() {
        const OneBam* ret = nullptr;
        if (!minHeap.empty()) {
            ret = minHeap.top().bam;
        }
        return ret;
    }
};

// 线程内排序
void mtInThreadSort(void* data, long idx, int tid) {
    Phase1PipelineArg& p = *(Phase1PipelineArg*)data;
    tid = idx;

    // 先按照坐标排序
    int startIdx = START_IDX(tid, p.numThread, p.allBams.Size());
    int stopIdx = STOP_IDX(tid, p.numThread, p.allBams.Size());

    if (nsgv::gSortArg.SORT_COORIDINATE) {
        std::sort(p.allBams.arr.begin() + startIdx, p.allBams.arr.begin() + stopIdx, [](const OneBam& b1, const OneBam& b2) {
            if (b1.tid == -1)
                return false;
            return b1.tid < b2.tid || (b1.tid == b2.tid && b1.pos < b2.pos) || (b1.tid == b2.tid && b1.pos == b2.pos && b1.qnameLen < b2.qnameLen);
        });
    } else if (nsgv::gSortArg.QUERY_NAME_TYPE == nsmd::QueryNameType::PICARD) {
        std::sort(p.allBams.arr.begin() + startIdx, p.allBams.arr.begin() + stopIdx, [&](const OneBam& b1, const OneBam& b2) {
            int cmp = strncmp((char*)(p.uncompressData.dataBuf + b1.offset + OneBam::QnameOffset),
                    (char*)(p.uncompressData.dataBuf + b2.offset + OneBam::QnameOffset), std::min(b1.qnameLen, b2.qnameLen));
            if (cmp == 0)
                return b1.qnameLen < b2.qnameLen;
            return cmp < 0;
        });
    }
}

void singleThreadSort(void* data) {
    Phase1PipelineArg& p = *(Phase1PipelineArg*)data;
    // 先按照坐标排序
    int startIdx = 0;
    int stopIdx = p.allBams.Size();

    if (nsgv::gSortArg.SORT_COORIDINATE) {
        std::sort(p.allBams.arr.begin() + startIdx, p.allBams.arr.begin() + stopIdx, [](const OneBam& b1, const OneBam& b2) {
            if (b1.tid == -1)
                return false;
            return b1.tid < b2.tid || (b1.tid == b2.tid && b1.pos < b2.pos) || (b1.tid == b2.tid && b1.pos == b2.pos && b1.qnameLen < b2.qnameLen);
        });
    } else if (nsgv::gSortArg.QUERY_NAME_TYPE == nsmd::QueryNameType::PICARD) {
        std::sort(p.allBams.arr.begin() + startIdx, p.allBams.arr.begin() + stopIdx, [&](const OneBam& b1, const OneBam& b2) {
            int cmp = strncmp((char*)(p.uncompressData.dataBuf + b1.offset + OneBam::QnameOffset),
                              (char*)(p.uncompressData.dataBuf + b2.offset + OneBam::QnameOffset), std::min(b1.qnameLen, b2.qnameLen));
            if (cmp == 0)
                return b1.qnameLen < b2.qnameLen;
            return cmp < 0;
        });
    }
}

void phase1Sort(void* data) {
    Phase1PipelineArg& p = *(Phase1PipelineArg*)data;
    PROF_G_BEG(sort);
    kt_for(p.numThread, mtInThreadSort, &p, p.numThread);
    // kt_for(1, mtInThreadSort, &p, p.numThread);
    PROF_G_END(sort);

    // spdlog::info("b1:{}, b2:{}, b3:{}, b4:{}, b5:{}", p.allBams.arr[0].tid, p.allBams.arr[1].tid, p.allBams.arr[2].tid, p.allBams.arr[3].tid,
    //              p.allBams.arr[4].tid);
    // spdlog::info("b1:{}, b2:{}, b3:{}, b4:{}, b5:{}", p.allBams.arr[0].pos, p.allBams.arr[1].pos, p.allBams.arr[2].pos, p.allBams.arr[3].pos,
    //              p.allBams.arr[4].pos);
    // spdlog::info("b1:{}, b2:{}, b3:{}, b4:{}, b5:{}", p.allBams.arr[0].offset, p.allBams.arr[1].offset, p.allBams.arr[2].offset, p.allBams.arr[3].offset,
    //              p.allBams.arr[4].offset);
}

template <class BamGreaterThan>
void doMergeLastData(Phase1PipelineArg& p, Phase1BamHeap<BamGreaterThan>& heap) {
    const OneBam* bam = nullptr;

    vector<OneBam> arr(p.allBams.Size());

    //p.lastBams.resize(p.allBams.Size());
    size_t i = 0;
    //uint64_t pos;
    while ((bam = heap.Pop()) != nullptr) {
        //pos += bam->tid;
        arr[i++] = *bam;
        //p.lastBams[i++] = bam;
    }
    //spdlog::info("bam num: {}, pos: {}", p.lastBams.size(), pos);
    spdlog::info("bam size: {}-{}, buf size: {}-{}", p.allBams.Capacity(), arr.size(), p.uncompressData.dataBufSize, p.uncompressData.usedBufSize);
    p.allBams.arr = arr; // 替换
}

void mergeLastData(void* data) {
    Phase1PipelineArg& p = *(Phase1PipelineArg*)data;

    gUncompressDataBuf = p.uncompressData.dataBuf;  // 主要用来比较bam的名字信息
    Phase1BamHeap<Phase1BamPosGreaterThan> posHeap;
    Phase1BamHeap<Phase1BamNameGreaterThan> nameHeap;
    if (nsgv::gSortArg.SORT_COORIDINATE) {
        posHeap.Init(&p);
        doMergeLastData(p, posHeap);
    } else {
        nameHeap.Init(&p);
        doMergeLastData(p, nameHeap);
    }
}

template <class BamGreaterThan>
bool doMergeSort(Phase1PipelineArg& p, Phase1BamHeap<BamGreaterThan>& heap) {
    auto &mergeCompressData = p.mergeCompressData[p.mergeOrder % p.MERGE_BUF_NUM];

    bool finish = false;
    const OneBam* bam = nullptr;

    int mergedBlockNum = 0;
    size_t bamBytes = 0;
    int singleBlockBytes = 0xff00;
    mergeCompressData.blockDataArr[mergedBlockNum].Clear();
    while ((bam = heap.Top()) != nullptr) {
        if (bamBytes + bam->wholeBamLen > singleBlockBytes) {
            mergedBlockNum++;
            if (mergedBlockNum >= p.mergeBlocksThreshold) {
                mergeCompressData.curIdx = mergedBlockNum; // 这里和下面finish时候不同，要注意
                break;
            }
            mergeCompressData.blockDataArr[mergedBlockNum].Clear(); // 清理，为添加bam数据做准备
            bamBytes = 0;
        }
#if 0
        mergeCompressData.blockDataArr[mergedBlockNum].blockBuf.MemCopy(p.uncompressData.dataBuf + bam->offset, bam->wholeBamLen);
#else
        sort_bam_id += 1;
        mergeCompressData.blockDataArr[mergedBlockNum].bamPtrArr.Add(bam);
#endif
        bamBytes += bam->wholeBamLen;  // for test
        heap.Pop();
    }
    if (bam == nullptr) {
        // 都处理完
        finish = true;
        mergeCompressData.curIdx = mergedBlockNum + 1;  // 记录有多少个block待压缩，主要是最后一些数据凑不满threshold
    }
   
    block_nums += mergedBlockNum;
    // spdlog::info("mergedBlockNum: {}, bamBytes: {}, blocks: {} bam id: {}", mergedBlockNum, bamBytes, block_nums, sort_bam_id);

    return finish;
}

void* phase1MergeSort(void* data) {
    Phase1PipelineArg& p = *(Phase1PipelineArg*)data;
    gUncompressDataBuf = p.uncompressData.dataBuf;  // 主要用来比较bam的名字信息
    Phase1BamHeap<Phase1BamPosGreaterThan> posHeap;
    Phase1BamHeap<Phase1BamNameGreaterThan> nameHeap;

    if (nsgv::gSortArg.SORT_COORIDINATE) {
        posHeap.Init(&p);
    } else {
        nameHeap.Init(&p);
    }
    // merge
    /* do the work */
    while (true) {
        // previous dependency
        yarn::DEPENDENCY_NOT_TO_BE(p.mergeSig, p.MERGE_BUF_NUM);
        PROF_G_BEG(merge);
        // bool finish = doMergeSort(p, heap);
        bool finish = false;
        if (nsgv::gSortArg.SORT_COORIDINATE) {
            finish = doMergeSort(p, posHeap);
        } else {
            finish = doMergeSort(p, nameHeap);
        }
        PROF_G_END(merge);

        if (finish) {
            yarn::SIGNAL_FINISH_ADD_ORDER(p.mergeSig, p.mergeFinish, p.mergeOrder);
            // yarn::SIGNAL_FINISH(p.mergeSig, p.mergeFinish);
            break;
        }

        // update self status
        yarn::UPDATE_SIG_ORDER(p.mergeSig, p.mergeOrder);
    }

    spdlog::info("End merge sort order: {}", p.mergeOrder);
    return nullptr;
}

void ReInitMergePhase(Phase1PipelineArg& p) {
    p.mergeOrder = 0;
    p.compressOrder = 0;
    p.writeOrder = 0;
    p.mergeFinish = 0;
    p.compressFinish = 0;
    yarn::INIT_SIG(p.mergeSig);
    yarn::INIT_SIG(p.compressSig);
}

void phase1MergeCompress(void* data) {
    Phase1PipelineArg& p = *(Phase1PipelineArg*)data;
    ReInitMergePhase(p);
    // 创建中间文件
    p.midFilePtr = fopen((p.midFileNamePrefix + std::to_string(p.midFileOrder)).c_str(), "wb");
    if (p.midFilePtr == nullptr) {
        spdlog::error("failed to create mid file: {}", p.midFileNamePrefix + std::to_string(p.midFileOrder));
        exit(1);
    }

    // 2-stage pipeline, merge-compress-write
    pthread_t tidArr[3];  // 2-stage pipeline
    pthread_create(&tidArr[0], NULL, phase1MergeSort, &p);
    pthread_create(&tidArr[1], NULL, phase1Compress, &p);
    pthread_create(&tidArr[2], NULL, phase1Write, &p);

    for (int i = 0; i < 3; ++i) pthread_join(tidArr[i], NULL);
    p.midFileOrder += 1; // 写完一次中间文件 +1
    fclose(p.midFilePtr);
}