/*
     Description: 写入最终排序后的文件

     Copyright : All right reserved by ICT

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

#include "phase_2_write.h"

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

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

static size_t all_bams = 0;

static void checkCompress(uint8_t* addr, uint64_t len) {
    size_t curReadPos = 0;
    int blockLen = 0;
    int maxBlockLen = 0;
    DataBuffer buf;
    buf.AllocMem(SINGLE_BLOCK_SIZE);

    while (curReadPos + BLOCK_HEADER_LENGTH <= len) { /* 确保能解析block长度 */
        blockLen = unpackInt16(&addr[curReadPos + 16]) + 1;
        if (blockLen > maxBlockLen) {
            maxBlockLen = blockLen;
        }
        if (curReadPos + blockLen <= len) {   /* 完整的block数据在buf里 */
            size_t dlen = SINGLE_BLOCK_SIZE;  // 65535
            uint32_t crc = le_to_u32(addr + curReadPos + blockLen - 8);
            int ret = bgzfUncompress(buf.data, &dlen, (Bytef*)(addr + curReadPos) + BLOCK_HEADER_LENGTH, blockLen - BLOCK_HEADER_LENGTH, crc);
            if (ret != 0) {
                spdlog::error("block len: {}, uncompressed len: {}", blockLen, dlen);
                exit(0);
            }

            curReadPos += blockLen;
        } else {
            spdlog::error("not valid compressed block: {}, {}", curReadPos + blockLen, len);
            break; /* 当前block数据不完整，一部分在还没读入的file数据里 */
        }
    }
    if (curReadPos != len) {
        spdlog::error("addr: {}, len: {}", curReadPos, len);
        exit(0);
    }
}

//////////////////////////////压缩
static void mtCompressBlock(void* data, long idx, int tid) {
    Phase2PipelineArg& p = *(Phase2PipelineArg*)data;
    auto& mergeData = p.mergeData[p.compressOrder % p.COMPRESS_BUF_NUM];
    auto& blockData = mergeData.blockDataArr[idx];
    auto& compressedData = mergeData.compressDataArr[idx];

    compressedData.ReAllocMem(SINGLE_BLOCK_SIZE);  // 压缩后的block数据不会超过单个block的大小
    compressedData.curLen = SINGLE_BLOCK_SIZE;
    bgzfCompress(compressedData.data, &compressedData.curLen, blockData.blockBuf.data, blockData.blockBuf.curLen, p.compressLevel);
    // spdlog::info("bytes: {}", blockData.blockBuf.curLen);
    // checkCompress(compressedData.data, compressedData.curLen);
}

static void doCompress(Phase2PipelineArg& p) {
    PROF_G_BEG(compress);
    auto& compressBuf = p.compressBuf[p.compressOrder % p.COMPRESS_BUF_NUM];
    auto& mergeData = p.mergeData[p.compressOrder % p.MERGE_BUF_NUM];
    compressBuf.Clear();

    kt_for(p.numThread, mtCompressBlock, &p, mergeData.Size());

    for (int i = 0; i < mergeData.Size(); ++i) {
        compressBuf.MemCopy(mergeData.compressDataArr[i].data, mergeData.compressDataArr[i].curLen);
        all_bams += mergeData.blockDataArr[i].bamNum;
    }
    mergeData.Clear();
    spdlog::info("compress order: {}, bytes: {}, bams: {}", p.compressOrder, compressBuf.curLen, all_bams);
    PROF_G_END(compress);

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

}

/* phase1Compress step- 压缩线程 */
void* phase2Compress(void* data) {
    Phase2PipelineArg& p = *(Phase2PipelineArg*)data;
    /* do the work */
    while (true) {
        // previous dependency
        yarn::DEPENDENCY_NOT_TO_BE(p.mergeSig, 0);
        yarn::DEPENDENCY_NOT_TO_BE(p.compressSig, p.COMPRESS_BUF_NUM);

        if (p.mergeFinish) {
            while (p.compressOrder < p.mergeOrder) {
                yarn::DEPENDENCY_NOT_TO_BE(p.compressSig, p.COMPRESS_BUF_NUM);
                doCompress(p);
                yarn::UPDATE_SIG_ORDER(p.compressSig, p.compressOrder);
            }
            yarn::SIGNAL_FINISH(p.compressSig, p.compressFinish);
            break;
        }

        doCompress(p);

        // update status
        yarn::CONSUME_SIGNAL(p.mergeSig);
        yarn::UPDATE_SIG_ORDER(p.compressSig, p.compressOrder);
    }

    spdlog::info("End phase2 compress order: {}", p.compressOrder);
    return nullptr;
}

/////////////////////////////// 写文件
static void doPhase2Write(Phase2PipelineArg& p) {
    PROF_G_BEG(write_final);
    DataBuffer& compressBuf = p.compressBuf[p.writeOrder % p.COMPRESS_BUF_NUM];
    // checkCompress(compressBuf.data, compressBuf.curLen);
    fwrite(compressBuf.data, 1, compressBuf.curLen, p.outFilePtr);
    PROF_G_END(write_final);
    // spdlog::info("phase2 write order: {}, bytes: {}", p.writeOrder, compressBuf.curLen);
}

void* phase2Write(void* data) {
    Phase2PipelineArg& p = *(Phase2PipelineArg*)data;
    /* do the work */
    while (true) {
        // previous dependency
        yarn::DEPENDENCY_NOT_TO_BE(p.compressSig, 0);

        if (p.compressFinish) {
            while (p.writeOrder < p.compressOrder) {
                doPhase2Write(p);
                p.writeOrder += 1;
            }
            break;
        }
        doPhase2Write(p);
        // update status
        yarn::CONSUME_SIGNAL(p.compressSig);
        p.writeOrder += 1;
    }

    spdlog::info("End phase2 write order: {}", p.writeOrder);
    return nullptr;
}