lean_lmdb

use rustler::{Atom, Binary, Encoder, Env as NifEnv, OwnedBinary, ResourceArc, Term};
#[cfg(all(debug_assertions, feature = "test-support"))]
use std::fs;
#[cfg(all(debug_assertions, feature = "test-support"))]
use std::path::Path;
use std::sync::Arc;

mod atoms;
mod crud;
mod lifecycle;
mod mutations;
mod scan;

type NifResult<T> = Result<T, rustler::Atom>;
type InternalResult<T> = Result<T, LifecycleError>;

use crud::{
    CrudError, clear_stale_readers, copy_bytes, decode_binary, decode_crud_database, delete_binary,
    map_heed_error, put_binary, validate_key,
};
use lifecycle::{
    DatabaseResource, EnvironmentResource, LifecycleError, canonical_environment_path,
    create_named_database, decode_database, decode_environment, open_named_database,
    open_shared_environment, registry, validate_database_name, validate_options,
};
use mutations::{
    CompareExchangeResult, ExpectedValue, ReplacementValue, compare_exchange_binary,
    decode_batch_operations, write_batch_operations,
};
use scan::{MAX_SCAN_BYTES, MAX_SCAN_LIMIT, ScanError, ScanMode, scan_binary};

fn version() -> (&'static str, &'static str) {
    ("lean_lmdb", env!("CARGO_PKG_VERSION"))
}

#[rustler::nif]
fn native_version() -> (&'static str, &'static str) {
    version()
}

#[rustler::nif]
fn system_page_size() -> usize {
    page_size::get()
}

#[rustler::nif(schedule = "DirtyIo")]
#[allow(clippy::too_many_arguments)]
fn environment_open(
    path: String,
    map_size: u64,
    max_dbs: u32,
    max_readers: u32,
    read_only: bool,
    create: bool,
    durability_code: u8,
    fixed_map: bool,
    read_ahead: bool,
) -> NifResult<ResourceArc<EnvironmentResource>> {
    let options = validate_options(
        map_size,
        max_dbs,
        max_readers,
        read_only,
        create,
        durability_code,
        fixed_map,
        read_ahead,
    )
    .map_err(LifecycleError::atom)?;
    let path = canonical_environment_path(&path, create).map_err(LifecycleError::atom)?;
    let state = open_shared_environment(path, options, create).map_err(LifecycleError::atom)?;
    Ok(ResourceArc::new(EnvironmentResource { state }))
}

#[rustler::nif(schedule = "DirtyIo")]
fn environment_sync(environment: Term<'_>) -> NifResult<()> {
    let environment = decode_environment(environment).map_err(LifecycleError::atom)?;
    if environment.state.options.read_only {
        return Err(atoms::read_only());
    }
    environment
        .state
        .env
        .force_sync()
        .map_err(|_| atoms::sync_failed())
}

#[rustler::nif(schedule = "DirtyIo")]
fn database_create<'a>(
    environment: Term<'a>,
    name: Binary<'a>,
) -> NifResult<ResourceArc<DatabaseResource>> {
    let environment = decode_environment(environment).map_err(LifecycleError::atom)?;
    let name = validate_database_name(name.as_slice()).map_err(LifecycleError::atom)?;
    let database = create_named_database(&environment.state, name).map_err(LifecycleError::atom)?;
    Ok(ResourceArc::new(DatabaseResource {
        state: Arc::clone(&environment.state),
        database,
        name: name.to_owned(),
    }))
}

#[rustler::nif(schedule = "DirtyIo")]
fn database_open<'a>(
    environment: Term<'a>,
    name: Binary<'a>,
) -> NifResult<ResourceArc<DatabaseResource>> {
    let environment = decode_environment(environment).map_err(LifecycleError::atom)?;
    let name = validate_database_name(name.as_slice()).map_err(LifecycleError::atom)?;
    let database = open_named_database(&environment.state, name).map_err(LifecycleError::atom)?;
    Ok(ResourceArc::new(DatabaseResource {
        state: Arc::clone(&environment.state),
        database,
        name: name.to_owned(),
    }))
}

enum LookupResult<'a> {
    Found(Binary<'a>),
    Missing,
}

impl Encoder for LookupResult<'_> {
    fn encode<'a>(&self, env: NifEnv<'a>) -> Term<'a> {
        match self {
            Self::Found(binary) => binary.encode(env),
            Self::Missing => atoms::not_found().encode(env),
        }
    }
}

#[rustler::nif(schedule = "DirtyIo")]
fn database_get<'a>(database: Term<'a>, key: Term<'a>) -> NifResult<LookupResult<'a>> {
    let env = database.get_env();
    let database = decode_crud_database(database).map_err(CrudError::atom)?;
    let key = decode_binary(key, CrudError::InvalidKey).map_err(CrudError::atom)?;
    let key = validate_key(&database.state, key.as_slice()).map_err(CrudError::atom)?;
    let transaction = database
        .state
        .env
        .read_txn()
        .map_err(map_heed_error)
        .map_err(CrudError::atom)?;
    let copied = match database
        .database
        .get(&transaction, key)
        .map_err(map_heed_error)
        .map_err(CrudError::atom)?
    {
        Some(value) => {
            let mut copied = OwnedBinary::new(value.len()).ok_or_else(atoms::out_of_memory)?;
            copied.as_mut_slice().copy_from_slice(value);
            Some(copied)
        }
        None => None,
    };
    transaction
        .commit()
        .map_err(map_heed_error)
        .map_err(CrudError::atom)?;
    Ok(match copied {
        Some(binary) => LookupResult::Found(binary.release(env)),
        None => LookupResult::Missing,
    })
}

fn decode_scan_bound(term: Term<'_>, max_key_size: usize) -> Result<Option<Vec<u8>>, ScanError> {
    if matches!(term.decode::<Atom>(), Ok(atom) if atom == atoms::nil()) {
        Ok(None)
    } else {
        let binary = term
            .decode::<Binary<'_>>()
            .map_err(|_| ScanError::InvalidMode)?;
        if binary.len() > max_key_size {
            return Err(ScanError::InvalidMode);
        }
        copy_bytes(binary.as_slice())
            .map(Some)
            .map_err(|error| match error {
                CrudError::OutOfMemory => ScanError::OutOfMemory,
                _ => ScanError::InvalidMode,
            })
    }
}

fn decode_scan_mode(term: Term<'_>, max_key_size: usize) -> Result<ScanMode, ScanError> {
    if let Ok((tag, prefix)) = term.decode::<(Atom, Binary<'_>)>() {
        if tag != atoms::prefix() || prefix.len() > max_key_size {
            return Err(ScanError::InvalidMode);
        }
        return copy_bytes(prefix.as_slice())
            .map(ScanMode::Prefix)
            .map_err(|_| ScanError::OutOfMemory);
    }
    if let Ok((tag, from, to, from_inclusive, to_inclusive)) =
        term.decode::<(Atom, Term<'_>, Term<'_>, bool, bool)>()
    {
        if tag != atoms::range() {
            return Err(ScanError::InvalidMode);
        }
        return Ok(ScanMode::Range {
            from: decode_scan_bound(from, max_key_size)?,
            to: decode_scan_bound(to, max_key_size)?,
            from_inclusive,
            to_inclusive,
        });
    }
    Err(ScanError::InvalidMode)
}

enum ScanContinuationNif<'a> {
    Done,
    Token(Binary<'a>),
}

impl Encoder for ScanContinuationNif<'_> {
    fn encode<'a>(&self, env: NifEnv<'a>) -> Term<'a> {
        match self {
            Self::Done => atoms::done().encode(env),
            Self::Token(token) => token.encode(env),
        }
    }
}

#[rustler::nif(schedule = "DirtyIo")]
fn database_scan<'a>(
    database: Term<'a>,
    mode: Term<'a>,
    limit: Term<'a>,
    max_bytes: Term<'a>,
    continuation: Term<'a>,
) -> NifResult<(Vec<(Binary<'a>, Binary<'a>)>, ScanContinuationNif<'a>)> {
    let env = database.get_env();
    let database = decode_crud_database(database).map_err(|_| atoms::invalid_database())?;
    let mode =
        decode_scan_mode(mode, database.state.env.max_key_size()).map_err(ScanError::atom)?;
    let limit = limit.decode::<u64>().map_err(|_| atoms::invalid_limit())?;
    if !(1..=MAX_SCAN_LIMIT).contains(&limit) {
        return Err(atoms::invalid_limit());
    }
    let max_bytes = max_bytes
        .decode::<u64>()
        .map_err(|_| atoms::invalid_max_bytes())?;
    if !(1..=MAX_SCAN_BYTES).contains(&max_bytes) {
        return Err(atoms::invalid_max_bytes());
    }
    let continuation_binary = if matches!(continuation.decode::<Atom>(), Ok(atom) if atom == atoms::nil())
    {
        None
    } else {
        Some(
            continuation
                .decode::<Binary<'a>>()
                .map_err(|_| atoms::invalid_continuation())?,
        )
    };
    let page = scan_binary(
        &database,
        &mode,
        limit,
        max_bytes,
        continuation_binary.as_ref().map(Binary::as_slice),
    )
    .map_err(ScanError::atom)?;

    let mut rows = Vec::new();
    rows.try_reserve_exact(page.rows.len())
        .map_err(|_| atoms::out_of_memory())?;
    for (key, value) in page.rows {
        let mut key_binary = OwnedBinary::new(key.len()).ok_or_else(atoms::out_of_memory)?;
        key_binary.as_mut_slice().copy_from_slice(&key);
        let mut value_binary = OwnedBinary::new(value.len()).ok_or_else(atoms::out_of_memory)?;
        value_binary.as_mut_slice().copy_from_slice(&value);
        rows.push((key_binary.release(env), value_binary.release(env)));
    }
    let continuation = match page.continuation {
        None => ScanContinuationNif::Done,
        Some(token) => {
            let mut binary = OwnedBinary::new(token.len()).ok_or_else(atoms::out_of_memory)?;
            binary.as_mut_slice().copy_from_slice(&token);
            ScanContinuationNif::Token(binary.release(env))
        }
    };
    Ok((rows, continuation))
}

#[rustler::nif(schedule = "DirtyIo")]
fn database_put<'a>(database: Term<'a>, key: Term<'a>, value: Term<'a>) -> NifResult<()> {
    let database = decode_crud_database(database).map_err(CrudError::atom)?;
    let key = decode_binary(key, CrudError::InvalidKey).map_err(CrudError::atom)?;
    let value = decode_binary(value, CrudError::InvalidValue).map_err(CrudError::atom)?;
    put_binary(&database, key.as_slice(), value.as_slice()).map_err(CrudError::atom)
}

#[rustler::nif(schedule = "DirtyIo")]
fn database_delete<'a>(database: Term<'a>, key: Term<'a>) -> NifResult<()> {
    let database = decode_crud_database(database).map_err(CrudError::atom)?;
    let key = decode_binary(key, CrudError::InvalidKey).map_err(CrudError::atom)?;
    delete_binary(&database, key.as_slice()).map_err(CrudError::atom)
}

#[rustler::nif(schedule = "DirtyIo")]
fn environment_write_batch(environment: Term<'_>, operations: Term<'_>) -> NifResult<()> {
    let environment = environment
        .decode::<ResourceArc<EnvironmentResource>>()
        .map_err(|_| atoms::invalid_environment())?;
    let operations = decode_batch_operations(&environment, operations).map_err(CrudError::atom)?;
    write_batch_operations(&environment.state, &operations).map_err(CrudError::atom)
}

enum CompareExchangeNifResult<'a> {
    Exchanged,
    ConflictMissing,
    ConflictValue(Binary<'a>),
}

impl Encoder for CompareExchangeNifResult<'_> {
    fn encode<'a>(&self, env: NifEnv<'a>) -> Term<'a> {
        match self {
            Self::Exchanged => atoms::ok().encode(env),
            Self::ConflictMissing => (atoms::conflict(), atoms::missing()).encode(env),
            Self::ConflictValue(value) => (atoms::conflict(), value).encode(env),
        }
    }
}

#[rustler::nif(schedule = "DirtyIo")]
fn database_compare_exchange<'a>(
    database: Term<'a>,
    key: Term<'a>,
    expected: Term<'a>,
    replacement: Term<'a>,
) -> NifResult<CompareExchangeNifResult<'a>> {
    let env = database.get_env();
    let database = decode_crud_database(database).map_err(CrudError::atom)?;
    let key = decode_binary(key, CrudError::InvalidKey).map_err(CrudError::atom)?;
    validate_key(&database.state, key.as_slice()).map_err(CrudError::atom)?;

    let expected = if matches!(expected.decode::<Atom>(), Ok(atom) if atom == atoms::missing()) {
        ExpectedValue::Missing
    } else if let Ok(value) = expected.decode::<Binary<'a>>() {
        ExpectedValue::Value(value.as_slice())
    } else {
        return Err(CrudError::InvalidExpected.atom());
    };
    let replacement = if matches!(replacement.decode::<Atom>(), Ok(atom) if atom == atoms::delete())
    {
        ReplacementValue::Delete
    } else if let Ok((tag, value)) = replacement.decode::<(Atom, Binary<'a>)>() {
        if tag != atoms::put() {
            return Err(CrudError::InvalidReplacement.atom());
        }
        ReplacementValue::Put(value.as_slice())
    } else {
        return Err(CrudError::InvalidReplacement.atom());
    };

    match compare_exchange_binary(
        &database,
        key.as_slice(),
        &expected,
        &replacement,
        |value| {
            let mut copied = OwnedBinary::new(value.len()).ok_or(CrudError::OutOfMemory)?;
            copied.as_mut_slice().copy_from_slice(value);
            Ok(copied)
        },
    )
    .map_err(CrudError::atom)?
    {
        CompareExchangeResult::Exchanged => Ok(CompareExchangeNifResult::Exchanged),
        CompareExchangeResult::Conflict(None) => Ok(CompareExchangeNifResult::ConflictMissing),
        CompareExchangeResult::Conflict(Some(value)) => {
            Ok(CompareExchangeNifResult::ConflictValue(value.release(env)))
        }
    }
}

#[rustler::nif(schedule = "DirtyIo")]
fn environment_clear_stale_readers(environment: Term<'_>) -> NifResult<usize> {
    let environment = decode_environment(environment).map_err(LifecycleError::atom)?;
    clear_stale_readers(&environment.state).map_err(CrudError::atom)
}

#[cfg(all(debug_assertions, feature = "test-support"))]
fn wait_for_test_release(release_path: &str) -> Result<(), rustler::Atom> {
    use std::time::{Duration, Instant};

    const TEST_TRANSACTION_TIMEOUT: Duration = Duration::from_secs(60);
    let deadline = Instant::now()
        .checked_add(TEST_TRANSACTION_TIMEOUT)
        .ok_or_else(atoms::test_timeout)?;
    loop {
        if Path::new(release_path).is_file() {
            return Ok(());
        }
        if Instant::now() >= deadline {
            return Err(atoms::test_timeout());
        }
        std::thread::sleep(Duration::from_millis(5));
    }
}

#[cfg(all(debug_assertions, feature = "test-support"))]
fn publish_test_marker(marker_path: &str) -> Result<(), rustler::Atom> {
    fs::write(marker_path, b"active").map_err(|_| atoms::io_error())
}

/// Test-only coordination NIF. The marker is published only after LMDB has
/// installed the active reader; the transaction is always released by the
/// bounded poll or by process termination.
#[cfg(all(debug_assertions, feature = "test-support"))]
#[rustler::nif(schedule = "DirtyIo")]
fn test_hold_read_transaction<'a>(
    database: Term<'a>,
    key: Term<'a>,
    marker_path: String,
    release_path: String,
) -> NifResult<()> {
    let database = decode_crud_database(database).map_err(CrudError::atom)?;
    let key = decode_binary(key, CrudError::InvalidKey).map_err(CrudError::atom)?;
    let key = validate_key(&database.state, key.as_slice()).map_err(CrudError::atom)?;
    let transaction = database
        .state
        .env
        .read_txn()
        .map_err(map_heed_error)
        .map_err(CrudError::atom)?;
    let _value = database
        .database
        .get(&transaction, key)
        .map_err(map_heed_error)
        .map_err(CrudError::atom)?;
    publish_test_marker(&marker_path)?;
    wait_for_test_release(&release_path)?;
    transaction
        .commit()
        .map_err(map_heed_error)
        .map_err(CrudError::atom)
}

/// Test-only coordination NIF. Acquiring the transaction before publishing the
/// marker proves that the LMDB writer lock is held while the test coordinates.
#[cfg(all(debug_assertions, feature = "test-support"))]
#[rustler::nif(schedule = "DirtyIo")]
fn test_hold_write_transaction<'a>(
    database: Term<'a>,
    key: Term<'a>,
    value: Term<'a>,
    marker_path: String,
    release_path: String,
) -> NifResult<()> {
    let database = decode_crud_database(database).map_err(CrudError::atom)?;
    let key = decode_binary(key, CrudError::InvalidKey).map_err(CrudError::atom)?;
    let key = validate_key(&database.state, key.as_slice()).map_err(CrudError::atom)?;
    let value = decode_binary(value, CrudError::InvalidValue).map_err(CrudError::atom)?;
    if database.state.options.read_only {
        return Err(atoms::read_only());
    }
    let mut transaction = database
        .state
        .env
        .write_txn()
        .map_err(map_heed_error)
        .map_err(CrudError::atom)?;
    database
        .database
        .put(&mut transaction, key, value.as_slice())
        .map_err(map_heed_error)
        .map_err(CrudError::atom)?;
    publish_test_marker(&marker_path)?;
    wait_for_test_release(&release_path)?;
    transaction
        .commit()
        .map_err(map_heed_error)
        .map_err(CrudError::atom)
}

#[rustler::nif]
fn environment_info(environment: Term<'_>) -> NifResult<(u64, String, usize, u32, u32, bool)> {
    let environment = decode_environment(environment).map_err(LifecycleError::atom)?;
    let state = &environment.state;
    let path = state
        .path
        .to_str()
        .ok_or(LifecycleError::InvalidPath)
        .map_err(LifecycleError::atom)?;
    Ok((
        state.id,
        path.to_owned(),
        state.options.map_size,
        state.options.max_dbs,
        state.options.max_readers,
        state.options.read_only,
    ))
}

#[rustler::nif]
fn database_info(database: Term<'_>) -> NifResult<(u64, String)> {
    let database = decode_database(database).map_err(LifecycleError::atom)?;
    // Touch the immutable DBI handle here so this diagnostic also proves that
    // the database resource—not merely its Elixir metadata—decoded correctly.
    let _database_handle = database.database;
    Ok((database.state.id, database.name.clone()))
}

#[rustler::nif(schedule = "DirtyIo")]
fn registry_info() -> (usize, usize, usize) {
    let mut entries = registry();
    let total = entries.len();
    let live = entries
        .values()
        .filter(|state| state.strong_count() != 0)
        .count();
    let stale = total - live;
    // Report stale entries once, then ensure the registry cannot accumulate them.
    entries.retain(|_, state| state.strong_count() != 0);
    (live, stale, total)
}

rustler::init!("Elixir.LeanLmdb.Native");

#[cfg(test)]
mod tests;