lean_lmdb

use crate::atoms;
use crate::crud::{CrudError, copy_bytes, map_heed_error};
use crate::lifecycle::{DatabaseResource, registry};
use rustler::Atom;
use std::collections::hash_map::RandomState;
use std::hash::{BuildHasher, Hasher};
use std::ops::Bound;
use std::sync::LazyLock;

pub(crate) static SCAN_TOKEN_HASHER: LazyLock<RandomState> = LazyLock::new(RandomState::new);
pub(crate) const MAX_SCAN_LIMIT: u64 = 10_000;
pub(crate) const MAX_SCAN_BYTES: u64 = 64 * 1024 * 1024;
pub(crate) const MAX_SCAN_TOKEN_BYTES: usize = 4096;
pub(crate) const SCAN_TOKEN_MAGIC: &[u8; 4] = b"LLS1";

#[derive(Clone, Debug, Eq, PartialEq)]
pub(crate) enum ScanMode {
    Prefix(Vec<u8>),
    Range {
        from: Option<Vec<u8>>,
        to: Option<Vec<u8>>,
        from_inclusive: bool,
        to_inclusive: bool,
    },
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum ScanError {
    InvalidDatabase,
    InvalidMode,
    InvalidLimit,
    InvalidMaxBytes,
    RowTooLarge,
    InvalidContinuation,
    StaleContinuation,
    OutOfMemory,
    MapResized,
    ReadersFull,
    Io,
    Database,
}

impl ScanError {
    pub(crate) fn atom(self) -> Atom {
        match self {
            Self::InvalidDatabase => atoms::invalid_database(),
            Self::InvalidMode => atoms::invalid_mode(),
            Self::InvalidLimit => atoms::invalid_limit(),
            Self::InvalidMaxBytes => atoms::invalid_max_bytes(),
            Self::RowTooLarge => atoms::row_too_large(),
            Self::InvalidContinuation => atoms::invalid_continuation(),
            Self::StaleContinuation => atoms::stale_continuation(),
            Self::OutOfMemory => atoms::out_of_memory(),
            Self::MapResized => atoms::map_resized(),
            Self::ReadersFull => atoms::readers_full(),
            Self::Io => atoms::io_error(),
            Self::Database => atoms::database_error(),
        }
    }
}

pub(crate) fn map_scan_heed_error(error: heed::Error) -> ScanError {
    match map_heed_error(error) {
        CrudError::OutOfMemory => ScanError::OutOfMemory,
        CrudError::MapResized => ScanError::MapResized,
        CrudError::ReadersFull => ScanError::ReadersFull,
        CrudError::Io => ScanError::Io,
        CrudError::InvalidDatabase => ScanError::InvalidDatabase,
        _ => ScanError::Database,
    }
}

pub(crate) fn scan_token_tag(bytes: &[u8]) -> u64 {
    // RandomState provides a process-local keyed SipHash-family hasher. The
    // integrity tag rejects accidental or unsophisticated token modification;
    // it is not a cryptographic authentication contract. Tokens remain
    // process-local and free of native pointers.
    let mut hasher = SCAN_TOKEN_HASHER.build_hasher();
    hasher.write(bytes);
    hasher.finish()
}

pub(crate) fn push_sized_bytes(output: &mut Vec<u8>, bytes: &[u8]) -> Result<(), ScanError> {
    let length = u32::try_from(bytes.len()).map_err(|_| ScanError::InvalidContinuation)?;
    output.extend_from_slice(&length.to_be_bytes());
    output.extend_from_slice(bytes);
    Ok(())
}

pub(crate) fn scan_binding(database_name: &str, mode: &ScanMode) -> Result<Vec<u8>, ScanError> {
    let mode_bytes = match mode {
        ScanMode::Prefix(prefix) => 1_usize
            .checked_add(4)
            .and_then(|size| size.checked_add(prefix.len())),
        ScanMode::Range { from, to, .. } => {
            [from, to].into_iter().try_fold(3_usize, |size, bound| {
                let field_size = bound.as_ref().map_or(Some(1), |bytes| {
                    1_usize.checked_add(4)?.checked_add(bytes.len())
                })?;
                size.checked_add(field_size)
            })
        }
    }
    .ok_or(ScanError::OutOfMemory)?;
    let capacity = 4_usize
        .checked_add(database_name.len())
        .and_then(|size| size.checked_add(mode_bytes))
        .ok_or(ScanError::OutOfMemory)?;
    let mut binding = Vec::new();
    binding
        .try_reserve_exact(capacity)
        .map_err(|_| ScanError::OutOfMemory)?;
    push_sized_bytes(&mut binding, database_name.as_bytes())?;
    match mode {
        ScanMode::Prefix(prefix) => {
            binding.push(0);
            push_sized_bytes(&mut binding, prefix)?;
        }
        ScanMode::Range {
            from,
            to,
            from_inclusive,
            to_inclusive,
        } => {
            binding.push(1);
            for bound in [from, to] {
                match bound {
                    None => binding.push(0),
                    Some(bytes) => {
                        binding.push(1);
                        push_sized_bytes(&mut binding, bytes)?;
                    }
                }
            }
            binding.push(u8::from(*from_inclusive));
            binding.push(u8::from(*to_inclusive));
        }
    }
    Ok(binding)
}

pub(crate) fn make_scan_token(
    environment_id: u64,
    database_name: &str,
    mode: &ScanMode,
    last_key: &[u8],
) -> Result<Vec<u8>, ScanError> {
    let binding = scan_binding(database_name, mode)?;
    let capacity = 4_usize
        .checked_add(8)
        .and_then(|size| size.checked_add(4))
        .and_then(|size| size.checked_add(binding.len()))
        .and_then(|size| size.checked_add(4))
        .and_then(|size| size.checked_add(last_key.len()))
        .and_then(|size| size.checked_add(8))
        .ok_or(ScanError::OutOfMemory)?;
    let mut token = Vec::new();
    token
        .try_reserve_exact(capacity)
        .map_err(|_| ScanError::OutOfMemory)?;
    token.extend_from_slice(SCAN_TOKEN_MAGIC);
    token.extend_from_slice(&environment_id.to_be_bytes());
    push_sized_bytes(&mut token, &binding)?;
    push_sized_bytes(&mut token, last_key)?;
    let tag = scan_token_tag(&token);
    token.extend_from_slice(&tag.to_be_bytes());
    Ok(token)
}

pub(crate) fn take_token_field<'a>(
    token: &'a [u8],
    offset: &mut usize,
) -> Result<&'a [u8], ScanError> {
    let length_end = offset
        .checked_add(4)
        .ok_or(ScanError::InvalidContinuation)?;
    let length_bytes: [u8; 4] = token
        .get(*offset..length_end)
        .ok_or(ScanError::InvalidContinuation)?
        .try_into()
        .map_err(|_| ScanError::InvalidContinuation)?;
    *offset = length_end;
    let end = offset
        .checked_add(u32::from_be_bytes(length_bytes) as usize)
        .ok_or(ScanError::InvalidContinuation)?;
    let field = token
        .get(*offset..end)
        .ok_or(ScanError::InvalidContinuation)?;
    *offset = end;
    Ok(field)
}

pub(crate) fn environment_id_is_live(environment_id: u64) -> bool {
    registry().values().any(|state| {
        state
            .upgrade()
            .is_some_and(|state| state.id == environment_id)
    })
}

pub(crate) fn scan_mode_contains_key(mode: &ScanMode, key: &[u8]) -> bool {
    match mode {
        ScanMode::Prefix(prefix) => key.starts_with(prefix),
        ScanMode::Range {
            from,
            to,
            from_inclusive,
            to_inclusive,
        } => {
            let above_from = from.as_ref().is_none_or(|from| {
                if *from_inclusive {
                    key >= from.as_slice()
                } else {
                    key > from.as_slice()
                }
            });
            let below_to = to.as_ref().is_none_or(|to| {
                if *to_inclusive {
                    key <= to.as_slice()
                } else {
                    key < to.as_slice()
                }
            });
            above_from && below_to
        }
    }
}

pub(crate) fn parse_scan_token(
    token: &[u8],
    environment_id: u64,
    database_name: &str,
    mode: &ScanMode,
    max_key_size: usize,
) -> Result<Vec<u8>, ScanError> {
    if token.len() < 4 + 8 + 4 + 4 + 8
        || token.len() > MAX_SCAN_TOKEN_BYTES
        || token.get(..4) != Some(SCAN_TOKEN_MAGIC.as_slice())
    {
        return Err(ScanError::InvalidContinuation);
    }
    let payload_end = token.len() - 8;
    let expected_tag = u64::from_be_bytes(
        token[payload_end..]
            .try_into()
            .map_err(|_| ScanError::InvalidContinuation)?,
    );
    if scan_token_tag(&token[..payload_end]) != expected_tag {
        return Err(ScanError::InvalidContinuation);
    }

    let token_environment = u64::from_be_bytes(
        token[4..12]
            .try_into()
            .map_err(|_| ScanError::InvalidContinuation)?,
    );
    let mut offset = 12;
    let binding = take_token_field(&token[..payload_end], &mut offset)?;
    let last_key = take_token_field(&token[..payload_end], &mut offset)?;
    if offset != payload_end {
        return Err(ScanError::InvalidContinuation);
    }
    if token_environment != environment_id {
        return Err(if environment_id_is_live(token_environment) {
            ScanError::InvalidContinuation
        } else {
            ScanError::StaleContinuation
        });
    }
    if binding != scan_binding(database_name, mode)?.as_slice()
        || last_key.is_empty()
        || last_key.len() > max_key_size
        || !scan_mode_contains_key(mode, last_key)
    {
        return Err(ScanError::InvalidContinuation);
    }
    copy_bytes(last_key).map_err(|error| match error {
        CrudError::OutOfMemory => ScanError::OutOfMemory,
        _ => ScanError::Database,
    })
}

pub(crate) fn prefix_upper_bound(prefix: &[u8]) -> Result<Option<Vec<u8>>, ScanError> {
    let Some(index) = prefix.iter().rposition(|byte| *byte != 0xff) else {
        return Ok(None);
    };
    let upper_len = index.checked_add(1).ok_or(ScanError::OutOfMemory)?;
    let mut upper = Vec::new();
    upper
        .try_reserve_exact(upper_len)
        .map_err(|_| ScanError::OutOfMemory)?;
    upper.extend_from_slice(&prefix[..=index]);
    upper[index] += 1;
    Ok(Some(upper))
}

pub(crate) fn scan_bounds_are_empty(start: &Bound<&[u8]>, end: &Bound<&[u8]>) -> bool {
    let (start_key, start_inclusive) = match start {
        Bound::Unbounded => {
            return matches!(end, Bound::Included(key) | Bound::Excluded(key) if key.is_empty());
        }
        Bound::Included(key) => (*key, true),
        Bound::Excluded(key) => (*key, false),
    };
    let (end_key, end_inclusive) = match end {
        Bound::Unbounded => return false,
        Bound::Included(key) => (*key, true),
        Bound::Excluded(key) => (*key, false),
    };
    start_key > end_key
        || (start_key == end_key && !(start_inclusive && end_inclusive))
        || end_key.is_empty()
}

#[derive(Debug, Eq, PartialEq)]
pub(crate) struct ScanPage {
    pub(crate) rows: Vec<(Vec<u8>, Vec<u8>)>,
    pub(crate) continuation: Option<Vec<u8>>,
}

pub(crate) fn scan_binary(
    database: &DatabaseResource,
    mode: &ScanMode,
    limit: u64,
    max_bytes: u64,
    continuation: Option<&[u8]>,
) -> Result<ScanPage, ScanError> {
    if !(1..=MAX_SCAN_LIMIT).contains(&limit) {
        return Err(ScanError::InvalidLimit);
    }
    if !(1..=MAX_SCAN_BYTES).contains(&max_bytes) {
        return Err(ScanError::InvalidMaxBytes);
    }
    let max_key_size = database.state.env.max_key_size();
    let resume = continuation
        .map(|token| parse_scan_token(token, database.state.id, &database.name, mode, max_key_size))
        .transpose()?;

    // No valid LMDB key can have an oversized prefix. Parse a supplied token first
    // so malformed continuations still have stable invalid-continuation behavior.
    if matches!(mode, ScanMode::Prefix(prefix) if prefix.len() > max_key_size) {
        return Ok(ScanPage {
            rows: Vec::new(),
            continuation: None,
        });
    }

    // An oversized lower bound lies immediately after its max-sized prefix in the
    // valid-key universe. Its successor is therefore the first possible key; an
    // all-FF prefix has no successor and makes the range empty.
    let oversized_lower = match mode {
        ScanMode::Range {
            from: Some(from), ..
        } if from.len() > max_key_size => match prefix_upper_bound(&from[..max_key_size])? {
            Some(successor) => Some(successor),
            None => {
                return Ok(ScanPage {
                    rows: Vec::new(),
                    continuation: None,
                });
            }
        },
        _ => None,
    };

    // Keep normalized storage alive while heed owns borrowed bounds. Prefix bounds
    // are materialized here rather than relying on a prefix API, including all-FF.
    let prefix_upper = match mode {
        ScanMode::Prefix(prefix) => prefix_upper_bound(prefix)?,
        _ => None,
    };
    let start: Bound<&[u8]> = if let Some(last_key) = resume.as_deref() {
        Bound::Excluded(last_key)
    } else if let Some(successor) = oversized_lower.as_deref() {
        Bound::Included(successor)
    } else {
        match mode {
            ScanMode::Prefix(prefix) if prefix.is_empty() => Bound::Unbounded,
            ScanMode::Prefix(prefix) => Bound::Included(prefix.as_slice()),
            ScanMode::Range { from: None, .. } => Bound::Unbounded,
            ScanMode::Range {
                from: Some(from), ..
            } if from.is_empty() => Bound::Unbounded,
            ScanMode::Range {
                from: Some(from),
                from_inclusive: true,
                ..
            } => Bound::Included(from.as_slice()),
            ScanMode::Range {
                from: Some(from), ..
            } => Bound::Excluded(from.as_slice()),
        }
    };
    let end: Bound<&[u8]> = match mode {
        ScanMode::Prefix(_) => prefix_upper
            .as_deref()
            .map_or(Bound::Unbounded, Bound::Excluded),
        ScanMode::Range { to: None, .. } => Bound::Unbounded,
        ScanMode::Range { to: Some(to), .. } if to.len() > max_key_size => {
            Bound::Included(&to[..max_key_size])
        }
        ScanMode::Range {
            to: Some(to),
            to_inclusive: true,
            ..
        } => Bound::Included(to.as_slice()),
        ScanMode::Range { to: Some(to), .. } => Bound::Excluded(to.as_slice()),
    };
    if scan_bounds_are_empty(&start, &end) {
        return Ok(ScanPage {
            rows: Vec::new(),
            continuation: None,
        });
    }
    let bounds = (start, end);
    let transaction = database.state.env.read_txn().map_err(map_scan_heed_error)?;
    let mut iterator = database
        .database
        .range(&transaction, &bounds)
        .map_err(map_scan_heed_error)?;
    let mut rows = Vec::new();
    rows.try_reserve(limit as usize)
        .map_err(|_| ScanError::OutOfMemory)?;
    let mut used_bytes = 0_u64;
    let mut has_more = false;
    for entry in iterator.by_ref() {
        let (key, value) = entry.map_err(map_scan_heed_error)?;
        let row_bytes = u64::try_from(key.len())
            .ok()
            .and_then(|bytes| bytes.checked_add(value.len() as u64))
            .ok_or(ScanError::OutOfMemory)?;
        if rows.len() as u64 == limit
            || (!rows.is_empty() && used_bytes.saturating_add(row_bytes) > max_bytes)
        {
            has_more = true;
            break;
        }
        if row_bytes > max_bytes {
            return Err(ScanError::RowTooLarge);
        }
        rows.push((
            copy_bytes(key).map_err(|_| ScanError::OutOfMemory)?,
            copy_bytes(value).map_err(|_| ScanError::OutOfMemory)?,
        ));
        used_bytes = used_bytes.saturating_add(row_bytes);
    }
    drop(iterator);
    transaction.commit().map_err(map_scan_heed_error)?;

    let continuation = if has_more {
        let last_key = rows.last().ok_or(ScanError::RowTooLarge)?.0.as_slice();
        Some(make_scan_token(
            database.state.id,
            &database.name,
            mode,
            last_key,
        )?)
    } else {
        None
    };
    Ok(ScanPage { rows, continuation })
}