use crate::CubeError;
use cubedatasketches::{HLLDataSketch, HLLUnionDataSketch};
use cubehll::HllSketch;
use cubezetasketch::HyperLogLogPlusPlus;

#[derive(Debug)]
pub enum Hll {
    Airlift(HllSketch),              // Compatible with Athena, Presto, etc.
    ZetaSketch(HyperLogLogPlusPlus), // Compatible with BigQuery.
    DataSketches(HLLDataSketch),     // Compatible with DataBricks
}

const DS_LIST_PREINTS: u8 = 2;
const DS_HASH_SET_PREINTS: u8 = 3;
const DS_HLL_PREINTS: u8 = 10;
const DS_SER_VER: u8 = 1;
const DS_FAMILY_ID: u8 = 7;

impl Hll {
    pub fn read(data: &[u8]) -> Result<Hll, CubeError> {
        if data.is_empty() {
            return Err(CubeError::internal(
                "invalid serialized HLL (empty data)".to_string(),
            ));
        }

        //  - must larger than 3 due to how protos are encoded in ZetaSketch.
        //  - represents the data format version and is <= 3 in AirLift.
        // -  checking first 3 bytes for figure out HLL from Apache DataSketches
        if (data[0] == DS_LIST_PREINTS
            || data[0] == DS_HASH_SET_PREINTS
            || data[0] == DS_HLL_PREINTS)
            && data[1] == DS_SER_VER
            && data[2] == DS_FAMILY_ID
        {
            return Ok(Hll::DataSketches(HLLDataSketch::read(data)?));
        } else if data[0] <= 3 {
            return Ok(Hll::Airlift(HllSketch::read(data)?));
        } else {
            return Ok(Hll::ZetaSketch(HyperLogLogPlusPlus::read(data)?));
        }
    }

    pub fn write(&self) -> Vec<u8> {
        match self {
            Self::Airlift(h) => h.write(),
            Self::ZetaSketch(h) => h.write(),
            Self::DataSketches(h) => h.write(),
        }
    }

    pub fn cardinality(&mut self) -> u64 {
        match self {
            Hll::Airlift(h) => h.cardinality(),
            Hll::ZetaSketch(h) => h.cardinality(),
            Hll::DataSketches(h) => h.cardinality(),
        }
    }
}

#[derive(Debug)]
pub enum HllUnion {
    Airlift(HllSketch),
    ZetaSketch(HyperLogLogPlusPlus),
    DataSketches(HLLUnionDataSketch),
}

impl HllUnion {
    pub fn new(hll: Hll) -> Result<Self, CubeError> {
        match hll {
            Hll::Airlift(h) => Ok(Self::Airlift(h)),
            Hll::ZetaSketch(h) => Ok(Self::ZetaSketch(h)),
            Hll::DataSketches(h) => {
                let mut union = HLLUnionDataSketch::new(h.get_lg_config_k())?;
                union.merge_with(h)?;

                Ok(Self::DataSketches(union))
            }
        }
    }

    pub fn write(&self) -> Vec<u8> {
        match self {
            Self::Airlift(h) => h.write(),
            Self::ZetaSketch(h) => h.write(),
            Self::DataSketches(h) => h.write(),
        }
    }

    pub fn is_compatible(&self, other: &Hll) -> bool {
        match (self, other) {
            (Self::Airlift(l), Hll::Airlift(r)) => l.index_bit_len() == r.index_bit_len(),
            (Self::ZetaSketch(l), Hll::ZetaSketch(r)) => l.is_compatible(r),
            (Self::DataSketches(l), Hll::DataSketches(r)) => {
                l.get_lg_config_k() == r.get_lg_config_k()
            }
            _ => return false,
        }
    }

    /// Clients are responsible for calling `is_compatible` before running this function.
    /// On error, `self` may end up in inconsistent state and must be discarded.
    pub fn merge_with(&mut self, other: Hll) -> Result<(), CubeError> {
        debug_assert!(self.is_compatible(&other));

        match (self, other) {
            (Self::Airlift(l), Hll::Airlift(r)) => l.merge_with(&r),
            (Self::ZetaSketch(l), Hll::ZetaSketch(r)) => l.merge_with(&r)?,
            (Self::DataSketches(l), Hll::DataSketches(r)) => l.merge_with(r)?,
            _ => return Err(CubeError::internal("incompatible HLL types".to_string())),
        }

        return Ok(());
    }

    /// The size of allocated memory used (not including `sizeof::<Self>()`).  Must be exact.
    pub fn allocated_size(&self) -> usize {
        match self {
            Self::Airlift(hll_sketch) => hll_sketch.allocated_size(),
            Self::ZetaSketch(hll_pp) => hll_pp.allocated_size(),
            Self::DataSketches(hll_uds) => hll_uds.allocated_size(),
        }
    }
}

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

    #[test]
    fn hll_detect_sketches_success() -> Result<(), CubeError> {
        {
            let mut sketch = Hll::read(&hex::decode(&"0201070C03080108320B1F05")?)?;
            assert!(matches!(sketch, Hll::DataSketches(_)));
            assert_eq!(sketch.cardinality(), 1);
        }

        {
            let mut sketch = Hll::read(&hex::decode(&"0301070C050800010800000064987A05A5456B06E9EBD51A4C116B08307C4E047258E51293723306176C6E06")?)?;
            assert!(matches!(sketch, Hll::DataSketches(_)));
            assert_eq!(sketch.cardinality(), 8);
        }

        {
            let mut sketch = Hll::read(&hex::decode(&" ")?)?;
            assert!(matches!(sketch, Hll::DataSketches(_)));
            assert_eq!(sketch.cardinality(), 589);
        }

        Ok(())
    }
}
