/*
     Description: bam是个外排序过程，整个过程包含两个阶段，这是第一个阶段
        * 第一阶段（phase-1）：
        * 1. 读入bam文件header（要注意如果是sam或者管道输入则不需要解压）
        * 2. phase-1-stage-1: 读入待解压的数据块，直到填满phase-1-stage-1的buffer，读入数据期间分割压缩的block
        * 3. phase-1-stage-2: 多线程解压数据块，每个线程解压后放在自己的buf里，乱序解压，线程buf满后，进行线程内的排序
        * 4. phase-1-stage-3: 线程间进行归并排序，然后压缩写入中间文件

     Copyright : All right reserved by ICT

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

#pragma once

#include <stdint.h>

#include <string>

#include "const_val.h"
#include "sam_io.h"
#include "sort.h"
#include "util/yarn.h"

using std::string;


// 把线程解压需要的数据放到一个结构体里
struct ThreadUncompressData {
    DataBuffer blockBuf; // 解压的block放在这里
    size_t blockNum = 0; // 解压的block数量
    BamArr bamArr; // 解析后的bam数据放在这里
    BamArr firstBam; // 连接上一个不完整的bam数据，当作当前block的第一个bam
    DataBuffer lastBamBuf;  // 最后那个不完整的bam解析时用到的缓冲区
    int64_t memOffset = 0;  // 这个线程解压的数据在全局解压数据中的偏移位置
    int64_t bamOffset = 0; // 这个线程解析的bam数据在全局解压数据中的偏移位置


    void Resize(int vecInitSize) {
        blockBuf.AllocMem(vecInitSize * SINGLE_BLOCK_SIZE);  // 每个线程的解压block数组初始大小，后续如果不够用会自动扩容
        bamArr.arr.reserve(vecInitSize);
        firstBam.arr.reserve(1);
        lastBamBuf.AllocMem(SINGLE_BLOCK_SIZE);
    }

    void Reset() {
        blockNum = 0;
        blockBuf.Clear();
        bamArr.Clear();
        firstBam.Clear();
        lastBamBuf.Clear();
        memOffset = 0;
        bamOffset = 0;
    }

    size_t GetBlockNum() const { return blockNum; }
    size_t GetBamNum() const { return bamArr.Size() + firstBam.Size(); }
};

// 第一阶段用到的数据结构
// 第一阶段的解压、排序、归并、压缩
struct ThreadUncompressWrap {
    vector<ThreadUncompressData> threadUncompressDataArr; // 每个线程一个，用来保存解压数据和解析bam数据

    ThreadUncompressWrap() {}
    ThreadUncompressWrap(int numThread) { Resize(numThread); }
    ThreadUncompressWrap(int numThread, int vecInitSize) { Resize(numThread, vecInitSize); }
    void Resize(int numThread) { Resize(numThread, 128); }
    void Resize(int numThread, int vecInitSize) {
        threadUncompressDataArr.resize(numThread);
        for (int i = 0; i < numThread; ++i) {
            threadUncompressDataArr[i].Resize(vecInitSize);
        }
    }
    void ResetBlockArr() {
        for (int i = 0; i < threadUncompressDataArr.size(); ++i) {
            threadUncompressDataArr[i].Reset();
        }
    }

    uint64_t GetTotalBlockNum() {
        uint64_t totalBlockNum = 0;
        for (int i = 0; i < threadUncompressDataArr.size(); ++i) {
            totalBlockNum += threadUncompressDataArr[i].GetBlockNum();
        }
        return totalBlockNum;
    }

    uint64_t GetTotalBamNum() {
        uint64_t totalBamNum = 0;
        for (int i = 0; i < threadUncompressDataArr.size(); ++i) {
            totalBamNum += threadUncompressDataArr[i].GetBamNum();
        }
        return totalBamNum;
    }
};

struct BlockBams {
    DataBuffer blockBuf; // 解压的block放在这里
    BamPtrArr bamPtrArr; // 解析后的bam数据放在这里
    int bamNum = 0; // 解析后的bam数量，调试用

    void Clear() {
        blockBuf.Clear();
        bamPtrArr.Clear();
        bamNum = 0;
    }
};

// 用于合并压缩的数据结构
struct MergeCompressData {
    vector<BlockBams> blockDataArr;     // 待压缩的数据
    vector<DataBuffer> compressDataArr; // 压缩后的数据
    int curIdx = 0;
    int curBytes = 0;
    int lastRoundIdx = 0; // 上一轮已经处理的block数量

    void Resize(int blockNum) {
        blockDataArr.resize(blockNum);
        compressDataArr.resize(blockNum);
    }

    int Size() { return curIdx; }

    void Clear() {
        curIdx = 0;
        curBytes = 0;
        lastRoundIdx = 0;
    }
};

/* 第一阶段的多线程流水线参数 */
struct Phase1PipelineArg {
    static const int READ_BUF_NUM = 2; // 读入的buf数量
    static const int UNCOMPRESS_BUF_NUM = 1; // 解压的buf数量，1表示解压和拷贝是串行，2，表示并行

    static const int MERGE_BUF_NUM = 2; // 归并的buf数量
    static const int COMPRESS_BUF_NUM = 2;  // 压缩的buf数量

    // common parameters
    int numThread = 0; // 线程数
    uint64_t singleThreadMemBytes = 0; // 单线程开辟的内存字节上限
    uint64_t uncompressBufBytes = 0; // 总的解压缓冲区大小
    uint64_t maxMemBytes = 0; // 最大可使用内存，参数
    int compressLevel = 1; // 中间文件压缩级别，参数

    // for read-uncompress-parse
    uint64_t readOrder = 0; // 读取文件轮次编号，与下边的uncompressOrder对应
    uint64_t uncompressOrder = 0; // 并行解压gz block, 包含排序（缓冲区满之后排序），以及合并之后的压缩
    uint64_t memCopyOrder = 0; // 串行拷贝解压数据到uncompressData的轮次编号，与上边的uncompressOrder对应

    volatile int readFinish = 0;
    volatile int uncompressFinish = 0;

    yarn::lock_t* readSig;
    yarn::lock_t* uncompressSig;

    ReadBuffer readData[READ_BUF_NUM];  // 用来读如数据，双缓冲
    ThreadUncompressWrap threadUncompressWrap[UNCOMPRESS_BUF_NUM];  // 每个thread一个，用来保存解压后的block数据
    DataBuffer lastRoundBuf;                                            // 上一轮不完整的数据放在这里

    UncompressBlockBuffer uncompressData;  // 所有线程共用一个，串行往这里添加解压后的block数据
    BamArr allBams; // 所有线程共用一个，串行往这里添加解析后的bam数据
    // vector<const OneBam*> lastBams;

    // 判断bam是否有效的阈值，暂时没用了
    int maxSeqLen = 0; // bam里seq的最大长度，初始值是int的最大值，后续会根据解压的bam数据更新这个值，作为判断bam是否合法的一个条件
    int maxBamLen = 0; // bam的最大长度，初始值是int的最大值，后续会根据解压的bam数据更新这个值，作为判断bam是否合法的一个条件

    // 主要用来统计信息，测试用
    uint64_t bamNum = 0; // 解压后的bam数量
    uint64_t blockNum = 0; // 解压后的block数量

    ///////////////////////////////////////////////////// for sort-merge-compress-write

    // for merge-compress-write
    int midFileOrder = 0;  // 用来给中间文件编号
    int mergeBlocksThreshold = 0; // 进行归并排序的，单次block数量阈值，参数

    uint64_t mergeOrder = 0;
    uint64_t compressOrder = 0;  // 排序后压缩，这个和下边的writeOder对应，跟上边的order不相关
    uint64_t writeOrder = 0; // 串行合并解压后的blocks，并解析每个bam的长度，达到内存阈值后，并行排序

    volatile int mergeFinish = 0;
    volatile int compressFinish = 0;

    yarn::lock_t* mergeSig;
    yarn::lock_t* compressSig;

    MergeCompressData mergeCompressData[MERGE_BUF_NUM]; // 归并压缩用到的数据结构
    DataBuffer compressBuf[COMPRESS_BUF_NUM]; // 压缩用到的buf
    string midFileNamePrefix; // 中间文件名前缀
    FILE* midFilePtr = nullptr; // 中间文件指针

    // 统计用
    vector<uint64_t> uncompressedDataSize; // nthread数组长度
    vector<uint64_t> compressedDataSize;

    Phase1PipelineArg() {
        readSig = yarn::NEW_LOCK(0);
        uncompressSig = yarn::NEW_LOCK(0);
        mergeSig = yarn::NEW_LOCK(0);
        compressSig = yarn::NEW_LOCK(0);
        lastRoundBuf.AllocMem(SINGLE_BLOCK_SIZE);
    }

    void SetThreadNum(int nThread) {
        numThread = nThread;
        // const int kThres = 1024;
        const int kThres = 2048;
        int threadBlocksThres = numThread * 32;
        mergeBlocksThreshold = threadBlocksThres < kThres ? threadBlocksThres : kThres;  // 大概64M的解压缩后的数据
        for (int i = 0; i < MERGE_BUF_NUM; ++i) {
            mergeCompressData[i].Resize(mergeBlocksThreshold);
        }
    }
};




// 排序第一阶段，并行流水线执行程序
void phase1Pipeline(Phase1PipelineArg &p);

std::string getFileNameWithoutExt(const std::string& filepath);

std::string getFilePathWithoutExt(const std::string& filepath);