md_analysis 0.2.1

use std::vec;

use rust_htslib::bam::{Record, ext::BamRecordExtensions};

/// 流水线中间消息类型
///
/// RecordData：从 BAM record 序列化而来，供 stats workers 使用。
/// 解决 rust_htslib::bam::Record 不 Send 的问题——所有数据转为纯值类型（String、Vec<u8>、Vec<u64>）。
use crate::{bam_record_ext::BamRecordExt, range::Range};

const COMPLEMENT_TABLE: [u8; 256] = {
    let mut table = [0u8; 256];
    // 先初始化为自身（即每个字符映射到自身）
    let mut i = 0;
    while i < 256 {
        table[i] = i as u8;
        i += 1;
    }
    // 设置互补映射
    table[b'A' as usize] = b'T';
    table[b'T' as usize] = b'A';
    table[b'C' as usize] = b'G';
    table[b'G' as usize] = b'C';
    table[b'a' as usize] = b't';
    table[b't' as usize] = b'a';
    table[b'c' as usize] = b'g';
    table[b'g' as usize] = b'c';
    table
};

/// 从 BAM record 提取的记录数据，可跨线程发送。
///
/// `bases` / `dw_bins` / `ar_bins` 的每个子 Vec 对应一个 region（first_n、last_n）。
/// slice_n == 0 时只有一个 "full" 元素包含全部位置。
#[derive(Clone, Debug)]
pub struct RecordData {
    #[allow(unused)]
    pub name: String,
    pub seq_len: usize,
    pub bases: String,
    pub dw: Vec<u32>,
    pub ar: Vec<u32>,
}

impl From<&Record> for RecordData {
    fn from(value: &Record) -> Self {
        let record_ext = BamRecordExt::new(value);
        let name = record_ext.get_qname();
        let seq_len = value.seq_len();
        let bases = record_ext.get_seq();
        let dw = record_ext.get_dw().unwrap();
        let ar = record_ext.get_ar().unwrap();
        Self {
            name,
            seq_len,
            bases,
            dw,
            ar,
        }
    }
}

impl RecordData {
    pub fn from_record_within_ref_range(
        value: &Record,
        range: &Range<usize>,
    ) -> Option<RecordData> {
        if value.is_unmapped() || value.is_supplementary() || value.is_secondary() {
            return None;
        }

        let mut rcursor = None;
        let is_rev = value.is_reverse();

        let record_ext = BamRecordExt::new(value);
        let name = record_ext.get_qname();
        let query_seq = record_ext.get_seq();
        let query_bytes = query_seq.as_bytes();

        let all_dw = record_ext.get_dw().unwrap();
        let all_ar = record_ext.get_ar().unwrap();

        let mut range_seq = String::new();
        let mut range_dw = vec![];
        let mut range_ar = vec![];

        let reference_align_end = record_ext.reference_end();
        let reference_align_start = record_ext.reference_start();

        for [qpos, rpos] in value.aligned_pairs_full() {
            if rpos.is_some() {
                rcursor = rpos;
            }

            if qpos.is_none() && rpos.is_none() {
                continue;
            }

            if rcursor.is_none() {
                continue;
            }

            if qpos.is_none() {
                continue;
            }

            let rcursor_ = rcursor.unwrap() as usize;

            if rcursor_ < reference_align_start {
                continue;
            }

            let qpos_ = qpos.unwrap() as usize;
            if range.within_range(rcursor_) {
                let mut ori_base = query_bytes[qpos_];
                if is_rev {
                    ori_base = COMPLEMENT_TABLE[ori_base as usize];
                }

                range_seq.push(ori_base as char);
                range_dw.push(all_dw[qpos_]);
                range_ar.push(all_ar[qpos_]);
            }

            if rcursor_ >= reference_align_end {
                break;
            }
        }
        let seq_len = range_seq.len();
        Some(Self {
            name,
            seq_len,
            bases: range_seq,
            dw: range_dw,
            ar: range_ar,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn record_data_type_exists() {
        let rd = RecordData {
            name: String::new(),
            seq_len: 0,
            bases: String::new(),
            dw: vec![],
            ar: vec![],
        };
        assert_eq!(rd.name, "");
        assert_eq!(rd.seq_len, 0);
    }

    #[test]
    fn record_data_debug_fmt() {
        let rd = RecordData {
            name: "read1".into(),
            seq_len: 100,
            bases: "ATGC".into(),
            dw: vec![1, 2],
            ar: vec![3, 4],
        };
        let s = format!("{rd:?}");
        assert!(s.contains("read1"));
        assert!(s.contains("bases"));
    }

    #[test]
    fn record_data_clone() {
        let rd = RecordData {
            name: "read2".into(),
            seq_len: 50,
            bases: "ACGT".into(),
            dw: vec![1],
            ar: vec![2],
        };
        let rd2 = rd.clone();
        assert_eq!(rd.name, rd2.name);
        assert_eq!(rd.seq_len, rd2.seq_len);
        assert_eq!(rd.bases, rd2.bases);
        assert_eq!(rd.dw, rd2.dw);
        assert_eq!(rd.ar, rd2.ar);
    }
}