/*
     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 "phase_2.h"
#include "util/profiling.h"

//////////////////////////////压缩
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);
}

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.blockDataArr.size());

    for (int i = 0; i < mergeData.blockDataArr.size(); ++i) {
        compressBuf.MemCopy(mergeData.compressDataArr[i].data, mergeData.compressDataArr[i].curLen);
    }
    mergeData.Clear();
    // spdlog::info("compress bytes: {}", compressBuf.curLen);
    PROF_G_END(compress);
}

/* 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];
    fwrite(compressBuf.data, 1, compressBuf.curLen, p.outFilePtr);
    PROF_G_END(write_final);
    spdlog::info("phase2 write order: {}", p.writeOrder);
}

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;
}