/*
     Description: 归并中间文件

     Copyright : All right reserved by ICT

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

#include "phase_2_merge.h"

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

#include "common_data.h"
#include "phase_2.h"
#include "util/profiling.h"

static size_t all_bams = 0;

/* bam 排序堆 */
struct Phase2BamArrIdIdx {
    int idx = 0; // 如果是第一阶段转移过来的buffer，用这个，此时file为nullptr
    uint8_t* addr = nullptr; // bam的原始地址，就是减去offset
    Phase2File* file = nullptr;
    const OneBam* bam = nullptr;
};

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

struct Phase2BamNameGreaterThan {
    bool operator()(const Phase2BamArrIdIdx& a, const Phase2BamArrIdIdx& b) const {
        int cmp = strncmp((char*)(a.addr + a.bam->offset + OneBam::QnameOffset),
                          (char*)(b.addr + 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 Phase2BamHeap {
    Phase2PipelineArg* p;
    priority_queue<Phase2BamArrIdIdx, vector<Phase2BamArrIdIdx>, GreaterThan> minHeap;
    size_t popNum = 0;
    size_t curIdx = 0;

    int Init(Phase2PipelineArg* _p) {
        p = _p;
        if (p == nullptr) {
            return -1;
        }
        // for (int i = 0; i < 4; ++i) {
        for (int i = 0; i < p->midFiles.size(); ++i) {
            auto& f = p->midFiles[i];
            f.mergeData.Front()->addr = f.mergeData.data.data;
            if (f.mergeData.Front()->addr == nullptr) {
                spdlog::error("null addr");
            }
            minHeap.push({0, f.mergeData.data.data, &f, f.mergeData.Front()});
            f.mergeData.Pop();
        }
        if (p->allBams.Size() > curIdx) {
            p->allBams.arr[curIdx].addr = p->uncompressData.dataBuf;
            if (p->allBams.arr[curIdx].addr == nullptr) {
                spdlog::error("null addr");
            }
            minHeap.push({0, p->uncompressData.dataBuf, nullptr, &p->allBams.arr[curIdx]});
        }
        return 0;
    }

    void ReInitFile(Phase2File *file) {
        auto &f = *file;
        if (f.mergeData.Front() != nullptr) {
            f.mergeData.Front()->addr = f.mergeData.data.data;
            if (f.mergeData.Front()->addr == nullptr) {
                spdlog::error("null addr");
            }
            minHeap.push({0, f.mergeData.data.data, &f, f.mergeData.Front()});
            f.mergeData.Pop();
        }
    }

    void Revert() {
        while (!minHeap.empty()) {
            auto minVal = minHeap.top();
            if (minVal.file != nullptr) {
                auto& mergeData = minVal.file->mergeData;
                mergeData.Revert();
            } else {
                --curIdx;
            }
            minHeap.pop();
        }
    }

    const OneBam* Pop() {
        const OneBam* ret = nullptr;
        if (!minHeap.empty()) {
            auto minVal = minHeap.top();
            minHeap.pop();
            ++popNum;
            ret = minVal.bam;
            if (minVal.file != nullptr) {
                auto& mergeData = minVal.file->mergeData;
                if (mergeData.Size() > 0) {
                    mergeData.Front()->addr = minVal.addr;
                    if (mergeData.Front()->addr == nullptr) {
                        spdlog::error("null addr");
                    }
                    minHeap.push({0, minVal.addr, minVal.file, mergeData.Front()});
                    mergeData.Pop();
                }
            } else {
                if (minVal.idx + 1 < p->allBams.Size()) {
                    curIdx = minVal.idx + 1;
                    p->allBams.arr[minVal.idx + 1].addr = minVal.addr;
                    minHeap.push({minVal.idx + 1, minVal.addr, nullptr, &p->allBams.arr[minVal.idx + 1]});
                }
            }
        }
        return ret;
    }

    const OneBam* Pop(Phase2File ** emptyFile) {
        const OneBam* ret = nullptr;
        if (!minHeap.empty()) {
            auto minVal = minHeap.top();
            minHeap.pop();
            ++popNum;
            ret = minVal.bam;
            if (minVal.file != nullptr) {
                auto& mergeData = minVal.file->mergeData;
                if (mergeData.Size() > 0) {
                    mergeData.Front()->addr = minVal.addr;
                    if (mergeData.Front()->addr == nullptr) {
                        spdlog::error("null addr");
                    }
                    minHeap.push({0, minVal.addr, minVal.file, mergeData.Front()});
                    // if(minVal.file->fid == 4 && minVal.file->copyOrder == 8) {
                    // if(minVal.file->fid == 4 && minVal.file->copyOrder == 6) {
                    //     spdlog::info("debug");
                    // }
                    // if (minVal.file->fid == 4 && mergeData.Size() == 1) {
                    //     spdlog::info("debug");
                    // }
                    mergeData.Pop();
                } else {
                    *emptyFile = minVal.file;
                }
            } else {
                if (minVal.idx + 1 < p->allBams.Size()) {
                    curIdx = minVal.idx + 1;
                    p->allBams.arr[minVal.idx + 1].addr = minVal.addr;
                    if (p->allBams.arr[minVal.idx + 1].addr == nullptr) {
                        spdlog::error("null addr");
                    }
                    minHeap.push({minVal.idx + 1, minVal.addr, nullptr, &p->allBams.arr[minVal.idx + 1]});
                }
            }
        }
        return ret;
    }

    const OneBam* Top() {
        const OneBam* ret = nullptr;
        if (!minHeap.empty()) {
            ret = minHeap.top().bam;
            auto &top = minHeap.top();
            int idx = 0;              // 如果是第一阶段转移过来的buffer，用这个，此时file为nullptr
            uint8_t* addr = nullptr;  // bam的原始地址，就是减去offset
            Phase2File* file = nullptr;
            const OneBam* bam = nullptr;

            // spdlog::info("idx: {}, addr: {}, file: {}, offset: {}, qname: {}", top.idx, (void*)top.addr, top.file != nullptr ? top.file->fileName : "null", top.bam->offset, std::string((char*)(top.addr + top.bam->offset + OneBam::QnameOffset), top.bam->qnameLen));
        }
        return ret;
    }
};

static void mtCopyBams(void* data, long idx, int tid) {
    Phase2PipelineArg& p = *(Phase2PipelineArg*)data;
    MergeCompressData& mergeData = p.mergeData[p.mergeOrder % p.MERGE_BUF_NUM];

    auto& bams = mergeData.blockDataArr[idx].bamPtrArr;
    auto& blockData = mergeData.blockDataArr[idx].blockBuf;

    // spdlog::info("bam size: {}", bams.Size());
    for (int i = 0; i < bams.Size(); ++i) {
        const OneBam* bp = bams.arr[i];
        blockData.MemCopy(bp->addr + bp->offset, bp->wholeBamLen);
        // check bam
        // CheckBam(bp->addr + bp->offset, bp->wholeBamLen);
    }
    // 要清理bams数组
    mergeData.blockDataArr[idx].bamNum += bams.Size();

    bams.Clear();
}

template <class BamGreaterThan>
static bool doPhase2Merge(Phase2PipelineArg& p, Phase2BamHeap<BamGreaterThan>& heap, bool *mergeFull = nullptr, Phase2File **emptyFilePtr = nullptr, bool copyFinish = false) {
    bool finish = false;

    auto& mergeData = p.mergeData[p.mergeOrder % p.MERGE_BUF_NUM];
    const OneBam* bam = nullptr;

    int singleBlockBytes = 0xff00;
    Phase2File* emptyFile = nullptr;

    if (emptyFilePtr != nullptr && *emptyFilePtr != nullptr) {
        heap.ReInitFile(*emptyFilePtr);
    }
    // spdlog::error("curIdx: {}, bytes: {}", mergeData.curIdx, mergeData.curBytes);
    if (mergeData.curBytes == 0) { // 如果上一轮用完了某个文件，则此时curBytes可能大于0，而且没满，此时不能清空
        mergeData.blockDataArr[mergeData.curIdx].Clear();
    }
    // mergeData.curBytes = 0;
    while ((bam = heap.Top()) != nullptr) {
        if (bam->addr == nullptr) {
            spdlog::info("null addr");
        }
        // spdlog::info("bam: {}, pos: {}, tid: {}, qname: {}", bam->offset, bam->pos, bam->tid, std::string((char*)(bam->addr + bam->offset + OneBam::QnameOffset), bam->qnameLen));
        // CheckBam(bam->addr + bam->offset, bam->wholeBamLen);

        if (mergeData.curBytes + bam->wholeBamLen > singleBlockBytes) {
            mergeData.curIdx++;
            if (mergeData.curIdx >= p.compressBlocksThreshold) {
                if (mergeFull != nullptr) {
                    *mergeFull = true;
                }
                // mergeData.curIdx = 0; // 调试用的
                break;
            }
            // spdlog::error("bamBytes: {}", mergeData.curBytes);
            mergeData.blockDataArr[mergeData.curIdx].Clear();  // 清理，为添加bam数据做准备
            mergeData.curBytes = 0;
        }
#if 0
        mergeData.blockDataArr[mergeData.curIdx].blockBuf.MemCopy(p.uncompressData.dataBuf + bam->offset, bam->wholeBamLen);
#else
        mergeData.blockDataArr[mergeData.curIdx].bamPtrArr.Add(bam);
        // spdlog::error("bamBytes-1: {} - {}", mergeData.curIdx, mergeData.curBytes);
        write_bam_id += 1;
#endif
        mergeData.curBytes += bam->wholeBamLen;  // for test
        if (copyFinish) {
            heap.Pop();
        } else {
            heap.Pop(&emptyFile);
            if (emptyFile != nullptr) {
                // spdlog::error("emptyfile, curIdx: {}, bytes: {}", mergeData.curIdx, mergeData.blockDataArr[mergeData.curIdx].blockBuf.curLen);
                // for (int n = 0; n < mergeData.curIdx; ++n) {
                //     spdlog::info("emptyfile, block len: {}, bam len: {}", mergeData.blockDataArr[n].blockBuf.curLen, bam->wholeBamLen);
                // }
                break;
            }
        }
    }
    // 并行拷贝bam数据
    if (p.mergeOrder == 43) {
        spdlog::info("curIdx: {}, bytes: {}", mergeData.curIdx, mergeData.curBytes);
    }
    spdlog::info("curIdx: {}", mergeData.curIdx);
    if (mergeData.lastRoundIdx != p.compressBlocksThreshold) { // 回头再想想

    }

    if (mergeData.curIdx < p.compressBlocksThreshold) {
        kt_for(p.numThread, mtCopyBams, &p, mergeData.curIdx + 1);
        // kt_for(p.numThread, mtCopyBams, &p, mergeData.curIdx);
    } else {
        kt_for(p.numThread, mtCopyBams, &p, mergeData.curIdx);
    }
    // if (mergeFull != nullptr && * mergeFull) {
    // if (mergeData.curIdx == 2) {
    //     for (int n = 0; n < mergeData.curIdx; ++n) {
    //         spdlog::info("n: {}, block len: {}, bam len: {}", n, mergeData.blockDataArr[n].blockBuf.curLen, bam->wholeBamLen);
    //     }
    // }
    if (bam == nullptr) {
        // 都处理完
        finish = true;
        if (mergeData.curIdx < p.compressBlocksThreshold) {
            mergeData.curIdx++; // 最后一次，需要包含curIdx
        }
    }

    if (emptyFile != nullptr) {
        // 这里应该不需要revert，因为弹出empty file的最后一个后，没有在调用top
        // heap.Revert();
        emptyFile->readyMergeBufNum--;
    }

    if (emptyFilePtr != nullptr) {
        *emptyFilePtr = emptyFile;
    }

// 调试用
#if 0
    yarn::CONSUME_SIGNAL(p.mergeSig);
#endif

    // spdlog::info("phase2 merge sort order: {}, write bam id: {}", p.mergeOrder, write_bam_id);
     if (mergeFull != nullptr && *mergeFull) {
         for (int i = 0; i <= mergeData.curIdx; ++i) {
             all_bams += mergeData.blockDataArr[i].bamNum;
         }
         spdlog::info("phase2 merge sort order: {}, write bam id: {}, all bams: {}", p.mergeOrder, write_bam_id, all_bams);
     }
    // spdlog::info("curIdx: {}", mergeData.curIdx);

    return finish;
}

void* phase2Merge(void* data) {

    Phase2PipelineArg& p = *(Phase2PipelineArg*)data;
    Phase2BamHeap<Phase2BamPosGreaterThan> posHeap;
    Phase2BamHeap<Phase2BamNameGreaterThan> nameHeap;
    bool firstInit = true;
    Phase2File* emptyFile = nullptr;
#if 1
    // merge
    /* do the work */
    while (true) {
        // previous dependency
        yarn::DEPENDENCY_NOT_TO_BE(p.mergeSig, p.MERGE_BUF_NUM);
        yarn::DEPENDENCY_NOT_TO_BE(p.copyMerge, 0);

        // bool finish = doMergeSort(p, heap);
        bool finish = false;
        bool mergeFull = false;
        // emptyFile = nullptr;

        if (firstInit) {
            if (nsgv::gSortArg.SORT_COORIDINATE) {
                posHeap.Init(&p);
            } else {
                nameHeap.Init(&p);
            }
            firstInit = false;
        }
        if (p.copyMergeFinish) {
            while (!finish) {
                mergeFull = false;
                PROF_G_BEG(phase2_merge);
                yarn::DEPENDENCY_NOT_TO_BE(p.mergeSig, p.MERGE_BUF_NUM);
                if (nsgv::gSortArg.SORT_COORIDINATE) {
                    finish = doPhase2Merge(p, posHeap, &mergeFull, &emptyFile, true);
                } else {
                    finish = doPhase2Merge(p, nameHeap, &mergeFull, &emptyFile, true);
                }
                if (mergeFull) {
                    yarn::UPDATE_SIG_ORDER(p.mergeSig, p.mergeOrder);
                }
                PROF_G_END(phase2_merge);
            }
            yarn::UPDATE_SIG_ORDER(p.mergeSig, p.mergeOrder);
            spdlog::info("last phase2 merge sort order: {}, write bam id: {}", p.mergeOrder, write_bam_id);
            yarn::SIGNAL_FINISH(p.mergeSig, p.mergeFinish);
            break;
        }
        PROF_G_BEG(phase2_merge);
        if (nsgv::gSortArg.SORT_COORIDINATE) {
            doPhase2Merge(p, posHeap, &mergeFull, &emptyFile);
        } else {
            doPhase2Merge(p, nameHeap, &mergeFull, &emptyFile);
        }
        PROF_G_END(phase2_merge);

        // update self status
        if (emptyFile != nullptr) {// 只有某个文件的buf消耗完了，才读入
            yarn::CONSUME_SIGNAL(p.copyMerge);
        }
        if (mergeFull) {
            yarn::UPDATE_SIG_ORDER(p.mergeSig, p.mergeOrder);
        }
    }
#endif
    spdlog::info("End phase2 merge sort order: {}", p.mergeOrder);
    return nullptr;
}