liteboxfs 0.1.0

use std::{
    iter::FusedIterator,
    ops::Range,
    ops::{self, RangeBounds},
};

use bitflags::bitflags;
use rusqlite::{ToSql, types::FromSql};

use crate::hash::BlockHash;

bitflags! {
    /// A block's encoding.
    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
    pub struct BlockEncoding: u32 {
        const NONE = 0;
        const ZSTD = 1;
    }
}

impl ToSql for BlockEncoding {
    fn to_sql(&self) -> rusqlite::Result<rusqlite::types::ToSqlOutput<'_>> {
        Ok(rusqlite::types::ToSqlOutput::from(self.bits()))
    }
}

impl FromSql for BlockEncoding {
    fn column_result(value: rusqlite::types::ValueRef<'_>) -> rusqlite::types::FromSqlResult<Self> {
        let bits = value.as_i64()? as u32;
        Self::from_bits(bits).ok_or_else(|| {
            rusqlite::types::FromSqlError::Other(format!("Invalid block encoding: {bits}").into())
        })
    }
}

/// The index of a block in a file.
pub type BlockIndex = usize;

/// The byte offset from the start of a block
pub type BlockOffset = usize;

/// The byte offset from the start of a file.
pub type FileOffset = u64;

/// The length of a block.
pub type BlockLen = usize;

/// The length of a hole
pub type HoleLen = u64;

/// The database ID of a file.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct FileId(i64);

impl From<i64> for FileId {
    fn from(value: i64) -> Self {
        FileId(value)
    }
}

impl From<FileId> for i64 {
    fn from(file_id: FileId) -> Self {
        file_id.0
    }
}

#[doc(hidden)]
impl ToSql for FileId {
    fn to_sql(&self) -> rusqlite::Result<rusqlite::types::ToSqlOutput<'_>> {
        Ok(rusqlite::types::ToSqlOutput::from(self.0))
    }
}

#[doc(hidden)]
impl FromSql for FileId {
    fn column_result(value: rusqlite::types::ValueRef<'_>) -> rusqlite::types::FromSqlResult<Self> {
        Ok(FileId::from(i64::column_result(value)?))
    }
}

/// The database ID of a block.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct BlockId(i64);

impl From<i64> for BlockId {
    fn from(value: i64) -> Self {
        BlockId(value)
    }
}

impl From<BlockId> for i64 {
    fn from(block_id: BlockId) -> Self {
        block_id.0
    }
}

#[doc(hidden)]
impl ToSql for BlockId {
    fn to_sql(&self) -> rusqlite::Result<rusqlite::types::ToSqlOutput<'_>> {
        Ok(rusqlite::types::ToSqlOutput::from(self.0))
    }
}

#[doc(hidden)]
impl FromSql for BlockId {
    fn column_result(value: rusqlite::types::ValueRef<'_>) -> rusqlite::types::FromSqlResult<Self> {
        Ok(BlockId::from(i64::column_result(value)?))
    }
}

// A "block" is the fundamental unit of data storage in the system. Data is written to the database
// in blocks, which may be either data blocks or holes (sparse regions). Blocks are identified by a
// hash and deduplicated based on that hash, so a block can belong to many files. Must of the
// complexity in the system comes from splitting and merging blocks as data is written.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum BlockSignature {
    Data { hash: BlockHash, len: BlockLen },
    Hole { len: HoleLen },
}

impl BlockSignature {
    pub fn len(&self) -> BlockLen {
        match self {
            BlockSignature::Data { len, .. } => *len,
            BlockSignature::Hole { len } => *len as BlockLen,
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Block {
    pub id: BlockId,
    pub signature: BlockSignature,
}

impl Block {
    pub fn len(&self) -> BlockLen {
        self.signature.len()
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct IndexedBlock {
    pub index: BlockIndex,
    pub block: Block,
}

/// The data contained in a block.
#[derive(Debug)]
pub enum BlockData {
    Data { bytes: Vec<u8> },
    Hole { len: HoleLen },
}

/// The list of blocks that comprise a file.
///
// This type is read-only because it's the data store's responsibility to keep track of the block
// list; users of the trait will always get a new block list from the data store rather than mutate
// it themself.
#[derive(Debug, Default, Clone)]
pub struct BlockList {
    list: Vec<Block>,
    cached_total_len: Option<FileOffset>,
}

impl BlockList {
    pub fn new(blocks: Vec<Block>) -> Self {
        Self {
            list: blocks,
            cached_total_len: None,
        }
    }

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

    pub fn is_empty(&self) -> bool {
        self.list.is_empty()
    }

    /// Return the total file length (sum of all block lengths).
    pub fn file_len(&mut self) -> FileOffset {
        if let Some(total_len) = self.cached_total_len {
            total_len
        } else {
            let total_len = self.list.iter().map(|b| b.len() as FileOffset).sum();
            self.cached_total_len = Some(total_len);
            total_len
        }
    }

    /// Return true if the entire block list consists only of holes (no data blocks).
    pub fn is_entirely_hole(&self) -> bool {
        !self.list.is_empty()
            && self
                .list
                .iter()
                .all(|b| matches!(b.signature, BlockSignature::Hole { .. }))
    }

    pub fn at_offset(&self, offset: FileOffset) -> Option<(BlockOffset, Block)> {
        let mut current_offset: FileOffset = 0;

        for block in &self.list {
            let block_len = block.len() as FileOffset;

            if current_offset + block_len > offset {
                return Some(((offset - current_offset) as BlockOffset, block.clone()));
            }

            current_offset += block_len;
        }

        None
    }

    /// Return the ranges of logical block indices that are entirely within holes.
    ///
    /// This only returns ranges for holes that are fully aligned to logical block boundaries.
    fn hole_ranges(&self, logical_block_size: BlockLen) -> Vec<Range<usize>> {
        let mut ranges = Vec::new();
        let mut current_offset: FileOffset = 0;

        for block in self.iter() {
            if let BlockSignature::Hole { len } = block.signature {
                let hole_start = current_offset;
                let hole_end = current_offset + len;

                // First logical block that is entirely within the hole.
                let first_aligned = hole_start.div_ceil(logical_block_size as FileOffset);

                // Last logical block index (exclusive) that is entirely within the hole.
                let last_aligned = hole_end / logical_block_size as FileOffset;

                if first_aligned < last_aligned {
                    ranges.push(first_aligned as usize..last_aligned as usize);
                }

                current_offset = hole_end;
            } else {
                current_offset += block.len() as u64;
            }
        }

        ranges
    }

    /// Return indices of logical blocks that contain any data (not pure holes).
    ///
    /// This iterates only through data blocks and boundary regions, not through all indices.
    pub fn data_indices(
        &self,
        num_logical_blocks: usize,
        logical_block_size: BlockLen,
    ) -> Vec<usize> {
        let hole_ranges = self.hole_ranges(logical_block_size);

        // Build a list of data indices by excluding hole ranges.
        let mut data_indices = Vec::new();
        let mut next_index = 0;

        for hole_range in hole_ranges {
            // Add indices before this hole range.
            for idx in next_index..hole_range.start.min(num_logical_blocks) {
                data_indices.push(idx);
            }
            next_index = hole_range.end;
        }

        // Add any remaining indices after the last hole.
        for idx in next_index..num_logical_blocks {
            data_indices.push(idx);
        }

        data_indices
    }

    pub fn iter(&'_ self) -> BlocksIter<'_> {
        BlocksIter {
            inner: self.list.iter(),
        }
    }
}

impl From<Vec<Block>> for BlockList {
    fn from(blocks: Vec<Block>) -> Self {
        BlockList::new(blocks)
    }
}

/// An abstraction for mutating a block list in-place.
///
/// This is used internally by the data store, but is not exposed to users of the data store trait.
#[derive(Debug)]
pub struct BlockListEditor {
    blocks: Vec<Block>,
    cached_total_len: Option<FileOffset>,
}

impl BlockListEditor {
    pub fn new(list: BlockList) -> Self {
        Self {
            blocks: list.list,
            cached_total_len: list.cached_total_len,
        }
    }

    pub fn splice(
        &mut self,
        range: impl ops::RangeBounds<usize>,
        replacement: impl IntoIterator<Item = Block>,
    ) {
        let replacement = replacement.into_iter().collect::<Vec<_>>();

        if let Some(total_len) = &mut self.cached_total_len {
            let range_start = match range.start_bound() {
                ops::Bound::Included(&start) => start,
                ops::Bound::Excluded(&start) => start + 1,
                ops::Bound::Unbounded => 0,
            };

            let range_end = match range.end_bound() {
                ops::Bound::Included(&end) => end + 1,
                ops::Bound::Excluded(&end) => end,
                ops::Bound::Unbounded => self.blocks.len(),
            };

            let removed_len: FileOffset = self
                .blocks
                .get(range_start..range_end)
                .unwrap_or(&[])
                .iter()
                .map(|block| block.len() as FileOffset)
                .sum();

            let added_len: FileOffset = replacement
                .iter()
                .map(|block| block.len() as FileOffset)
                .sum();

            *total_len = *total_len - removed_len + added_len;
        }

        self.blocks.splice(range, replacement);
    }

    pub fn insert(&mut self, index: BlockIndex, block: Block) {
        if let Some(total_len) = &mut self.cached_total_len {
            *total_len += block.len() as FileOffset;
        }

        self.blocks.insert(index, block);
    }

    pub fn drain<R: RangeBounds<usize>>(&mut self, range: R) {
        if let Some(total_len) = &mut self.cached_total_len {
            let range_start = match range.start_bound() {
                ops::Bound::Included(&start) => start,
                ops::Bound::Excluded(&start) => start + 1,
                ops::Bound::Unbounded => 0,
            };

            let range_end = match range.end_bound() {
                ops::Bound::Included(&end) => end + 1,
                ops::Bound::Excluded(&end) => end,
                ops::Bound::Unbounded => self.blocks.len(),
            };

            let drained_len: FileOffset = self
                .blocks
                .get(range_start..range_end)
                .unwrap_or(&[])
                .iter()
                .map(|block| block.len() as FileOffset)
                .sum();

            *total_len -= drained_len;
        }

        self.blocks.drain(range);
    }

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

    pub fn finish(self) -> BlockList {
        BlockList {
            list: self.blocks,
            cached_total_len: self.cached_total_len,
        }
    }
}

#[derive(Debug, Clone)]
pub struct BlocksIter<'a> {
    inner: std::slice::Iter<'a, Block>,
}

impl<'a> Iterator for BlocksIter<'a> {
    type Item = &'a Block;

    fn next(&mut self) -> Option<Self::Item> {
        self.inner.next()
    }
}

impl<'a> FusedIterator for BlocksIter<'a> {}

#[derive(Debug, Clone)]
pub struct BlocksIntoIter {
    inner: std::vec::IntoIter<Block>,
}

impl Iterator for BlocksIntoIter {
    type Item = Block;

    fn next(&mut self) -> Option<Self::Item> {
        self.inner.next()
    }
}

impl FusedIterator for BlocksIntoIter {}

impl IntoIterator for BlockList {
    type Item = Block;

    type IntoIter = BlocksIntoIter;

    fn into_iter(self) -> Self::IntoIter {
        BlocksIntoIter {
            inner: self.list.into_iter(),
        }
    }
}