/*
     Description: 对原始bam进行部分数据解析，以更方便后续处理

     Copyright : All right reserved by ICT

     Author : Zhang Zhonghai
     Date : 2025/12/29
*/
#pragma once

#include <cstdint>

#include "bam_wrap.h"
#include "stable_array.h"
#include "base_utils.h"

// 对cigar进行简单包装
struct Cigar {
    char op = '0';
    int len = 0;
    // 该操作符是否消耗read的碱基
    static bool ConsumeReadBases(char cigar) { return cigar == 'M' || cigar == '=' || cigar == 'X' || cigar == 'I' || cigar == 'S'; }
    // 该操作符是否消耗参考基因组的碱基
    static bool ConsumeRefBases(char cigar) { return cigar == 'M' || cigar == '=' || cigar == 'X' || cigar == 'D' || cigar == 'N'; }
};

// 在read中某个特定位置对应的cigar操作符
struct ReadIdxCigar {
    int readIdx = -1;    // 在read序列中的位置
    char cigarOp = '0';  // 当前位置对应的cigar
};

// 对原始bam进行部分数据解析，以更方便后续处理
struct SamData {
    int64_t rid = 0;           // for debug

    // 这些数据用于最开始clipping等操作
    int read_len = 0;          // read长度，各种clip之后的长度
    int cigar_start = 0;       // cigar起始位置，闭区间
    int cigar_end = 0;         // cigar结束位置，开区间
    int first_cigar_clip = 0;  // 第一个cigar, clip的数量，切左侧
    int last_cigar_clip = 0;   // 最后一个cigar, clip的数量，切右侧
    int left_clip = 0;         // 左侧被切掉的碱基长度，BI和BD质量分数也会用到
    int right_clip = 0;        // 右侧被切掉的碱基长度
    int ref_offset = 0;        // 切除adapter和softclip之后(softclip应该不影响)，相对原始ref比对位置（contig_pos）的偏移量

    // 记录一下bqsr运算过程中用到的数据，回头提前计算一下，修正现在的复杂逻辑
    static constexpr int READ_INDEX_NOT_FOUND = -1;

    BamWrap* bw;
    int64_t start_pos;              // 因为soft clip都被切掉了，这里的softstart应该就是切掉之后的匹配位点，闭区间
    int64_t end_pos;                // 同上，闭区间
    StableArray<char> bases; // 处理之后的read的碱基数组形式
    StableArray<uint8_t> base_quals;  // 对应的质量分数
    StableArray<uint8_t> ins_quals;   // insert质量分数, BI (大部分应该都没有)
    StableArray<uint8_t> del_quals;   // delete质量分数, BD (大部分应该都没有)
    StableArray<Cigar> cigars;

    int64_t& softStart() { return start_pos; }
    int64_t& softEnd() { return end_pos; }

    // functions

    // 初始化
    void init() {
        bases.clear();
        base_quals.clear();
        ins_quals.clear();
        del_quals.clear();
        cigars.clear();
        rid = 0;
        read_len = 0;
        cigar_start = 0;
        cigar_end = 0;
        first_cigar_clip = 0;
        last_cigar_clip = 0;
        left_clip = 0;
        right_clip = 0;
        ref_offset = 0;
        bw = nullptr;
        start_pos = 0;
        end_pos = 0;
    }

    // 初步解析bam
    void parseBasic(BamWrap *_bw) {
        bw = _bw;
        read_len = BamWrap::BamEffectiveLength(bw->b);
        cigar_end = bw->b->core.n_cigar;
    }

    // 应用各种转换操作，更新SamData中的数据
    void applyTransformations() {
        const uint32_t* cigar = bam_get_cigar(bw->b);
        const bam1_core_t& bc = bw->b->core;
        int i = 0;

        // 计算ref_offset，就是相对比对的position，要将ref右移多少
        for (i = 0; i < cigar_start; ++i) {
            if (BaseUtils::consumeRefBases(bam_cigar_opchr(cigar[i]))) {
                ref_offset += bam_cigar_oplen(cigar[i]);
            }
        }
        if (BaseUtils::consumeRefBases(bam_cigar_opchr(cigar[i]))) {
            ref_offset += first_cigar_clip;
        }

        // 计算clip处理之后，剩余的碱基
        bases.resize(read_len);
        base_quals.resize(read_len);
        uint8_t* seq = bam_get_seq(bw->b);
        uint8_t* quals = bam_get_qual(bw->b);
        for (i = 0; i < read_len; ++i) {
            bases[i] = BaseUtils::cBaseToChar[bam_seqi(seq, i + left_clip)];
            base_quals[i] = quals[i + left_clip];
        }

        // 计算read两端clip之后的softstart和softend，其实S之前都被切掉了
        int64_t softStart = bw->b->core.pos + ref_offset;
        int64_t softEnd = softStart - 1;  // 闭区间
        bool rightTail = false;
        for (i = cigar_start; i < cigar_end; ++i) {
            const char c = bam_cigar_opchr(cigar[i]);
            int len = bam_cigar_oplen(cigar[i]);
            if (i == cigar_start) len -= first_cigar_clip;
            if (i == cigar_end - 1) len -= last_cigar_clip;
            // if (c == 'S' || c == 'I')
            if (c == 'S') softStart -= len;
            else if (c != 'H') rightTail = true;
            if (rightTail && (BaseUtils::consumeRefBases(c) || c == 'S')) softEnd += len;
        }
        this->softStart() = softStart; // 其实这里的softStart就是start
        this->softEnd() = softEnd; // 同上

        // 计算clip之后的cigar，其实可以考虑下边的代码和上边的换一下位置
        for (i = cigar_start; i < cigar_end; ++i) {
            char c = bam_cigar_opchr(cigar[i]);
            int len = bam_cigar_oplen(cigar[i]);
            if (i == cigar_start) len -= first_cigar_clip;
            if (i == cigar_end - 1) len -= last_cigar_clip;
            // if ((i == sd.cigar_start || i == sd.cigar_end - 1) && c == 'D') // 跳过开头的deletion
            if (i == cigar_start && c == 'D') {  // 跳过开头的deletion
                c = 'H';
                ref_offset += len;
                start_pos += len; // 更新起始位置
            } else if (i == cigar_end - 1 && c == 'D') { // 跳过结尾的deletion
                c = 'H';
                softEnd -= len; // 更新结束位置
            }
            cigars.push_back({c, len});
        }
    }

    // 给定一个ref pos，返回对应的read index和cigar操作
    ReadIdxCigar getReadIndexForReferenceCoordinate(int64_t refPos) {
        ReadIdxCigar rc;
        if (refPos < start_pos)
            return rc;
        int firstReadPosOfElement = 0;         // inclusive
        int firstRefPosOfElement = start_pos;  // inclusive
        int lastReadPosOfElement = 0;          // exclusive
        int lastRefPosOfElement = start_pos;   // exclusive
        // advance forward through all the cigar elements until we bracket the reference coordinate
        for (auto& cigar : cigars) {
            firstReadPosOfElement = lastReadPosOfElement;
            firstRefPosOfElement = lastRefPosOfElement;
            lastReadPosOfElement += Cigar::ConsumeReadBases(cigar.op) ? cigar.len : 0;
            lastRefPosOfElement += Cigar::ConsumeRefBases(cigar.op) || cigar.op == 'S' ? cigar.len : 0;
            if (firstRefPosOfElement <= refPos && refPos < lastRefPosOfElement) {  // refCoord falls within this cigar element
                int readPosAtRefCoord = firstReadPosOfElement + (Cigar::ConsumeReadBases(cigar.op) ? (refPos - firstRefPosOfElement) : 0);
                rc.cigarOp = cigar.op;
                rc.readIdx = readPosAtRefCoord;
                return rc;
            }
        }
        return rc;
    }
};