use std::{
    collections::HashMap,
    fs::{self, File},
    path::PathBuf,
    str::FromStr,
};

use anyhow::{anyhow, Context};
use chrono::{DateTime, Utc};
use glob::glob;
use log::warn;
use pandora_lib_bindings::{
    progs::cramino::{Cramino, CraminoRes},
    utils::RunBin,
};
use rayon::prelude::*;
use rust_htslib::bam::{record::Cigar, Read};
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Bam {
    pub id: String,
    pub time_point: String,
    pub reference_genome: String,
    pub bam_type: BamType,
    pub path: PathBuf,
    pub modified: DateTime<Utc>,
    pub cramino: Option<CraminoRes>,
    pub composition: Vec<(String, f64)>, // acquisition id
}

#[derive(Debug, PartialEq, Clone, Deserialize, Serialize)]
pub enum BamType {
    WGS,
    Panel(String),
    ChIP(String),
}

impl Bam {
    pub fn new(path: PathBuf) -> anyhow::Result<Self> {
        let stem = path
            .clone()
            .file_stem()
            .context(format!("Can't parse stem from {}", path.display()))?
            .to_string_lossy()
            .to_string();
        let stem: Vec<&str> = stem.split('_').collect();

        if stem.len() > 4 || stem.len() < 3 {
            return Err(anyhow!("Error in bam name: {}", path.display()));
        }

        let id = stem[0].to_string();
        let time_point = stem[1].to_string();

        let reference_genome = stem
            .last()
            .context("Can't get last from stem {stem}")?
            .to_string();

        let bam_type = if stem.len() == 4 {
            match stem[2] {
                "oncoT" => BamType::Panel("oncoT".to_string()),
                "H3K27ac" => BamType::ChIP("H3K27ac".to_string()),
                "H3K4me3" => BamType::ChIP("H3K4me3".to_string()),
                _ => return Err(anyhow!("Error in bam name: {}", path.display())),
            }
        } else {
            BamType::WGS
        };
        let file_metadata = fs::metadata(&path)?;
        let modified = file_metadata.modified()?;

        let tp_dir = path
            .parent()
            .context("Can't parse parent from: {bam_path}")?;

        let json_path = format!(
            "{}/{id}_{time_point}_{reference_genome}_info.json",
            tp_dir.to_string_lossy()
        );
        let json_file = PathBuf::from(&json_path);

        if json_file.exists() && json_file.metadata()?.modified()? > modified {
            return Self::load_json(&json_path);
        }

        let cramino_path = format!(
            "{}/{id}_{time_point}_{reference_genome}_cramino.txt",
            tp_dir.to_string_lossy()
        );

        let cramino = if bam_type == BamType::WGS {
            if !PathBuf::from_str(&cramino_path)?.exists() {
                return Err(anyhow!("Cramino file missing {cramino_path}"));
            }
            let mut cramino = Cramino::default().with_result_path(&cramino_path);
            cramino
                .parse_results()
                .context(format!("Error while parsing cramino for {cramino_path}"))?;

            if let Some(cramino) = cramino.results {
                Some(cramino)
            } else {
                return Err(anyhow!("Cramino results parsing failed"));
            }
        } else {
            None
        };

        let composition = bam_compo(path.to_string_lossy().as_ref(), 20000).context(format!(
            "Error while reading BAM composition for {}",
            path.display()
        ))?;

        let s = Self {
            path,
            cramino,
            id: id.to_string(),
            time_point: time_point.to_string(),
            bam_type,
            reference_genome,
            composition,
            modified: modified.into(),
        };
        s.save_json(&json_path)?;
        Ok(s)
    }

    pub fn load_json(path: &str) -> anyhow::Result<Self> {
        let f = File::open(path)?;
        let s: Self = serde_json::from_reader(f)?;
        Ok(s)
    }

    pub fn save_json(&self, path: &str) -> anyhow::Result<()> {
        let f = File::create(path)?;
        serde_json::to_writer(f, self)?;
        Ok(())
    }
}

#[derive(Debug)]
pub struct BamCollection {
    pub bams: Vec<Bam>,
}

impl BamCollection {
    pub fn new(result_dir: &str) -> Self {
        load_bam_collection(result_dir)
    }

    pub fn by_acquisition_id(&self) -> HashMap<String, Vec<&Bam>> {
        let mut acq: HashMap<String, Vec<&Bam>> = HashMap::new();
        for bam in self.bams.iter() {
            for (acq_id, _) in bam.composition.iter() {
                if let Some(entry) = acq.get_mut(acq_id) {
                    entry.push(bam);
                } else {
                    acq.insert(acq_id.to_string(), vec![bam]);
                }
            }
        }
        acq
    }

    pub fn get(&self, id: &str, time_point: &str) -> Vec<&Bam> {
        self.bams
            .iter()
            .filter(|b| b.id == id && b.time_point == time_point)
            .collect()
    }

    pub fn by_id_completed(&self, min_diag_cov: f32, min_mrd_cov: f32) -> Vec<Bam> {
        self.bams
            .iter()
            .filter(|b| matches!(b.bam_type, BamType::WGS))
            .filter(|b| match &b.cramino {
                Some(cramino) => match b.time_point.as_str() {
                    "diag" => cramino.mean_length >= min_diag_cov as f64,
                    "mrd" => cramino.mean_length >= min_mrd_cov as f64,
                    _ => false,
                },
                _ => false,
            })
            .cloned()
            .collect()
    }

    pub fn ids(&self) -> Vec<String> {
        let mut ids: Vec<String> = self.bams.iter().map(|b| b.id.clone()).collect();
        ids.sort();
        ids.dedup();
        ids
    }
}

pub fn load_bam_collection(result_dir: &str) -> BamCollection {
    let pattern = format!("{}/*/*/*.bam", result_dir);
    let bams = glob(&pattern)
        .expect("Failed to read glob pattern")
        .par_bridge()
        .filter_map(|entry| {
            match entry {
                Ok(path) => match Bam::new(path) {
                    Ok(bam) => return Some(bam),
                    Err(e) => warn!("{e}"),
                },
                Err(e) => warn!("Error: {:?}", e),
            }
            None
        })
        .collect();

    BamCollection { bams }
}

pub fn bam_compo(file_path: &str, sample_size: usize) -> anyhow::Result<Vec<(String, f64)>> {
    let mut bam = rust_htslib::bam::Reader::from_path(file_path)?;

    let mut rgs: HashMap<String, u64> = HashMap::new();
    for result in bam.records().filter_map(Result::ok).take(sample_size) {
        if let rust_htslib::bam::record::Aux::String(s) = result.aux(b"RG")? {
            *rgs.entry(s.to_string()).or_default() += 1;
        }
    }

    Ok(rgs
        .into_iter()
        .map(|(rg, n)| (rg.to_string(), n as f64 * 100.0 / sample_size as f64))
        .map(|(rg, p)| {
            let (a, _) = rg.split_once('_').unwrap();
            (a.to_string(), p)
        })
        .collect())
}

// 0-based inclusif
pub fn nt_pileup(
    bam: &mut rust_htslib::bam::IndexedReader,
    chr: &str,
    position: u32,
    with_next_ins: bool,
) -> anyhow::Result<Vec<u8>> {
    use rust_htslib::{bam, bam::Read};
    let mut bases = Vec::new();
    bam.fetch((chr, position, position + 1))?;
    let mut bam_pileup = Vec::new();
    for p in bam.pileup() {
        let pileup = p.context(format!(
            "Can't pileup bam at position {}:{} (0-based)",
            chr, position
        ))?;
        let cur_position = pileup.pos();
        if cur_position == position {
            for alignment in pileup.alignments() {
                match alignment.indel() {
                    bam::pileup::Indel::Ins(_len) => bam_pileup.push(b'I'),
                    bam::pileup::Indel::Del(_len) => bam_pileup.push(b'D'),
                    _ => {
                        let record = alignment.record();
                        if record.seq_len() > 0 {
                            if let Some(b) = base_at(&record, position, with_next_ins)? {
                                bases.push(b);
                            }
                        } else if alignment.is_del() {
                            bases.push(b'D');
                        }
                    }
                }
            }
        }
    }
    Ok(bases)
}

pub fn base_at(
    record: &rust_htslib::bam::record::Record,
    at_pos: u32,
    with_next_ins: bool,
) -> anyhow::Result<Option<u8>> {
    let cigar = record.cigar();
    let seq = record.seq();
    let pos = cigar.pos() as u32;

    let mut read_i = 0u32;
    // let at_pos = at_pos - 1;
    let mut ref_pos = pos;
    if ref_pos > at_pos {
        return Ok(None);
    }

    for (id, op) in cigar.iter().enumerate() {
        let (add_read, add_ref) = match *op {
            Cigar::Match(len) | Cigar::Equal(len) | Cigar::Diff(len) => (len, len),
            Cigar::Ins(len) => (len, 0),
            Cigar::Del(len) => (0, len),
            Cigar::RefSkip(len) => (0, len),
            Cigar::SoftClip(len) => (len, 0),
            Cigar::HardClip(_) | Cigar::Pad(_) => (0, 0),
        };
        // If at the end of the op len and next op is Ins return I
        if with_next_ins && ref_pos + add_read == at_pos + 1 {
            if let Some(Cigar::Ins(_)) = cigar.get(id + 1) {
                return Ok(Some(b'I'));
            }
        }

        if ref_pos + add_ref > at_pos {
            // Handle deletions directly
            if let Cigar::Del(_) = *op {
                return Ok(Some(b'D'));
            } else if let Cigar::RefSkip(_) = op {
                return Ok(None);
            } else {
                let diff = at_pos - ref_pos;
                let p = read_i + diff;
                return Ok(Some(seq[p as usize]));
            }
        }

        read_i += add_read;
        ref_pos += add_ref;
    }
    Ok(None)
}

pub fn counts_at(
    bam: &mut rust_htslib::bam::IndexedReader,
    chr: &str,
    position: u32,
) -> anyhow::Result<HashMap<String, i32>> {
    let p = nt_pileup(bam, chr, position, false)?
        .iter()
        .map(|e| String::from_utf8(vec![*e]).unwrap())
        .collect::<Vec<_>>();
    let mut counts = HashMap::new();
    for item in p.iter() {
        *counts.entry(item.to_string()).or_insert(0) += 1;
    }
    Ok(counts)
}

pub fn ins_pileup(
    bam: &mut rust_htslib::bam::IndexedReader,
    chr: &str,
    position: u32,
) -> anyhow::Result<Vec<String>> {
    let mut bases = Vec::new();
    bam.fetch((chr, position, position + 10))?;
    for p in bam.pileup() {
        let pileup = p.context(format!(
            "Can't pileup bam at position {}:{} (0-based)",
            chr, position
        ))?;
        let cur_position = pileup.pos();
        // Ins in the next position
        if cur_position == position + 1 {
            // debug!("{cur_position}");
            for alignment in pileup.alignments() {
                let record = alignment.record();
                if record.seq_len() > 0 {
                    if let Some(b) = ins_at(&record, position)? {
                        bases.push(b);
                    }
                }
            }
        }
    }
    Ok(bases)
}

pub fn ins_at(
    record: &rust_htslib::bam::record::Record,
    at_pos: u32,
) -> anyhow::Result<Option<String>> {
    use rust_htslib::bam::record::Cigar;

    let cigar = record.cigar();
    let seq = record.seq();
    let pos = cigar.pos() as u32;

    let mut read_i = 0u32;
    let mut ref_pos = pos;
    if ref_pos > at_pos {
        return Ok(None);
    }
    // debug!(
    //     "read: {}",
    //     String::from_utf8(record.qname().to_vec()).unwrap()
    // );

    for op in cigar.iter() {
        let (add_read, add_ref) = match *op {
            Cigar::Match(len) | Cigar::Equal(len) | Cigar::Diff(len) => (len, len),
            Cigar::Ins(len) => (len, 0),
            Cigar::Del(len) => (0, len),
            Cigar::RefSkip(len) => (0, len),
            Cigar::SoftClip(len) => (len, 0),
            Cigar::HardClip(_) | Cigar::Pad(_) => (0, 0),
        };

        if ref_pos + add_read > at_pos && ref_pos + add_read < at_pos + 10 {
            if let Cigar::Ins(v) = *op {
                // debug!(
                //     "ins size {v} @ {} (corrected {})",
                //     ref_pos + add_read,
                //     (ref_pos + add_read) - v - 1
                // );

                if (ref_pos + add_read) - v - 1 == at_pos {
                    let inserted_seq =
                        seq.as_bytes()[read_i as usize..(read_i + v) as usize].to_vec();
                    return Ok(Some(String::from_utf8(inserted_seq)?));
                }
            }
        }

        read_i += add_read;
        ref_pos += add_ref;
    }
    Ok(None)
}

pub fn counts_ins_at(
    bam: &mut rust_htslib::bam::IndexedReader,
    chr: &str,
    position: u32,
) -> anyhow::Result<HashMap<String, i32>> {
    let p = ins_pileup(bam, chr, position)?;
    let mut counts = HashMap::new();
    for item in p.iter() {
        *counts.entry(item.to_string()).or_insert(0) += 1;
    }
    Ok(counts)
}
