use std::{convert::Infallible, sync::Arc};

use brk_types::{Height, StoredU16, StoredU64};
use parking_lot::RwLock;
use vecdb::{
    AnyVec, CachedReadableVec, PrintableIndex, ReadableVec, TypedVec, VecIndex, Version,
    short_type_name,
};

const CHECKPOINT_INTERVAL: usize = 256;

pub struct CachedBlockCountReader {
    block: Box<dyn CachedReadableVec<Height, StoredU16>>,
    checkpoints: Arc<RwLock<Checkpoints>>,
}

struct Checkpoints {
    block: Arc<[StoredU16]>,
    cumulative: Arc<[u64]>,
}

impl CachedBlockCountReader {
    pub(crate) fn new(block: Box<dyn CachedReadableVec<Height, StoredU16>>) -> Self {
        Self {
            block,
            checkpoints: Arc::new(RwLock::new(Checkpoints {
                block: Arc::from([]),
                cumulative: Arc::from([0]),
            })),
        }
    }

    pub fn cumulative_at(&self, index: usize) -> Option<StoredU64> {
        let (block, checkpoints) = self.snapshot();
        (index < block.len())
            .then(|| StoredU64::from(Self::sum_before(&block, &checkpoints, index + 1)))
    }

    pub fn for_each_cumulative(&self, from: usize, to: usize, mut each: impl FnMut(StoredU64)) {
        self.try_fold_cumulative(from, to, (), |(), value| {
            each(value);
            Ok::<_, Infallible>(())
        })
        .unwrap();
    }

    pub fn for_each_rolling_sum(
        &self,
        from: usize,
        starts: &[Height],
        mut each: impl FnMut(StoredU64),
    ) {
        let (block, checkpoints) = self.snapshot();
        let to = (from + starts.len()).min(block.len());
        if from >= to {
            return;
        }

        let starts = &starts[..to - from];
        let mut start = starts[0].to_usize();
        let mut cumulative = Self::sum_before(&block, &checkpoints, from);
        let mut before_start = Self::sum_before(&block, &checkpoints, start);

        for (offset, next_start) in starts.iter().enumerate() {
            let next_start = next_start.to_usize();
            debug_assert!(next_start >= start);
            debug_assert!(next_start <= from + offset);

            if next_start >= start {
                for value in &block[start..next_start] {
                    before_start += Self::as_u64(value);
                }
            } else {
                before_start = Self::sum_before(&block, &checkpoints, next_start);
            }
            start = next_start;

            cumulative += Self::as_u64(&block[from + offset]);
            each(StoredU64::from(cumulative - before_start));
        }
    }

    pub fn rolling_sum_at(&self, start: usize, end: usize) -> Option<StoredU64> {
        let (block, checkpoints) = self.snapshot();
        if start > end || end >= block.len() {
            return None;
        }

        Some(StoredU64::from(
            Self::sum_before(&block, &checkpoints, end + 1)
                - Self::sum_before(&block, &checkpoints, start),
        ))
    }

    fn snapshot(&self) -> (Arc<[StoredU16]>, Arc<[u64]>) {
        let block = self.block.cached();

        {
            let checkpoints = self.checkpoints.read();
            if Arc::ptr_eq(&checkpoints.block, &block) {
                return (block, checkpoints.cumulative.clone());
            }
        }

        let cumulative = Self::build_checkpoints(&block);
        let mut checkpoints = self.checkpoints.write();
        if !Arc::ptr_eq(&checkpoints.block, &block) {
            checkpoints.block = block.clone();
            checkpoints.cumulative = cumulative;
        }

        (block, checkpoints.cumulative.clone())
    }

    fn try_fold_cumulative<B, E>(
        &self,
        from: usize,
        to: usize,
        init: B,
        mut fold: impl FnMut(B, StoredU64) -> Result<B, E>,
    ) -> Result<B, E> {
        let (block, checkpoints) = self.snapshot();
        let to = to.min(block.len());
        if from >= to {
            return Ok(init);
        }

        let mut accumulator = init;
        let mut cumulative = Self::sum_before(&block, &checkpoints, from);

        for value in &block[from..to] {
            cumulative += Self::as_u64(value);
            accumulator = fold(accumulator, StoredU64::from(cumulative))?;
        }

        Ok(accumulator)
    }

    fn build_checkpoints(block: &[StoredU16]) -> Arc<[u64]> {
        let mut checkpoints = Vec::with_capacity(block.len() / CHECKPOINT_INTERVAL + 1);
        let mut cumulative = 0;
        checkpoints.push(cumulative);

        for (index, value) in block.iter().enumerate() {
            cumulative += Self::as_u64(value);
            if (index + 1) % CHECKPOINT_INTERVAL == 0 {
                checkpoints.push(cumulative);
            }
        }

        checkpoints.into()
    }

    #[inline(always)]
    fn sum_before(block: &[StoredU16], checkpoints: &[u64], end: usize) -> u64 {
        let end = end.min(block.len());
        let checkpoint = end / CHECKPOINT_INTERVAL;
        let from = checkpoint * CHECKPOINT_INTERVAL;
        checkpoints[checkpoint] + block[from..end].iter().map(Self::as_u64).sum::<u64>()
    }

    #[inline(always)]
    fn as_u64(value: &StoredU16) -> u64 {
        u64::from(**value)
    }
}

impl Clone for CachedBlockCountReader {
    fn clone(&self) -> Self {
        Self {
            block: self.block.cached_boxed_clone(),
            checkpoints: self.checkpoints.clone(),
        }
    }
}

impl AnyVec for CachedBlockCountReader {
    fn version(&self) -> Version {
        self.block.version()
    }

    fn name(&self) -> &str {
        self.block.name()
    }

    fn len(&self) -> usize {
        self.block.len()
    }

    fn index_type_to_string(&self) -> &'static str {
        <Height as PrintableIndex>::to_string()
    }

    fn region_names(&self) -> Vec<String> {
        Vec::new()
    }

    fn value_type_to_size_of(&self) -> usize {
        size_of::<StoredU64>()
    }

    fn value_type_to_string(&self) -> &'static str {
        short_type_name::<StoredU64>()
    }
}

impl TypedVec for CachedBlockCountReader {
    type I = Height;
    type T = StoredU64;
}

impl ReadableVec<Height, StoredU64> for CachedBlockCountReader {
    fn read_into_at(&self, from: usize, to: usize, buf: &mut Vec<StoredU64>) {
        buf.reserve(to.saturating_sub(from));
        self.for_each_cumulative(from, to, |value| buf.push(value));
    }

    fn for_each_range_dyn_at(&self, from: usize, to: usize, each: &mut dyn FnMut(StoredU64)) {
        self.for_each_cumulative(from, to, each);
    }

    fn fold_range_at<B, F: FnMut(B, StoredU64) -> B>(
        &self,
        from: usize,
        to: usize,
        init: B,
        mut fold: F,
    ) -> B {
        self.try_fold_cumulative(from, to, init, |accumulator, value| {
            Ok::<_, Infallible>(fold(accumulator, value))
        })
        .unwrap()
    }

    fn try_fold_range_at<B, E, F: FnMut(B, StoredU64) -> Result<B, E>>(
        &self,
        from: usize,
        to: usize,
        init: B,
        fold: F,
    ) -> Result<B, E> {
        self.try_fold_cumulative(from, to, init, fold)
    }

    fn collect_one_at(&self, index: usize) -> Option<StoredU64> {
        self.cumulative_at(index)
    }

    fn read_sorted_into_at(&self, indices: &[usize], out: &mut Vec<StoredU64>) {
        let (block, checkpoints) = self.snapshot();
        out.reserve(indices.len());
        indices
            .iter()
            .take_while(|&&index| index < block.len())
            .for_each(|&index| {
                out.push(StoredU64::from(Self::sum_before(
                    &block,
                    &checkpoints,
                    index + 1,
                )));
            });
    }
}

#[cfg(test)]
mod tests {
    use brk_types::{Height, StoredU16, StoredU64, Version};
    use vecdb::{AnyStoredVec, CachedVec, Database, EagerVec, ImportableVec, PcoVec, WritableVec};

    use super::*;

    #[test]
    fn reconstructs_cumulative_and_rolling_counts() {
        let path =
            std::env::temp_dir().join(format!("brk-cached-block-count-{}", std::process::id()));
        let db = Database::open(&path).unwrap();
        let mut block: EagerVec<PcoVec<Height, StoredU16>> =
            EagerVec::forced_import(&db, "count", Version::ONE).unwrap();

        let mut expected = Vec::new();
        let mut total = 0_u64;
        for index in 0..600 {
            let value = (index % 7) as u16;
            total += u64::from(value);
            expected.push(StoredU64::from(total));
            block.push(StoredU16::new(value));
        }
        block.write().unwrap();

        let mut block = CachedVec::wrap(block);
        let count = CachedBlockCountReader::new(block.read_only_cached_boxed_clone());

        assert_eq!(count.cumulative_at(599), Some(expected[599]));

        let mut reconstructed = Vec::new();
        count.for_each_cumulative(250, 270, |value| reconstructed.push(value));
        assert_eq!(reconstructed, expected[250..270]);

        let starts = (250..270)
            .map(|index| Height::new(index - 10))
            .collect::<Vec<_>>();
        let mut rolling = Vec::new();
        count.for_each_rolling_sum(250, &starts, |value| rolling.push(value));
        let expected = (250..270)
            .map(|index| {
                let current = expected[index];
                let previous = expected[index - 11];
                current - previous
            })
            .collect::<Vec<_>>();
        assert_eq!(rolling, expected);

        block.inner.truncate_if_needed_at(0).unwrap();
        for _ in 0..600 {
            block.inner.push(StoredU16::new(1));
        }
        block.inner.write().unwrap();
        block.clear();

        assert_eq!(count.cumulative_at(599), Some(StoredU64::from(600_u64)));

        drop(count);
        drop(block);
        drop(db);
        std::fs::remove_dir_all(path).unwrap();
    }
}
