/*
     Description: 中间文件读入和解析等，包含读入数据，解压，拷贝到merge缓冲区

     Copyright : All right reserved by ICT

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

#include "phase_2_read.h"

#include <klib/kthread.h>
#include <spdlog/spdlog.h>
#include <stdint.h>
#include <stdio.h>
#include <zlib.h>

#include "common_data.h"
#include "const_val.h"
#include "phase_1_read.h"
#include "phase_2.h"
#include "sam_io.h"
#include "sort.h"
#include "util/profiling.h"
#include "util/yarn.h"

//////////////////////////////////////////////// 读入中间文件
extern size_t ReadBamFile(ReadBuffer& readData, DataBuffer& halfBlock, FILE* fpr);

bool doReadMidFiles(Phase2PipelineArg& p) {
    PROF_G_BEG(phase2_read);
    bool finish = false;
    int finishNum = 0;
    for (int i = 0; i < p.midFiles.size(); ++i) {
        auto& f = p.midFiles[i];
        if (f.readyReadBufNum < Phase2File::READ_BUF_NUM && !f.finish) {  // 还有空余缓冲区
            auto& readData = f.readData[f.readOrder % Phase2File::READ_BUF_NUM];
            size_t readState = ReadBamFile(readData, f.halfBlock, f.fp);
            if (readState == 0) {
                f.finish = 1;
            } else {
                f.readOrder += 1;
                f.readyReadBufNum += 1;
                readData.uncompressBlocks.ClearResize(readData.startAddrArr.size());
            }
            //spdlog::info("read size: {}", readState);
        }
        finishNum += f.finish;
        //spdlog::info("fid: {}, readyNum: {}", i, f.readyReadBufNum);
    }
    spdlog::info("phase2 read order: {}, finishNum: {}, files: {}", p.readOrder, finishNum, p.midFiles.size());
    finish = finishNum == p.midFiles.size();
    PROF_G_END(phase2_read);
    return finish;
}

void* phase2ReadMidFile(void* data) {
    Phase2PipelineArg& p = *(Phase2PipelineArg*)data;

    /* do the work */
    while (true) {
        // previous dependency
        yarn::DEPENDENCY_NOT_TO_BE(p.readSig, Phase2File::READ_BUF_NUM);
        PROF_G_BEG(merge);
        bool finish = doReadMidFiles(p);
        PROF_G_END(merge);

        if (finish) {
            yarn::SIGNAL_FINISH(p.readSig, p.readFinish);
            break;
        }
        // update self status
        yarn::UPDATE_SIG_ORDER(p.readSig, p.readOrder);
    }

    spdlog::info("End read mid file order: {}", p.readOrder);
    return nullptr;
}

//////////////////////////////////////////////// 解析中间文件

extern size_t ParseAddAllBams(uint8_t* dataAddr, size_t startOffset, size_t endOffset, BamArr& bamArr, size_t* nextBamStartPtr = nullptr,
                       size_t* lastPosPtr = nullptr);

static void mtUncompressBlockBatch(void* data, long idx, int tid) {
    //PROF_T_BEG(mem_copy);

    Phase2PipelineArg& p = *(Phase2PipelineArg*)data;

    tid = idx;  // 静态分配任务，此时用idx代替tid
    int startIdx = START_IDX(tid, p.numThread, p.curBlockNum);
    int stopIdx = STOP_IDX(tid, p.numThread, p.curBlockNum);

    int blockNum = 0;
    int startFileIdx = -1;
    int stopFileIdx = 0;
    // 确定范围
    for (int i = 0; i < p.midFiles.size(); ++i) {
        auto& f = p.midFiles[i];
        if (f.needUncompress) {  // 还有空余缓冲区
            blockNum += f.readData[f.uncompressOrder % Phase2File::UNCOMPRESSS_BUF_NUM].startAddrArr.size();
        }
        if (startFileIdx == -1 && startIdx < blockNum) { startFileIdx = i; }
        if (blockNum >= stopIdx) {
            stopFileIdx = i; // 包含
            break;
        }
    }
    blockNum = 0;
    for (int i = 0; i < startFileIdx; ++i) {
        auto& f = p.midFiles[i];
        if (f.needUncompress) {  // 还有空余缓冲区
            auto& readData = f.readData[f.uncompressOrder % Phase2File::UNCOMPRESSS_BUF_NUM];
            blockNum += readData.startAddrArr.size();
        }
    }
    for (int i = startFileIdx; i <= stopFileIdx; ++i) {
        auto& f = p.midFiles[i];
        int start = 0;
        int stop = 0;
        if (f.needUncompress) {  // 还有空余缓冲区
            auto& readData = f.readData[f.uncompressOrder % Phase2File::UNCOMPRESSS_BUF_NUM];
            start = MAX(0, startIdx - blockNum);
            stop = MIN(readData.startAddrArr.size(), stopIdx - blockNum);

            // 解压并解析block
            for (int j = start; j < stop; ++j) {
                auto& blockData = readData.uncompressBlocks.arr[j];
                blockData.Clear();
                auto& blockBuf = blockData.block;
                blockBuf.ReAllocMem(SINGLE_BLOCK_SIZE);

                uint8_t* block = readData.startAddrArr[j];
                size_t dlen = SINGLE_BLOCK_SIZE;  // 65535
                int block_length = unpackInt16(&block[16]) + 1;
                uint32_t crc = le_to_u32(block + block_length - 8);
                //if (block_length > 30000)
                //    spdlog::error("block len: {}", block_length);
                int ret = bgzfUncompress(blockBuf.data + blockBuf.curLen, &dlen, (Bytef*)block + BLOCK_HEADER_LENGTH,
                                         block_length - BLOCK_HEADER_LENGTH, crc);
                if (ret != 0) {
                    spdlog::error("uncompress error, block id: {}, len: {}, ret: {}", idx, block_length, ret);
                    exit(0);
                }
                blockBuf.curLen += dlen;

                // 解析
                int bam_num = ParseAddAllBams(blockBuf.data, 0, blockBuf.curLen, blockData.bams);
                // spdlog::error(" bams: {} - {}", bam_num, blockData.bams.Size());
                read_bam_id += bam_num;
            }

            //spdlog::info("tid: {}, blockNum: {}, fid: {}, blocks: {}-{}, block range: {}-{}, {}", tid, blockNum, i, start, stop, startIdx, stopIdx,
            //             stopIdx - blockNum);

            blockNum += readData.startAddrArr.size();
        }
    }
}

// 计算对应的每个block在uncompress缓冲区的偏移量，并准备好那么多的空间
static void calculateOffset(Phase2PipelineArg& p) {
    for (int i = 0; i < p.midFiles.size(); ++i) {
        uint64_t offset = 0;
        int bamOffset = 0;
        auto& f = p.midFiles[i];
        auto& readData = f.readData[f.uncompressOrder % Phase2File::UNCOMPRESSS_BUF_NUM];
        for (int j = 0; j < readData.uncompressBlocks.Size(); ++j) {
            auto& block = readData.uncompressBlocks.arr[j];
            block.offset = offset;
            block.bamOffset = bamOffset;
            offset += block.block.curLen;
            bamOffset += block.bams.Size();
        }
        readData.bamNum = bamOffset;

        if (f.needUncompress) {
            auto& uncompressData = f.uncompressBuf[f.uncompressOrder % Phase2File::UNCOMPRESSS_BUF_NUM];
            uncompressData.Clear();
            uncompressData.bamArr.ReSize(bamOffset);
            uncompressData.blockBuf.ReSize(offset);
        }
    }
}

static void mtCopyToUncompressBuf(void* data, long idx, int tid) {
    Phase2PipelineArg& p = *(Phase2PipelineArg*)data;

    tid = idx;  // 静态分配任务，此时用idx代替tid
    int startIdx = START_IDX(tid, p.numThread, p.curBlockNum);
    int stopIdx = STOP_IDX(tid, p.numThread, p.curBlockNum);

    int blockNum = 0;
    int startFileIdx = -1;
    int stopFileIdx = 0;
    // 确定范围
    for (int i = 0; i < p.midFiles.size(); ++i) {
        auto& f = p.midFiles[i];
        if (f.needUncompress) {  // 还有空余缓冲区
            blockNum += f.readData[f.uncompressOrder % Phase2File::UNCOMPRESSS_BUF_NUM].startAddrArr.size();
        }
        if (startFileIdx == -1 && startIdx < blockNum) {
            startFileIdx = i;
        }
        if (blockNum >= stopIdx) {
            stopFileIdx = i;  // 包含
            break;
        }
    }
    blockNum = 0;
    for (int i = 0; i < startFileIdx; ++i) {
        auto& f = p.midFiles[i];
        if (f.needUncompress) {  // 还有空余缓冲区
            auto& readData = f.readData[f.uncompressOrder % Phase2File::UNCOMPRESSS_BUF_NUM];
            blockNum += readData.startAddrArr.size();
        }
    }

    for (int i = startFileIdx; i <= stopFileIdx; ++i) {
        auto& f = p.midFiles[i];
        int start = 0;
        int stop = 0;
        if (f.needUncompress) {  // 还有空余缓冲区
            auto& readData = f.readData[f.uncompressOrder % Phase2File::UNCOMPRESSS_BUF_NUM];
            auto& uncompressData = f.uncompressBuf[f.uncompressOrder % Phase2File::UNCOMPRESSS_BUF_NUM];

            start = MAX(0, startIdx - blockNum);
            stop = MIN(readData.startAddrArr.size(), stopIdx - blockNum);

            // 拷贝
            for (int j = start; j < stop; ++j) {
                auto& blockData = readData.uncompressBlocks.arr[j];
                auto& blockBuf = blockData.block;
                uncompressData.blockBuf.MemCopy(blockData.offset, blockData.block.data, blockData.block.curLen);
                for (int m = blockData.bamOffset, n = 0; n < blockData.bams.Size(); ++m, ++n) {
                    auto& bam = blockData.bams.arr[n];
                    bam.offset += blockData.offset;
                    uncompressData.bamArr.arr[m] = bam;
                }
            }
            blockNum += readData.startAddrArr.size();
        }
    }
}

bool doPhase2Uncompress(Phase2PipelineArg& p) {
    // spdlog::info("phase2 uncompress order: {}", p.uncompressOrder);
    PROF_G_BEG(phase2_uncompress);
    bool hasUncompress = false;
    int blockNum = 0;
    for (int i = 0; i < p.midFiles.size(); ++i) {
        auto& f = p.midFiles[i];
        f.needUncompress = false;
        if (f.readyReadBufNum > 0 && f.readyUncompressBufNum < Phase2File::UNCOMPRESSS_BUF_NUM) {  // 还有空余缓冲区
            blockNum += f.readData[f.uncompressOrder % Phase2File::UNCOMPRESSS_BUF_NUM].startAddrArr.size();
            f.needUncompress = true;
        }
    }
    p.curBlockNum = blockNum;
    block_id += blockNum;

#if 1
    kt_for(p.numThread, mtUncompressBlockBatch, &p, p.numThread);
    calculateOffset(p);
    kt_for(p.numThread, mtCopyToUncompressBuf, &p, p.numThread);
#endif

    for (int i = 0; i < p.midFiles.size(); ++i) {
        auto& f = p.midFiles[i];
        if (f.needUncompress) {  // 还有空余缓冲区
            f.readyReadBufNum -= 1;
            f.readyUncompressBufNum += 1;
            f.uncompressBuf[f.uncompressOrder % Phase2File::UNCOMPRESSS_BUF_NUM].startIdx = 0;  // 刚读入的，全部bam都还没拷贝到merge
            f.uncompressOrder += 1;
            // f.uncompressBuf->startIdx = 0; // 刚读入的，全部bam都还没拷贝到merge, 这个是错误的，buf是个数组，这个一直把第一个buf的startidx设置为0，bug来源


            hasUncompress = true;
        }

        // for test
        // f.readyUncompressBufNum = 0;
    }
    PROF_G_END(phase2_uncompress);

    spdlog::info("phase2 uncompress order: {}, read bams: {}, blocks: {}", p.uncompressOrder, read_bam_id, block_id);

    return hasUncompress;

    // exit(0);
}

void* phase2Uncompress(void* data) {
    Phase2PipelineArg& p = *(Phase2PipelineArg*)data;
    /* 2. do the work */
    while (true) {
        // previous dependency
        yarn::DEPENDENCY_NOT_TO_BE(p.readSig, 0);
        yarn::DEPENDENCY_NOT_TO_BE(p.uncompressSig, Phase2File::UNCOMPRESSS_BUF_NUM);

        if (p.readFinish) {
            while (p.uncompressOrder < p.readOrder) {
                //yarn::DEPENDENCY_NOT_TO_BE(p.uncompressSig, Phase2File::UNCOMPRESSS_BUF_NUM);
                doPhase2Uncompress(p);
                yarn::UPDATE_SIG_ORDER(p.uncompressSig, p.uncompressOrder);
            }
            yarn::SIGNAL_FINISH(p.uncompressSig, p.uncompressFinish);
            break;
        }
        doPhase2Uncompress(p);
        // if (doPhase2Uncompress(p)) { // 应该不需要判断？
        //  update status
        yarn::CONSUME_SIGNAL(p.readSig);
        yarn::UPDATE_SIG_ORDER(p.uncompressSig, p.uncompressOrder);
        // }
    }

    // spdlog::info("phase2 uncompress end order: {}", p.uncompressOrder);
    return nullptr;
}

//////////////////////////////////////////////// 拷贝到merge的缓冲区

int doPhase2CopyToMerge(Phase2PipelineArg& p) {
    PROF_G_BEG(phase2_copyToMerge);
    int usedUncompress = 0;
    p.uncompressReadyNum = 0;
    int maxUncompressReady = 0;
    size_t copied = 0;
#if 1
    for (int i = 0; i < p.midFiles.size(); ++i) {
        int usedUncompressInFile = 0; // 使用了多少uncompress缓冲区
        auto& f = p.midFiles[i];
        if (f.readyUncompressBufNum > 0 && f.readyMergeBufNum < f.COPY_BUF_NUM) {  // 有数据
            // 拷贝到merge数据里的循环数组和循环缓冲区
            auto& uncompressBuf = f.uncompressBuf[f.copyOrder % f.UNCOMPRESSS_BUF_NUM];
            auto& mergeData = f.mergeData;

            if (!mergeData.initialized) {
                // 初始化的时候确定merge buf的大小，已经对应的bam个数，以后每轮按照buf和bam的最小值拷贝
                mergeData.InitSize(uncompressBuf.blockBuf.curLen, uncompressBuf.bamArr.Size());
            }

            int copiedNum = mergeData.CopyBams(uncompressBuf.blockBuf, uncompressBuf.bamArr, uncompressBuf.startIdx, f.copyOrder, i);
            copied += copiedNum;
            uncompressBuf.startIdx += copiedNum;
            if (uncompressBuf.Size() == 0) {
                f.readyUncompressBufNum -= 1;
                usedUncompressInFile += 1;
                f.copyOrder += 1;
                // 检查还能不能继续添加数据
                if (f.readyUncompressBufNum > 0) {
                    auto& uncompressBuf = f.uncompressBuf[f.copyOrder % f.UNCOMPRESSS_BUF_NUM];
                    if (mergeData.hasSpace(uncompressBuf.Front()->wholeBamLen)) { // merge有空间，uncompress有数据，那就继续添加到merge
                        int copiedNum = mergeData.CopyBams(uncompressBuf.blockBuf, uncompressBuf.bamArr, uncompressBuf.startIdx, f.copyOrder, i);
                        copied += copiedNum;
                        uncompressBuf.startIdx += copiedNum;
                        if (uncompressBuf.Size() == 0) {
                            usedUncompressInFile += 1;
                            f.readyUncompressBufNum -= 1;
                            f.copyOrder += 1;
                        }
                    }
                }
            }
            f.readyMergeBufNum += 1; // 当前文件的merge data数据准备好了

            // check bam
            // for (int j = 0; j < mergeData.Size(); ++j) {
            //     auto& bam = mergeData.bams[j];
            //     CheckBam(bam.addr + bam.offset, bam.wholeBamLen);
            // }

            //spdlog::info("copy num: {}", copiedNum);
            maxUncompressReady = MAX(maxUncompressReady, f.readyUncompressBufNum);
// 调试用
#if 0
            // for test，消耗mergedata
            //f.mergeData.bams.Clear();
            //f.mergeData.data.Clear();
            int clearSize = mergeData.Size() / 3;
            //int clearSize = 110000;
            //int clearSize = mergeData.Size();
            for (int j = 0; j < clearSize; ++j) {
                mergeData.Pop();
            }
            f.readyMergeBufNum -= 1;
#endif
        }

        usedUncompress = MAX(usedUncompressInFile, usedUncompress);
        p.uncompressReadyNum += f.readyUncompressBufNum;
    }
#else
    p.midFiles[p.copyMergeOrder % p.midFiles.size()].readyUncompressBufNum -= 1;
#endif
    copy_bam_id += copied;

    PROF_G_END(phase2_copyToMerge);

    // for test，模拟消耗merge缓冲区的数据
    // 调试用
#if 0
    yarn::CONSUME_SIGNAL(p.copyMerge);
#endif

    usedUncompress = MAX(usedUncompress, Phase2File::UNCOMPRESSS_BUF_NUM - maxUncompressReady);

    //spdlog::info("copy to merge order: {}-{}, copied: {}, ready num: {}, max ready: {}, last id: {}", p.copyMergeOrder, usedUncompress, copied,
    //             p.uncompressReadyNum, maxUncompressReady, copy_bam_id);


    return usedUncompress;
}

void* phase2CopyToMergeBuf(void* data) {
    Phase2PipelineArg& p = *(Phase2PipelineArg*)data;
    /* 2. do the work */
    while (true) {
        // previous dependency
        yarn::DEPENDENCY_NOT_TO_BE(p.uncompressSig, 0);
        yarn::DEPENDENCY_NOT_TO_BE(p.copyMerge, Phase2File::COPY_BUF_NUM);

        if (p.uncompressFinish) {
            while (p.uncompressReadyNum > 0 || yarn::PEEK_LOCK(p.uncompressSig) > 0) {  // 一直有没处理完的解压数据
                yarn::DEPENDENCY_NOT_TO_BE(p.copyMerge, Phase2File::COPY_BUF_NUM);
                int usedUncompress = doPhase2CopyToMerge(p);
                usedUncompress = MIN(usedUncompress, yarn::PEEK_LOCK(p.uncompressSig));
                yarn::CONSUME_SIGNAL_BY(p.uncompressSig, usedUncompress);
                yarn::UPDATE_SIG_ORDER(p.copyMerge, p.copyMergeOrder);
            }
            yarn::SIGNAL_FINISH(p.copyMerge, p.copyMergeFinish);
            break;
        }
        int usedUncompress = doPhase2CopyToMerge(p);
        usedUncompress = MIN(usedUncompress, yarn::PEEK_LOCK(p.uncompressSig));

        // update status
        yarn::CONSUME_SIGNAL_BY(p.uncompressSig, usedUncompress);
        //yarn::CONSUME_SIGNAL(p.uncompressSig);
        yarn::UPDATE_SIG_ORDER(p.copyMerge, p.copyMergeOrder);
    }

    spdlog::info("phase2 copyMerge order: {}", p.copyMergeOrder);
    return nullptr;
}