use crate::{
    annotation::Annotations,
    positions::{GenomePosition, GetGenomePosition, VcfPosition},
    runners::Run,
    variant::variant_collection::VariantCollection,
};
use anyhow::{anyhow, Context, Ok};
use rayon::prelude::*;
use serde::{Deserialize, Serialize};
use std::{cmp::Ordering, collections::HashSet, fmt, hash::Hash, str::FromStr};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VcfVariant {
    pub position: GenomePosition,
    pub id: String,
    pub reference: ReferenceAlternative,
    pub alternative: ReferenceAlternative,
    pub quality: Option<f32>,
    pub filter: Filter,
    pub infos: Infos,
    pub formats: Formats,
}

impl PartialEq for VcfVariant {
    fn eq(&self, other: &Self) -> bool {
        // Nota bene: id, filter, info, format and quality is intentionally not compared
        self.position == other.position
            && self.reference == other.reference
            && self.alternative == other.alternative
    }
}

impl Hash for VcfVariant {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.position.hash(state);
        self.reference.hash(state);
        self.alternative.hash(state);
    }
}

impl Eq for VcfVariant {}

impl FromStr for VcfVariant {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> anyhow::Result<Self> {
        let v: Vec<&str> = s.split('\t').collect();
        let vcf_position: VcfPosition = (
            *v.first().ok_or(anyhow!("Can't get contig from: {s}"))?,
            *v.get(1).ok_or(anyhow!("Can't get position from: {s}"))?,
        )
            .try_into()
            .context(format!("Can't parse position from: {s}"))?;

        let formats = if v.len() == 10 {
            (
                *v.get(8).ok_or(anyhow!("Can't parse formats from: {s}"))?,
                *v.get(9).ok_or(anyhow!("Can't parse formats from: {s}"))?,
            )
                .try_into()
                .context(format!("Can't parse formats from: {s}"))?
        } else {
            Formats::default()
        };

        Ok(Self {
            position: vcf_position.into(),
            id: v
                .get(2)
                .ok_or(anyhow!("Can't parse id from: {s}"))?
                .to_string(),
            reference: v
                .get(3)
                .ok_or(anyhow!("Can't parse reference from: {s}"))?
                .parse()
                .context(format!("Can't parse reference from: {s}"))?,
            alternative: v
                .get(4)
                .ok_or(anyhow!("Can't parse alternative from: {s}"))?
                .parse()
                .context(format!("Can't parse alternative from: {s}"))?,
            quality: v
                .get(5)
                .map(|s| s.parse::<f32>().ok()) // Try to parse as f64; returns Option<f64>
                .unwrap_or(None),
            filter: v
                .get(6)
                .ok_or(anyhow!("Can't parse filter from: {s}"))?
                .parse()
                .context(format!("Can't parse filter from: {s}"))?,
            infos: v
                .get(7)
                .ok_or(anyhow!("Can't parse infos from: {s}"))?
                .parse()
                .context(format!("Can't parse infos from: {s}"))?,
            formats,
        })
    }
}

// #CHROM  POS     ID      REF     ALT     QUAL    FILTER  INFO    FORMAT  ADJAGBA_diag
impl VcfVariant {
    pub fn into_vcf_row(&self) -> String {
        let vcf_position: VcfPosition = self.position.clone().into();
        let (contig, position) = vcf_position.into();

        let mut columns = vec![
            contig,
            position,
            self.id.to_string(),
            self.reference.to_string(),
            self.alternative.to_string(),
            self.quality
                .map(|v| v.to_string())
                .unwrap_or(".".to_string()),
            self.filter.to_string(),
            self.infos.to_string(),
        ];

        if !self.formats.0.is_empty() {
            let (format, values) = self.formats.clone().into();
            columns.push(format);
            columns.push(values);
        }

        columns.join("\t")
    }

    pub fn commun_deepvariant_clairs(&self) -> VcfVariant {
        VcfVariant {
            position: self.position.clone(),
            id: self.id.clone(),
            reference: self.reference.clone(),
            alternative: self.alternative.clone(),
            quality: self.quality,
            filter: Filter::Other(".".to_string()),
            infos: Infos(vec![Info::Empty]),
            formats: self.formats.commun_deepvariant_clairs(),
        }
    }

    pub fn hash_variant(&self) -> u128 {
        // Create a new BLAKE3 hasher
        let mut hasher = blake3::Hasher::new();
        // Update the hasher with the fields of the variant
        hasher.update(&self.position.contig.to_ne_bytes()); // Convert position to bytes
        hasher.update(&self.position.position.to_ne_bytes()); // Convert position to bytes
        hasher.update(self.reference.to_string().as_bytes()); // Reference string as bytes
        hasher.update(self.alternative.to_string().as_bytes()); // Alternative string as bytes

        // Finalize the hash and get the output
        let hash_output = hasher.finalize();

        // Convert the first 16 bytes of the hash output to a u128
        let mut array = [0u8; 16];
        array.copy_from_slice(&hash_output.as_bytes()[..16]);

        // Convert to u128
        u128::from_ne_bytes(array)
    }

    pub fn alteration_category(&self) -> AlterationCategory {
        match (&self.reference, &self.alternative) {
            (ReferenceAlternative::Nucleotide(_), ReferenceAlternative::Nucleotide(_)) => {
                AlterationCategory::Snv
            }
            (ReferenceAlternative::Nucleotide(_), ReferenceAlternative::Nucleotides(_)) => {
                AlterationCategory::Ins
            }
            (ReferenceAlternative::Nucleotide(_), ReferenceAlternative::Unstructured(_)) => {
                AlterationCategory::Other
            }
            (ReferenceAlternative::Nucleotides(_), ReferenceAlternative::Nucleotide(_)) => {
                AlterationCategory::Del
            }
            (ReferenceAlternative::Nucleotides(a), ReferenceAlternative::Nucleotides(b))
                if a.len() < b.len() =>
            {
                AlterationCategory::Ins
            }
            (ReferenceAlternative::Nucleotides(a), ReferenceAlternative::Nucleotides(b))
                if a.len() > b.len() =>
            {
                AlterationCategory::Del
            }
            (ReferenceAlternative::Nucleotides(_), ReferenceAlternative::Nucleotides(_)) => {
                AlterationCategory::Rep
            }
            (ReferenceAlternative::Nucleotides(_), ReferenceAlternative::Unstructured(_)) => {
                AlterationCategory::Other
            }
            (ReferenceAlternative::Unstructured(_), ReferenceAlternative::Nucleotide(_)) => {
                AlterationCategory::Other
            }
            (ReferenceAlternative::Unstructured(_), ReferenceAlternative::Nucleotides(_)) => {
                AlterationCategory::Other
            }
            (ReferenceAlternative::Unstructured(_), ReferenceAlternative::Unstructured(_)) => {
                AlterationCategory::Other
            }
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
pub enum AlterationCategory {
    Snv,
    Ins,
    Del,
    Rep,
    Other,
}

impl VariantId for VcfVariant {
    fn variant_id(&self) -> String {
        format!("{}_{}>{}", self.position, self.reference, self.alternative)
    }
}

impl GetGenomePosition for VcfVariant {
    fn position(&self) -> &GenomePosition {
        &self.position
    }
}

impl PartialOrd for VcfVariant {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for VcfVariant {
    fn cmp(&self, other: &Self) -> Ordering {
        self.position.cmp(&other.position)
    }
}

/// Info
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
pub struct Infos(Vec<Info>);

impl FromStr for Infos {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> anyhow::Result<Self> {
        Ok(Self(
            s.split(";")
                .map(Info::from_str)
                .collect::<Result<Vec<Info>, _>>()
                .map_err(|e| anyhow::anyhow!("Failed to parse info: {e}"))?,
        ))
    }
}

impl fmt::Display for Infos {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{}",
            self.0
                .iter()
                .map(|e| e.to_string())
                .collect::<Vec<String>>()
                .join(";")
        )
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum Info {
    Empty,
    H,
    F,
    P,
    FAU(u32),
    FCU(u32),
    FGU(u32),
    FTU(u32),
    RAU(u32),
    RCU(u32),
    RGU(u32),
    RTU(u32),
}

impl FromStr for Info {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> anyhow::Result<Self> {
        if s.contains("=") {
            let (key, value) = s
                .split_once('=')
                .context(format!("Can't split with `=` {s}"))?;
            Ok(match key {
                "FAU" => Info::FAU(
                    value
                        .parse()
                        .context(format!("Can't parse into u32: {value}"))?,
                ),
                "FCU" => Info::FCU(
                    value
                        .parse()
                        .context(format!("Can't parse into u32: {value}"))?,
                ),
                "FGU" => Info::FGU(
                    value
                        .parse()
                        .context(format!("Can't parse into u32: {value}"))?,
                ),
                "FTU" => Info::FTU(
                    value
                        .parse()
                        .context(format!("Can't parse into u32: {value}"))?,
                ),
                "RAU" => Info::RAU(
                    value
                        .parse()
                        .context(format!("Can't parse into u32: {value}"))?,
                ),
                "RCU" => Info::RCU(
                    value
                        .parse()
                        .context(format!("Can't parse into u32: {value}"))?,
                ),
                "RGU" => Info::RGU(
                    value
                        .parse()
                        .context(format!("Can't parse into u32: {value}"))?,
                ),
                "RTU" => Info::RTU(
                    value
                        .parse()
                        .context(format!("Can't parse into u32: {value}"))?,
                ),

                _ => Info::Empty,
            })
        } else {
            Ok(match s {
                "H" => Info::H,
                "F" => Info::F,
                "P" => Info::P,

                _ => Info::Empty,
            })
        }
    }
}

impl fmt::Display for Info {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Info::Empty => write!(f, "."),
            Info::H => write!(f, "H"),
            Info::F => write!(f, "F"),
            Info::P => write!(f, "P"),
            Info::FAU(v) => write!(f, "FAU={v}"),
            Info::FCU(v) => write!(f, "FCU={v}"),
            Info::FGU(v) => write!(f, "FGU={v}"),
            Info::FTU(v) => write!(f, "FTU={v}"),
            Info::RAU(v) => write!(f, "RAU={v}"),
            Info::RCU(v) => write!(f, "RCU={v}"),
            Info::RGU(v) => write!(f, "RGU={v}"),
            Info::RTU(v) => write!(f, "RTU={v}"),
        }
    }
}

/// Format
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum Format {
    // DeepVariant
    GT(String),
    GQ(u32),
    DP(u32),
    AD(Vec<u32>),
    VAF(f32),
    PL(Vec<u32>),

    // Clairs
    AF(f32),
    NAF(u32),
    NDP(u32),
    NAD(Vec<u32>),
    AU(u32),
    CU(u32),
    GU(u32),
    TU(u32),
    NAU(u32),
    NCU(u32),
    NGU(u32),
    NTU(u32),

    Other((String, String)), // (key, value)
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
pub struct Formats(Vec<Format>);

impl TryFrom<(&str, &str)> for Formats {
    type Error = anyhow::Error;

    fn try_from((k, v): (&str, &str)) -> anyhow::Result<Self> {
        let keys: Vec<&str> = k.split(':').collect();
        let values: Vec<&str> = v.split(':').collect();

        if keys.len() != values.len() {
            anyhow::bail!("Mismatch between keys and values count for {k} {v}");
        }

        Ok(Self(
            keys.into_iter()
                .zip(values)
                .map(|(key, value)| Format::try_from((key, value)))
                .collect::<Result<Vec<Format>, _>>()
                .map_err(|e| anyhow::anyhow!("Failed to parse format: {e}"))?,
        ))
    }
}

impl From<Formats> for (String, String) {
    fn from(formats: Formats) -> Self {
        let mut keys = Vec::new();
        let mut values = Vec::new();

        for format in formats.0 {
            let (key, value): (String, String) = format.into();
            keys.push(key);
            values.push(value);
        }

        (keys.join(":"), values.join(":"))
    }
}

impl TryFrom<(&str, &str)> for Format {
    type Error = anyhow::Error;

    fn try_from((key, value): (&str, &str)) -> anyhow::Result<Self> {
        Ok(match key {
            "GT" => Format::GT(value.to_string()),
            "GQ" => Format::GQ(value.parse().context(format!("Can't parse GQ: {value}"))?),
            "DP" => Format::DP(value.parse().context(format!("Can't parse DP: {value}"))?),
            "AD" => Format::AD(
                value
                    .split(',')
                    .map(|e| e.parse().context("Failed to parse AD"))
                    .collect::<anyhow::Result<Vec<_>>>()?,
            ),
            "VAF" => Format::VAF(value.parse().context(format!("Can't parse VAF: {value}"))?),
            "PL" => Format::PL(
                value
                    .split(',')
                    .map(|e| e.parse().context("Failed to parse AD"))
                    .collect::<anyhow::Result<Vec<_>>>()?,
            ),
            _ => Format::Other((key.to_string(), value.to_string())),
        })
    }
}

impl From<Format> for (String, String) {
    fn from(format: Format) -> Self {
        let concat = |values: Vec<u32>| -> String {
            values
                .iter()
                .map(|v| v.to_string())
                .collect::<Vec<_>>()
                .join(",")
        };

        match format {
            Format::GT(value) => ("GT".to_string(), value),
            Format::GQ(value) => ("GQ".to_string(), value.to_string()),
            Format::DP(value) => ("DP".to_string(), value.to_string()),
            Format::AD(values) => ("AD".to_string(), concat(values)),
            Format::VAF(value) => ("VAF".to_string(), value.to_string()),
            Format::PL(values) => ("PL".to_string(), concat(values)),
            Format::Other((key, value)) => (key, value),
            Format::AF(value) => ("AF".to_string(), value.to_string()),
            Format::NAF(value) => ("NAF".to_string(), value.to_string()),
            Format::NDP(value) => ("NDP".to_string(), value.to_string()),
            Format::NAD(values) => ("NAD".to_string(), concat(values)),
            Format::AU(value) => ("AU".to_string(), value.to_string()),
            Format::CU(value) => ("CU".to_string(), value.to_string()),
            Format::GU(value) => ("GU".to_string(), value.to_string()),
            Format::TU(value) => ("TU".to_string(), value.to_string()),
            Format::NAU(value) => ("NAU".to_string(), value.to_string()),
            Format::NCU(value) => ("NCU".to_string(), value.to_string()),
            Format::NGU(value) => ("NGU".to_string(), value.to_string()),
            Format::NTU(value) => ("NTU".to_string(), value.to_string()),
        }
    }
}

impl Formats {
    pub fn commun_deepvariant_clairs(&self) -> Self {
        let filtered_vec: Vec<Format> = self
            .0
            .clone()
            .into_iter()
            .map(|e| {
                if let Format::VAF(v) = e {
                    Format::AF(v)
                } else {
                    e
                }
            })
            .filter(|format| {
                matches!(
                    format,
                    Format::GT(_) | Format::GQ(_) | Format::DP(_) | Format::AD(_) | Format::AF(_)
                )
            })
            .collect();

        Formats(filtered_vec)
    }
}

/// Filter
#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
pub enum Filter {
    PASS,
    Other(String),
}

impl FromStr for Filter {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> anyhow::Result<Self> {
        match s {
            "PASS" => Ok(Filter::PASS),
            _ => Ok(Filter::Other(s.to_string())),
        }
    }
}

impl fmt::Display for Filter {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Filter::PASS => write!(f, "PASS"),
            Filter::Other(ref s) => write!(f, "{}", s),
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, Hash)]
pub enum ReferenceAlternative {
    Nucleotide(Base),
    Nucleotides(Vec<Base>),
    Unstructured(String),
}

impl FromStr for ReferenceAlternative {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> anyhow::Result<Self> {
        let possible_bases = s.as_bytes().iter();
        let mut res: Vec<Base> = Vec::new();
        for &base in possible_bases {
            match base.try_into() {
                std::result::Result::Ok(b) => res.push(b),
                Err(_) => {
                    return Ok(Self::Unstructured(s.to_string()));
                }
            }
        }

        if res.len() == 1 {
            Ok(Self::Nucleotide(res.pop().unwrap()))
        } else {
            Ok(Self::Nucleotides(res))
        }
    }
}

impl fmt::Display for ReferenceAlternative {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let string = match self {
            ReferenceAlternative::Nucleotide(b) => b.to_string(),
            ReferenceAlternative::Nucleotides(bases) => bases
                .iter()
                .fold(String::new(), |acc, e| format!("{}{}", acc, e)),
            ReferenceAlternative::Unstructured(s) => s.to_string(),
        };
        write!(f, "{}", string)
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, Hash)]
pub enum Base {
    A,
    T,
    C,
    G,
    N,
}

impl TryFrom<u8> for Base {
    type Error = anyhow::Error;
    fn try_from(base: u8) -> anyhow::Result<Self> {
        match base {
            b'A' => Ok(Base::A),
            b'T' => Ok(Base::T),
            b'C' => Ok(Base::C),
            b'G' => Ok(Base::G),
            b'N' => Ok(Base::N),
            _ => Err(anyhow::anyhow!(
                "Unknown base: {}",
                String::from_utf8_lossy(&[base])
            )),
        }
    }
}

impl Base {
    pub fn into_u8(self) -> u8 {
        match self {
            Base::A => b'A',
            Base::T => b'T',
            Base::C => b'C',
            Base::G => b'G',
            Base::N => b'N',
        }
    }
}

impl fmt::Display for Base {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        // Use `self.number` to refer to each positional data point.
        let str = match self {
            Base::A => "A",
            Base::T => "T",
            Base::C => "C",
            Base::G => "G",
            Base::N => "N",
        };
        write!(f, "{}", str)
    }
}

pub trait Variants {
    fn variants(&self, annotations: &Annotations) -> anyhow::Result<VariantCollection>;
}

pub trait VariantId {
    fn variant_id(&self) -> String;
}

pub trait RunnerVariants: Run + Variants + Send + Sync {}

pub fn load_variants(
    iterable: &mut [Box<dyn RunnerVariants + Send + Sync>],
    annotations: &Annotations,
) -> anyhow::Result<Vec<VariantCollection>> {
    // First, run all items in parallel
    iterable
        .par_iter_mut()
        .try_for_each(|runner| runner.run())?;

    // Then, collect variants from all items in parallel
    let variants: Vec<VariantCollection> = iterable
        .par_iter()
        .map(|runner| runner.variants(annotations))
        .collect::<anyhow::Result<Vec<_>>>()?;

    Ok(variants)
}

pub fn parallel_intersection<T: Hash + Eq + Clone + Send + Sync>(
    vec1: &[T],
    vec2: &[T],
) -> (Vec<T>, Vec<T>, Vec<T>) {
    let set1: HashSet<_> = vec1.par_iter().cloned().collect();
    let set2: HashSet<_> = vec2.par_iter().cloned().collect();

    let common: Vec<T> = set1
        .par_iter()
        .filter(|item| set2.contains(item))
        .cloned()
        .collect();

    let only_in_first: Vec<T> = set1
        .par_iter()
        .filter(|item| !set2.contains(item))
        .cloned()
        .collect();

    let only_in_second: Vec<T> = set2
        .par_iter()
        .filter(|item| !set1.contains(item))
        .cloned()
        .collect();

    (common, only_in_first, only_in_second)
}
