/*
     Description: 排序入口

     Copyright : All right reserved by ICT

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

#include "sort.h"

#include <fcntl.h>
#include <htslib/sam.h>
#include <htslib/thread_pool.h>
#include <klib/kthread.h>
#include <math.h>
#include <pthread.h>
#include <spdlog/spdlog.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/mman.h>
#include <unistd.h>
#include <zlib.h>

#include <algorithm>
#include <string>

#include "common_data.h"
#include "const_val.h"
#include "global_vars.h"
#include "phase_1.h"
#include "phase_1_read.h"
#include "phase_1_uncompress.h"
#include "phase_1_write.h"
#include "phase_2.h"
#include "sam_io.h"
#include "sort_args.h"
#include "util/profiling.h"
#include "util/yarn.h"

using std::string;

#define PARAM_BLOCK_SIZE 16L * 1024 * 1024


/* 对sam文件进行排序 */
static void samSortFirstPipe() {

}

static void bamSortPipeline() {
    /* set up phase 1 parameters */
    Phase1PipelineArg p1;
    p1.SetThreadNum(nsgv::gSortArg.NUM_THREADS);
    // 平均每线程4M缓冲区，累加起来，用来读入文件（BAM/SAM）（相对解压之后的缓冲区，大小可以忽略）
    const size_t kReadBufSize = 4L * 1024 * 1024 * p1.numThread;
    // p1.midFileNamePrefix = getFileNameWithoutExt(nsgv::gSortArg.OUTPUT_FILE) + "_sort_mid.";
    p1.midFileNamePrefix = getFilePathWithoutExt(nsgv::gSortArg.OUTPUT_FILE) + "_sort_mid.";  // 跟输出文件路径相同
    spdlog::info("prefix: {}", p1.midFileNamePrefix);
    p1.maxMemBytes = nsgv::gSortArg.MAX_MEM;

    p1.uncompressBufBytes = p1.maxMemBytes;  // 可以再考虑一下

    for (int i = 0; i < p1.UNCOMPRESS_BUF_NUM; ++i) {
        p1.threadUncompressWrap[i].Resize(p1.numThread);  // 每个线程的解压block数组初始大小，后续如果不够用会自动扩容
    }
    // 每个线程解压数据的内存上限，平均每线程4M读入数据，解压后大概5倍，留点余量，所以乘以5，除以线程数
    p1.singleThreadMemBytes = kReadBufSize * BAM_COMPRESS_RATIAO / p1.numThread;  
    spdlog::info("max mem: {}, uncompress mem: {}, single thread mem: {}", nsgv::gSortArg.MAX_MEM, p1.uncompressBufBytes,
                 p1.singleThreadMemBytes);
    for (int i = 0; i < p1.READ_BUF_NUM; ++i) {
        p1.readData[i].Resize(kReadBufSize);
    };
    // 根据最大内存参数，计算初始化开辟的空间
    p1.uncompressData.Resize(p1.uncompressBufBytes);

    // 运行第一阶段
    phase1Pipeline(p1);

    /////////////////////////////////////////////////////

#if 1
    // 第二阶段参数初始化
    Phase2PipelineArg p2(p1.uncompressData, p1.allBams, p1.midFileNamePrefix, p1.midFileOrder, p1.numThread);

    p2.outFilePtr = fopen(nsgv::gSortArg.OUTPUT_FILE.c_str(), "wb");
    if (p2.outFilePtr == nullptr) {
        spdlog::error("failed to create output file: {}", nsgv::gSortArg.OUTPUT_FILE.c_str());
        exit(1);
    }
    // 处理bam的header
    fwrite(nsgv::gInHdr.compressed.data, 1, nsgv::gInHdr.compressed.curLen, p2.outFilePtr);
    // 运行第二阶段
    phase2Pipeline(p2);
    // 写结尾
    fwrite("\037\213\010\4\0\0\0\0\0\377\6\0\102\103\2\0\033\0\3\0\0\0\0\0\0\0\0\0", 1, 28, p2.outFilePtr);
    fclose(p2.outFilePtr);
#endif

}

// 排序的入口函数，entry function
int doSort() {

#if 1
    gfp[0] = fopen("output/f0.txt", "w");
    gfp[1] = fopen("output/f1.txt", "w");
    gfp[2] = fopen("output/f2.txt", "w");
    gfp[3] = fopen("output/f3.txt", "w");
#endif

#if 1
    nsgv::gIsBigEndian = ed_is_big();
    
    // 并行流水线排序bam文件
    bamSortPipeline();

    spdlog::info("OneBam size: {}", sizeof(OneBam));
#else

    /* 打开输入bam文件 */
    samFile *inBamFp = sam_open_format(nsgv::gSortArg.INPUT_FILE.c_str(), "r", nullptr);
    if (!inBamFp) {
        spdlog::error("[{}] load sam/bam file failed.\n", __func__);
        return -1;
    }
    hts_set_opt(inBamFp, HTS_OPT_BLOCK_SIZE, PARAM_BLOCK_SIZE);
    sam_hdr_t *inBamHdr = sam_hdr_read(inBamFp);  // 读取header
    
    htsThreadPool htsPoolRead = {NULL, 0};        // 多线程读取，创建线程池
    htsPoolRead.pool = hts_tpool_init(nsgv::gSortArg.NUM_THREADS);
    if (!htsPoolRead.pool) {
        spdlog::error("[{}] failed to set up thread pool", __LINE__);
        sam_close(inBamFp);
        return -1;
    }
    hts_set_opt(inBamFp, HTS_OPT_THREAD_POOL, &htsPoolRead);

    // 测试读取bam的速度

    int max_bam_len = 0;
    bam1_t *bamp = bam_init1();
    uint64_t bam_num = 0;
    while (sam_read1(inBamFp, inBamHdr, bamp) >= 0) {
        bam_num++;
        if (max_bam_len < bamp->l_data) max_bam_len = bamp->l_data;
        if (bamp->l_data > 1000) {
            spdlog::info("large record len: {}", bamp->l_data);
        }
        if (bam_num % 10000000 == 0) {
            spdlog::info("bam num: {}, max bam len: {}", bam_num, max_bam_len);
        }

        // fprintf(gfp[0], "%d-%ld\n", bamp->core.tid, bamp->core.pos);
    }
    sam_close(inBamFp);
    spdlog::info("max record len: {}", max_bam_len);
    spdlog::info("bam num: {}", bam_num);

#endif

#if 1
    fclose(gfp[0]);
    fclose(gfp[1]);
    fclose(gfp[2]);
    fclose(gfp[3]);
#endif

    return 0;
}
