/*
     Description: 排序相关的数据结构

     Copyright : All right reserved by ICT

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

#pragma once

#include <inttypes.h>

#include <queue>
#include <vector>

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

using std::priority_queue;
using std::vector;

#define START_IDX(i, nt, nele) ((i) * (nele) / (nt))
#define STOP_IDX(i, nt, nele) (((i) + 1) * (nele) / (nt))

#undef MAX
#undef MIN
#define MAX(x, y)((x) > (y) ? (x) : (y))
#define MIN(x, y)((x) < (y) ? (x) : (y))

/* for step-1 read data from bam file */
struct ReadBuffer {
    uint8_t *dataBuf = nullptr;
    uint8_t *blockBuf = nullptr; // 用来保存上一轮没能完整读取的block
    size_t readBufSize = 0; // 读入的buf大小
    vector<uint8_t *> startAddrArr; // 存放每个block的起始地址
    ReadBuffer() { }
    ReadBuffer(size_t readBufSize_) {
        readBufSize = readBufSize_;
        dataBuf = (uint8_t *)malloc(readBufSize);
        blockBuf = (uint8_t *)malloc(SINGLE_BLOCK_SIZE);
    }
    ~ReadBuffer() {
        if (dataBuf) {
            free(dataBuf);
            dataBuf = nullptr;
        }
        if (blockBuf) {
            free(blockBuf);
            blockBuf = nullptr;
        }
        readBufSize = 0;
        startAddrArr.clear();
    }
    void Resize(size_t readBufSize_) {
        if (dataBuf) free(dataBuf);
        if (blockBuf) free(blockBuf);
        readBufSize = readBufSize_;
        dataBuf = (uint8_t *)malloc(readBufSize);
        blockBuf = (uint8_t *)malloc(SINGLE_BLOCK_SIZE);
    }
};

/* 用来串行保存解压后的block数据, 跟ReadBuffer差不多*/
struct UncompressBlockBuffer {
    uint8_t *dataBuf = nullptr; // 用来保存解压后的block数据，串行往这里添加解压后的block数据
    uint64_t dataBufSize = 0; // 存放的解压之后的block的buf大小
    uint64_t usedBufSize = 0; // 已经使用的buf大小
    uint64_t lastEndPos = 0; // 上一轮计算完整bam结束位置，也是下一轮第一个bam（不完整）的起始数据

    int avgBamSize = 0; // bam的平均长度，用来指导后续的解压和排序
    int avgBamNumPerBlock = 0; // 每个block里bam的平均数量，用来指导后续的解压和排序

    UncompressBlockBuffer() { }
    UncompressBlockBuffer(uint64_t dataBufSize_) {
        dataBufSize = dataBufSize_;
        dataBuf = (uint8_t *)malloc(dataBufSize);
    }
    ~UncompressBlockBuffer() {
        if (dataBuf) {
            free(dataBuf);
            dataBuf = nullptr;
        }
        dataBufSize = 0;
        usedBufSize = 0;
        lastEndPos = 0;
    }
    void Resize(uint64_t dataBufSize_) {
        if (dataBuf) free(dataBuf);
        dataBufSize = dataBufSize_;
        dataBuf = (uint8_t *)malloc(dataBufSize);
    }

    void ReAllocate(uint64_t dataBufSize_) {
        if (dataBufSize < dataBufSize_) {
            dataBufSize = dataBufSize_;
            dataBuf = (uint8_t*)realloc(dataBuf, dataBufSize);
        }
    }

    void Clear() {
        usedBufSize = 0;
        lastEndPos = 0;
    }

    uint8_t *HandOverBuf() {
        uint8_t *handOverBuf = dataBuf;
        dataBuf = nullptr; // 交出buf的所有权，外部需要负责释放内存
        return handOverBuf;
    }

    void NextRound() {
        usedBufSize -= lastEndPos;
        memcpy(dataBuf, dataBuf + lastEndPos, usedBufSize);
        lastEndPos = 0;
    }
};

/* 对vector的一个包装，避免频繁内存分配和释放 */
template <typename T>
struct FastVector {
    vector<T> arr;
    size_t curIdx = 0;  // 当前已经添加的元素数量，也是下一个要添加的元素的索引
    T& Add() {
        if (curIdx < arr.size())
            return arr[curIdx++];
        else {
#if 0
            arr.resize((arr.size() + 1) << 1);
            return arr[curIdx++];

#else
            arr.push_back(T());
            curIdx++;
            return arr.back();
#endif
        }
    }

    T& Get(size_t idx) { return arr[idx]; }

    T* Back() {
        if (curIdx == 0) {
            return nullptr;
        }
        return &arr[curIdx - 1];
    }

    void Add(const T& item) {
        T &newItem = Add();
        newItem = item;
    }
    void AddSize(size_t num) {
        if (curIdx + num > arr.size()) {
            arr.resize(curIdx + num);
        }
    }
    void ReAllocate(size_t num) {
        if (num > arr.size()) {
            arr.resize(num);
        }
    }

    void ReSize(size_t num) {
        ReAllocate(num);
        curIdx = num;
    }

    void ClearResize(size_t num) {
        if (num > arr.size()) {
            arr.clear();
            arr.resize(num);
        }
        curIdx = num;
    }

    size_t Size() const { return curIdx; }
    size_t Capacity() const { return arr.size(); }
    void Clear() { curIdx = 0; }
};

/* for step-2 parallel uncompress gz blocks，还有缓存满了之后的线程内排序，用来排序的*/
struct OneBam {
    static constexpr int QnameOffset = 36;  // 距离该sam记录开头地址的偏移量，包含了4字节bam长度，32字节的bam core信息，后面紧跟着qname字符串
    uint32_t qnameLen;    // 序列名字长度
    uint32_t wholeBamLen = 0; // 包含4字节bam长度的所有bam内容长度，uint16只适合二代
    int32_t tid = 0;     // 比对到的染色体
    int64_t pos = 0;     // mapping 位置
    uint8_t *addr = nullptr; // 首地址
    uint64_t offset = 0;  // 距离首地址的偏移量
};

typedef FastVector<OneBam> BamArr;
typedef FastVector<const OneBam*> BamPtrArr;

extern bool CheckBam(uint8_t* addr, int len);