linkrs-storage 0.1.0

use linkrs_core::value::{DateTimeValue, DateValue, TimeValue};
use linkrs_core::{DataType, StorageError, StorageResult, Value};

use super::{ensure_bitmap_len, ColumnStorage};
use bitvec::prelude::*;

/// Maximum `FixedString(n)` byte width stored inline in this column.
///
/// Short fixed strings live here as zero-padded `n`-byte slots; wider or
/// zero-width declarations stay on the variable-width base where the
/// length-prefixed layout and dictionary encoding fit better.
pub const FIXED_STRING_INLINE_LIMIT: usize = 32;

/// Maximum `VectorDense(n)` dimension stored as a fixed `n * 4`-byte slot
/// in this column. Small dense vectors ride the fixed base (no per-row
/// offsets or length prefix); wider or unsized vectors stay on the
/// variable-width base. The bound keeps one chunk window bounded
/// (64 dims = 256 bytes per row).
pub const VECTOR_DENSE_FIXED_MAX_DIM: usize = 64;

/// Column storage for fixed-width (primitive) types.
///
/// Values are stored in a flat `Vec<u8>` with direct offset calculation:
/// `offset = row_idx * element_size`.
/// This provides O(1) random access without any branching on type.
#[derive(Debug, Clone)]
pub struct FixedWidthColumn {
    pub(super) data: Vec<u8>,
    pub(super) data_type: DataType,
    pub(super) element_size: usize,
    pub(super) null_bitmap: Option<BitVec<u8, Lsb0>>,
    pub(super) row_count: usize,
    /// O(1) count of null rows, maintained incrementally on set/resize/clear
    /// so `used_memory_size` does not rescan the bitmap (which is O(n)).
    pub(super) null_count: usize,
}

impl FixedWidthColumn {
    pub fn new(data_type: DataType, nullable: bool) -> Self {
        let elem_size = element_size(&data_type);
        Self {
            data: Vec::new(),
            data_type: data_type.clone(),
            element_size: elem_size,
            null_bitmap: if nullable { Some(BitVec::new()) } else { None },
            row_count: 0,
            null_count: 0,
        }
    }
}

impl ColumnStorage for FixedWidthColumn {
    fn get(&self, row_idx: usize) -> Option<Value> {
        match self.try_get(row_idx) {
            Ok(value) => value,
            Err(e) => {
                log::warn!(
                    "fixed column row {} lenient read failed: {}; reading as missing",
                    row_idx,
                    e
                );
                None
            }
        }
    }

    fn try_get(&self, row_idx: usize) -> StorageResult<Option<Value>> {
        if self.is_null(row_idx) {
            return Ok(None);
        }
        if row_idx >= self.row_count {
            return Ok(None);
        }
        if let DataType::FixedString(limit) = &self.data_type {
            return try_read_fixed_string(&self.data, row_idx * self.element_size, *limit);
        }
        if let DataType::VectorDense(dim) = &self.data_type {
            return try_read_fixed_vector(&self.data, row_idx * self.element_size, *dim);
        }
        let offset = row_idx * self.element_size;
        if offset + self.element_size > self.data.len() {
            return Err(StorageError::deserialize_error(format!(
                "fixed column truncated at row {}: offset {} + size {} exceeds {} bytes",
                row_idx,
                offset,
                self.element_size,
                self.data.len()
            )));
        }
        let Some(raw) = read_fixed_value(&self.data, offset, self.element_size) else {
            return Err(StorageError::deserialize_error(format!(
                "fixed column undecodable at row {}: element size {}",
                row_idx, self.element_size
            )));
        };
        Ok(Some(convert_to_type(raw, &self.data_type)))
    }

    fn set(&mut self, row_idx: usize, value: Option<&Value>) -> StorageResult<()> {
        let was_null = self
            .null_bitmap
            .as_ref()
            .map(|b| row_idx < b.len() && b[row_idx])
            .unwrap_or(false);

        // Fixed-string bounds are checked before growing the buffer so a
        // rejected write leaves no zeroed slot behind (same observable
        // behavior as the variable-width path).
        if let DataType::FixedString(limit) = &self.data_type {
            if let Some(v) = value {
                let byte_len = match v {
                    Value::FixedString(s) => Some(s.len()),
                    Value::String(s) => Some(s.len()),
                    _ => None,
                };
                match byte_len {
                    Some(len) if len > *limit => {
                        return Err(StorageError::invalid_input(format!(
                            "FixedString({}) cannot hold {} bytes",
                            limit, len
                        )));
                    }
                    Some(_) => {}
                    None => {
                        return Err(StorageError::type_mismatch(v.data_type(), v.data_type()));
                    }
                }
            }
        }

        // Fixed-dense-vector dimension is checked before growing the buffer
        // so a rejected write leaves no zeroed slot behind.
        if let DataType::VectorDense(dim) = &self.data_type {
            if let Some(v) = value {
                match v {
                    Value::Vector(vec) => {
                        let actual = vec.dimension();
                        if actual != *dim {
                            return Err(StorageError::invalid_input(format!(
                                "VectorDense({}) cannot hold dimension {}",
                                dim, actual
                            )));
                        }
                    }
                    _ => {
                        return Err(StorageError::type_mismatch(v.data_type(), v.data_type()));
                    }
                }
            }
        }

        let offset = row_idx * self.element_size;
        if offset + self.element_size > self.data.len() {
            self.data.resize(offset + self.element_size, 0);
        }

        let old_count = self.row_count;
        let mut new_gaps = 0usize;
        match value {
            Some(v) => {
                if let DataType::FixedString(limit) = &self.data_type {
                    write_fixed_string(&mut self.data, offset, *limit, v)?;
                } else if let DataType::VectorDense(dim) = &self.data_type {
                    write_fixed_vector(&mut self.data, offset, *dim, v)?;
                } else {
                    write_fixed_value(&mut self.data, offset, self.element_size, v)?;
                }
                if let Some(ref mut bitmap) = self.null_bitmap {
                    ensure_bitmap_len(bitmap, row_idx + 1);
                    for gap in old_count..row_idx {
                        if !bitmap[gap] {
                            bitmap.set(gap, true);
                            new_gaps += 1;
                        }
                    }
                    bitmap.set(row_idx, false);
                }
            }
            None => {
                if let Some(ref mut bitmap) = self.null_bitmap {
                    ensure_bitmap_len(bitmap, row_idx + 1);
                    for gap in old_count..row_idx {
                        if !bitmap[gap] {
                            bitmap.set(gap, true);
                            new_gaps += 1;
                        }
                    }
                    bitmap.set(row_idx, true);
                }
            }
        }

        if self.null_bitmap.is_some() {
            self.null_count += new_gaps;
            match value {
                Some(v) if !v.is_null() => {
                    if was_null {
                        self.null_count = self.null_count.saturating_sub(1);
                    }
                }
                _ => {
                    // A null write either extends the row count (this row was
                    // not counted before) or replaces a previously non-null
                    // slot; a row already null stays counted once.
                    if row_idx >= old_count || !was_null {
                        self.null_count += 1;
                    }
                }
            }
        }

        if row_idx >= self.row_count {
            self.row_count = row_idx + 1;
        }

        Ok(())
    }

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

    fn is_null(&self, row_idx: usize) -> bool {
        self.null_bitmap
            .as_ref()
            .map(|b| row_idx < b.len() && b[row_idx])
            .unwrap_or(false)
    }

    fn reserve(&mut self, additional: usize) {
        self.data.reserve(additional * self.element_size.max(1));
        if let Some(ref mut bitmap) = self.null_bitmap {
            bitmap.reserve(additional);
        }
    }

    fn memory_usage(&self) -> usize {
        let data_size = self.data.len();
        let bitmap_size = self
            .null_bitmap
            .as_ref()
            .map(|b| b.as_raw_slice().len())
            .unwrap_or(0);
        data_size + bitmap_size
    }

    fn clear(&mut self) {
        self.data.clear();
        if let Some(ref mut bitmap) = self.null_bitmap {
            bitmap.clear();
        }
        self.row_count = 0;
        self.null_count = 0;
    }

    fn resize(&mut self, new_count: usize) {
        let old_count = self.row_count;
        self.data.resize(new_count * self.element_size, 0);
        if let Some(ref mut bitmap) = self.null_bitmap {
            bitmap.resize(new_count, false);
            for i in old_count..new_count {
                bitmap.set(i, true);
            }
        }
        if let Some(bitmap) = &self.null_bitmap {
            if new_count >= old_count {
                self.null_count += new_count - old_count;
            } else {
                // Shrink: recompute (rare; happens during load/compaction).
                self.null_count = bitmap.count_ones();
            }
        }
        self.row_count = new_count;
    }

    fn null_bitmap(&self) -> Option<&BitVec<u8, Lsb0>> {
        self.null_bitmap.as_ref()
    }

    fn null_count(&self) -> usize {
        self.null_count
    }

    fn load_data_from_raw(
        &mut self,
        data: Vec<u8>,
        _offsets: Vec<u32>,
        null_bitmap_raw: Option<Vec<u8>>,
        bitmap_bit_len: usize,
    ) {
        self.data = data;
        let elem_size = self.element_size.max(1);
        let remainder = self.data.len() % elem_size;
        if remainder != 0 {
            self.data
                .resize(self.data.len() + (elem_size - remainder), 0);
        }
        self.null_bitmap = null_bitmap_raw.map(|raw| {
            let mut bv = BitVec::from_vec(raw);
            bv.resize(bitmap_bit_len, false);
            bv
        });
        self.null_count = self
            .null_bitmap
            .as_ref()
            .map(|b| b.count_ones())
            .unwrap_or(0);
        self.row_count = self.data.len() / elem_size;
    }

    fn get_flush_data(&self) -> (Vec<u8>, Vec<u32>, Option<BitVec<u8, Lsb0>>) {
        (self.data.clone(), Vec::new(), self.null_bitmap.clone())
    }
}

/// Returns the element size for fixed-width data types.
/// Returns 0 for variable-length types.
pub fn element_size(data_type: &DataType) -> usize {
    match data_type {
        DataType::Bool => 1,
        DataType::SmallInt => 2,
        DataType::Int => 4,
        DataType::BigInt => 8,
        DataType::Float => 4,
        DataType::Double => 8,
        DataType::Date => 12,
        DataType::Time => 8,
        DataType::DateTime => 28,
        DataType::Uuid => 16,
        DataType::FixedString(n) if *n >= 1 && *n <= FIXED_STRING_INLINE_LIMIT => *n,
        DataType::VectorDense(dim) if *dim >= 1 && *dim <= VECTOR_DENSE_FIXED_MAX_DIM => {
            *dim * std::mem::size_of::<f32>()
        }
        _ => 0,
    }
}

/// Strict fixed-vector read: dimension and length mismatches are errors.
pub(crate) fn try_read_fixed_vector(
    data: &[u8],
    offset: usize,
    dim: usize,
) -> StorageResult<Option<Value>> {
    if dim == 0 {
        return Err(StorageError::deserialize_error(
            "fixed vector with zero dimension".to_string(),
        ));
    }
    let Some(bytes) = dim.checked_mul(std::mem::size_of::<f32>()) else {
        return Err(StorageError::deserialize_error(format!(
            "fixed vector dimension overflow: {}",
            dim
        )));
    };
    if offset + bytes > data.len() {
        return Err(StorageError::deserialize_error(format!(
            "fixed vector truncated: offset {} + {} bytes exceeds {} bytes",
            offset,
            bytes,
            data.len()
        )));
    }
    let mut out = Vec::with_capacity(dim);
    for i in 0..dim {
        let chunk: [u8; 4] = data[offset + i * 4..offset + (i + 1) * 4]
            .try_into()
            .map_err(|_| {
                StorageError::deserialize_error(format!("fixed vector component {} undecodable", i))
            })?;
        out.push(f32::from_le_bytes(chunk));
    }
    Ok(Some(Value::Vector(linkrs_core::value::VectorValue::dense(
        out,
    ))))
}

/// Write one fixed-dense-vector slot as little-endian `f32` components.
///
/// Only dense vectors of exactly `dim` components are accepted; sparse
/// vectors and dimension mismatches are rejected with the same shape as
/// the variable-width path.
pub(crate) fn write_fixed_vector(
    data: &mut [u8],
    offset: usize,
    dim: usize,
    value: &Value,
) -> StorageResult<()> {
    let dense: &[f32] = match value {
        Value::Vector(vec) => match vec.as_dense() {
            Some(dense) => dense,
            None => {
                return Err(StorageError::invalid_input(format!(
                    "VectorDense({}) cannot hold a sparse vector",
                    dim
                )));
            }
        },
        _ => {
            return Err(StorageError::type_mismatch(
                value.data_type(),
                value.data_type(),
            ));
        }
    };
    if dense.len() != dim {
        return Err(StorageError::invalid_input(format!(
            "VectorDense({}) cannot hold dimension {}",
            dim,
            dense.len()
        )));
    }
    let bytes = dim * std::mem::size_of::<f32>();
    if offset + bytes > data.len() {
        return Err(StorageError::invalid_input(format!(
            "Column data buffer too small: offset={}, limit={}, data_len={}",
            offset,
            bytes,
            data.len()
        )));
    }
    for (i, component) in dense.iter().enumerate() {
        data[offset + i * 4..offset + (i + 1) * 4].copy_from_slice(&component.to_le_bytes());
    }
    Ok(())
}

/// Strict fixed-string read with declared-length semantics.
///
/// Never-written windows are reported by the caller before this runs.
/// Truncated slots and invalid UTF-8 payloads are storage errors carrying
/// the row window so queries fail loudly instead of observing silent nulls.
pub(crate) fn try_read_fixed_string(
    data: &[u8],
    offset: usize,
    limit: usize,
) -> StorageResult<Option<Value>> {
    if limit == 0 {
        return Err(StorageError::deserialize_error(
            "fixed string slot with zero declared length".to_string(),
        ));
    }
    if offset + limit > data.len() {
        return Err(StorageError::deserialize_error(format!(
            "fixed string truncated: offset {} + limit {} exceeds {} bytes",
            offset,
            limit,
            data.len()
        )));
    }
    let mut end = offset + limit;
    while end > offset && data[end - 1] == 0 {
        end -= 1;
    }
    let s = String::from_utf8(data[offset..end].to_vec()).map_err(|e| {
        StorageError::deserialize_error(format!("fixed string invalid UTF-8: {}", e))
    })?;
    Ok(Some(Value::FixedString(s)))
}

/// Write one zero-padded inline fixed-string slot.
///
/// Accepts `String` as well as `FixedString` because primary-key mirrors
/// materialize text ids as `Value::String` regardless of the column
/// declaration. Byte length beyond `limit` is rejected with the same
/// message as the variable-width path.
pub(crate) fn write_fixed_string(
    data: &mut [u8],
    offset: usize,
    limit: usize,
    value: &Value,
) -> StorageResult<()> {
    let bytes: &[u8] = match value {
        Value::FixedString(s) => s.as_bytes(),
        Value::String(s) => s.as_bytes(),
        _ => {
            return Err(StorageError::type_mismatch(
                value.data_type(),
                value.data_type(),
            ));
        }
    };
    if bytes.len() > limit {
        return Err(StorageError::invalid_input(format!(
            "FixedString({}) cannot hold {} bytes",
            limit,
            bytes.len()
        )));
    }
    // Zero-padded slots cannot round-trip a trailing zero byte: the reader
    // strips padding, so such a value would lose bytes silently. Reject
    // loudly instead of storing a corrupted value.
    if bytes.last().is_some_and(|b| *b == 0) {
        return Err(StorageError::invalid_input(format!(
            "FixedString({}) cannot hold a value with trailing zero byte",
            limit
        )));
    }
    if offset + limit > data.len() {
        return Err(StorageError::invalid_input(format!(
            "Column data buffer too small: offset={}, limit={}, data_len={}",
            offset,
            limit,
            data.len()
        )));
    }
    data[offset..offset + bytes.len()].copy_from_slice(bytes);
    data[offset + bytes.len()..offset + limit].fill(0);
    Ok(())
}

pub(crate) fn write_fixed_value(
    data: &mut [u8],
    offset: usize,
    element_size: usize,
    value: &Value,
) -> StorageResult<()> {
    let required_size = match value {
        Value::Bool(_) => 1,
        Value::SmallInt(_) => 2,
        Value::Int(_) => 4,
        Value::BigInt(_) => 8,
        Value::Float(_) => 4,
        Value::Double(_) => 8,
        Value::Date(_) => 12,
        Value::Time(_) => 8,
        Value::DateTime(_) => 28,
        Value::Uuid(_) => 16,
        _ => {
            return Err(StorageError::type_mismatch(
                value.data_type(),
                value.data_type(),
            ));
        }
    };

    if offset + required_size > data.len() {
        return Err(StorageError::invalid_input(format!(
            "Column data buffer too small: offset={}, required_size={}, data_len={}, element_size={}",
            offset,
            required_size,
            data.len(),
            element_size
        )));
    }

    match value {
        Value::Bool(b) => {
            data[offset] = if *b { 1 } else { 0 };
        }
        Value::SmallInt(i) => {
            data[offset..offset + 2].copy_from_slice(&i.to_le_bytes());
        }
        Value::Int(i) => {
            data[offset..offset + 4].copy_from_slice(&i.to_le_bytes());
        }
        Value::BigInt(i) => {
            data[offset..offset + 8].copy_from_slice(&i.to_le_bytes());
        }
        Value::Float(f) => {
            data[offset..offset + 4].copy_from_slice(&f.to_le_bytes());
        }
        Value::Double(d) => {
            data[offset..offset + 8].copy_from_slice(&d.to_le_bytes());
        }
        Value::Date(d) => {
            data[offset..offset + 4].copy_from_slice(&d.year.to_le_bytes());
            data[offset + 4..offset + 8].copy_from_slice(&d.month.to_le_bytes());
            data[offset + 8..offset + 12].copy_from_slice(&d.day.to_le_bytes());
        }
        Value::Time(t) => {
            let micros = t.hour as u64 * 3_600_000_000
                + t.minute as u64 * 60_000_000
                + t.sec as u64 * 1_000_000
                + t.microsec as u64;
            data[offset..offset + 8].copy_from_slice(&micros.to_le_bytes());
        }
        Value::DateTime(dt) => {
            data[offset..offset + 4].copy_from_slice(&dt.year.to_le_bytes());
            data[offset + 4..offset + 8].copy_from_slice(&dt.month.to_le_bytes());
            data[offset + 8..offset + 12].copy_from_slice(&dt.day.to_le_bytes());
            data[offset + 12..offset + 16].copy_from_slice(&dt.hour.to_le_bytes());
            data[offset + 16..offset + 20].copy_from_slice(&dt.minute.to_le_bytes());
            data[offset + 20..offset + 24].copy_from_slice(&dt.sec.to_le_bytes());
            data[offset + 24..offset + 28].copy_from_slice(&dt.microsec.to_le_bytes());
        }
        Value::Uuid(u) => {
            data[offset..offset + 16].copy_from_slice(u.as_bytes());
        }
        _ => {
            return Err(StorageError::type_mismatch(
                value.data_type(),
                value.data_type(),
            ));
        }
    }
    Ok(())
}

pub(crate) fn read_fixed_value(data: &[u8], offset: usize, element_size: usize) -> Option<Value> {
    if offset + element_size > data.len() {
        return None;
    }

    match element_size {
        1 => Some(Value::Bool(data[offset] != 0)),
        2 => {
            let bytes: [u8; 2] = data[offset..offset + 2].try_into().ok()?;
            Some(Value::SmallInt(i16::from_le_bytes(bytes)))
        }
        4 => {
            let bytes: [u8; 4] = data[offset..offset + 4].try_into().ok()?;
            Some(Value::Int(i32::from_le_bytes(bytes)))
        }
        8 => {
            let bytes: [u8; 8] = data[offset..offset + 8].try_into().ok()?;
            Some(Value::BigInt(i64::from_le_bytes(bytes)))
        }
        12 => {
            let year_bytes: [u8; 4] = data[offset..offset + 4].try_into().ok()?;
            let month_bytes: [u8; 4] = data[offset + 4..offset + 8].try_into().ok()?;
            let day_bytes: [u8; 4] = data[offset + 8..offset + 12].try_into().ok()?;
            Some(Value::Date(DateValue {
                year: i32::from_le_bytes(year_bytes),
                month: u32::from_le_bytes(month_bytes),
                day: u32::from_le_bytes(day_bytes),
            }))
        }
        28 => {
            let year_bytes: [u8; 4] = data[offset..offset + 4].try_into().ok()?;
            let month_bytes: [u8; 4] = data[offset + 4..offset + 8].try_into().ok()?;
            let day_bytes: [u8; 4] = data[offset + 8..offset + 12].try_into().ok()?;
            let hour_bytes: [u8; 4] = data[offset + 12..offset + 16].try_into().ok()?;
            let minute_bytes: [u8; 4] = data[offset + 16..offset + 20].try_into().ok()?;
            let sec_bytes: [u8; 4] = data[offset + 20..offset + 24].try_into().ok()?;
            let microsec_bytes: [u8; 4] = data[offset + 24..offset + 28].try_into().ok()?;
            Some(Value::DateTime(DateTimeValue {
                year: i32::from_le_bytes(year_bytes),
                month: u32::from_le_bytes(month_bytes),
                day: u32::from_le_bytes(day_bytes),
                hour: u32::from_le_bytes(hour_bytes),
                minute: u32::from_le_bytes(minute_bytes),
                sec: u32::from_le_bytes(sec_bytes),
                microsec: u32::from_le_bytes(microsec_bytes),
            }))
        }
        16 => {
            let bytes: [u8; 16] = data[offset..offset + 16].try_into().ok()?;
            Some(Value::Uuid(linkrs_core::value::UuidValue::from_bytes(
                bytes,
            )))
        }
        _ => None,
    }
}

/// Convert a raw read_fixed_value result to the correct Value variant based on the declared DataType.
/// This handles ambiguous element sizes where multiple types share the same width.
pub(crate) fn convert_to_type(raw: Value, data_type: &DataType) -> Value {
    match (data_type, &raw) {
        (DataType::Double, Value::BigInt(n)) => Value::Double(f64::from_bits(*n as u64)),
        (DataType::Float, Value::Int(n)) => Value::Float(f32::from_bits(*n as u32)),
        (DataType::Float, Value::BigInt(n)) => Value::Float(f32::from_bits(*n as u32)),
        (DataType::Time, Value::BigInt(n)) => {
            let micros = *n as u64;
            let hour = (micros / 3_600_000_000) as u32;
            let rem = micros % 3_600_000_000;
            let minute = (rem / 60_000_000) as u32;
            let rem = rem % 60_000_000;
            let sec = (rem / 1_000_000) as u32;
            let microsec = (rem % 1_000_000) as u32;
            Value::Time(TimeValue {
                hour,
                minute,
                sec,
                microsec,
            })
        }
        _ => raw,
    }
}