reifydb-sdk 0.9.0

// SPDX-License-Identifier: Apache-2.0
// Copyright (c) 2026 ReifyDB

pub mod utils;

use std::ops::Bound;

use reifydb_codec::{
	key::encoded::EncodedKey,
	row::operator::{EncodedOperatorRow, OperatorState, decode},
};
use reifydb_core::key::operator_state::GroupStateKey;
use reifydb_value::{error::Error as ValueError, value::datetime::DateTime};

use crate::{
	error::{Result, SdkError},
	flow::operator::{
		context::{GuestContext, GuestState},
		extern_c::binding::{context::ExternCContext, state as extern_c},
	},
};

pub struct State<'a> {
	ctx: &'a mut ExternCContext,
}

impl<'a> State<'a> {
	pub(crate) fn new(ctx: &'a mut ExternCContext) -> Self {
		Self {
			ctx,
		}
	}

	pub fn get<T: OperatorState>(&self, key: &GroupStateKey) -> Result<Option<T>> {
		match self.get_bytes(key)? {
			Some(row) => decode_payload(&row).map(Some),
			None => Ok(None),
		}
	}

	pub fn set<T: OperatorState>(&mut self, key: &GroupStateKey, value: &T) -> Result<()> {
		let row = encode_payload(value, self.written_at())?;
		extern_c::set(self.ctx, key.as_encoded(), &row.into_bytes())
	}

	pub fn remove(&mut self, key: &GroupStateKey) -> Result<()> {
		extern_c::remove(self.ctx, key.as_encoded())
	}

	pub fn contains(&self, key: &GroupStateKey) -> Result<bool> {
		Ok(extern_c::get(self.ctx, key.as_encoded())?.is_some())
	}

	pub fn clear(&mut self) -> Result<()> {
		extern_c::clear(self.ctx)
	}

	pub fn scan_prefix<T: OperatorState>(&self, prefix: &GroupStateKey) -> Result<Vec<(GroupStateKey, T)>> {
		extern_c::prefix(self.ctx, prefix.as_encoded())?
			.into_iter()
			.map(|(k, row)| {
				Ok((
					framed(k)?,
					decode_payload(&EncodedOperatorRow::try_from(row).map_err(ValueError::from)?)?,
				))
			})
			.collect()
	}

	pub fn get_many<T: OperatorState>(&self, keys: &[GroupStateKey]) -> Result<Vec<(GroupStateKey, T)>> {
		let raw: Vec<EncodedKey> = keys.iter().map(|k| k.as_encoded().clone()).collect();
		extern_c::get_many(self.ctx, &raw)?
			.into_iter()
			.map(|(k, row)| {
				Ok((
					framed(k)?,
					decode_payload(&EncodedOperatorRow::try_from(row).map_err(ValueError::from)?)?,
				))
			})
			.collect()
	}

	pub fn keys_with_prefix(&self, prefix: &GroupStateKey) -> Result<Vec<GroupStateKey>> {
		extern_c::prefix(self.ctx, prefix.as_encoded())?.into_iter().map(|(k, _)| framed(k)).collect()
	}

	pub fn range<T: OperatorState>(
		&self,
		start: Bound<&GroupStateKey>,
		end: Bound<&GroupStateKey>,
	) -> Result<Vec<(GroupStateKey, T)>> {
		extern_c::range(self.ctx, start.map(GroupStateKey::as_encoded), end.map(GroupStateKey::as_encoded))?
			.into_iter()
			.map(|(k, row)| {
				Ok((
					framed(k)?,
					decode_payload(&EncodedOperatorRow::try_from(row).map_err(ValueError::from)?)?,
				))
			})
			.collect()
	}

	pub fn get_bytes(&self, key: &GroupStateKey) -> Result<Option<EncodedOperatorRow>> {
		match extern_c::get(self.ctx, key.as_encoded())? {
			Some(row) => Ok(Some(EncodedOperatorRow::try_from(row).map_err(ValueError::from)?)),
			None => Ok(None),
		}
	}

	pub fn set_bytes(&mut self, key: &GroupStateKey, payload: EncodedOperatorRow) -> Result<()> {
		extern_c::set(self.ctx, key.as_encoded(), &payload.into_bytes())
	}

	pub fn get_many_bytes_visit(
		&self,
		keys: &[GroupStateKey],
		visit: &mut dyn FnMut(GroupStateKey, EncodedOperatorRow) -> Result<()>,
	) -> Result<()> {
		let raw: Vec<EncodedKey> = keys.iter().map(|k| k.as_encoded().clone()).collect();
		for (k, row) in extern_c::get_many(self.ctx, &raw)? {
			visit(framed(k)?, EncodedOperatorRow::try_from(row).map_err(ValueError::from)?)?;
		}
		Ok(())
	}

	pub fn range_bytes_visit(
		&self,
		start: Bound<&GroupStateKey>,
		end: Bound<&GroupStateKey>,
		visit: &mut dyn FnMut(GroupStateKey, EncodedOperatorRow) -> Result<()>,
	) -> Result<()> {
		for (k, row) in extern_c::range(
			self.ctx,
			start.map(GroupStateKey::as_encoded),
			end.map(GroupStateKey::as_encoded),
		)? {
			visit(framed(k)?, EncodedOperatorRow::try_from(row).map_err(ValueError::from)?)?;
		}
		Ok(())
	}

	#[inline]
	fn written_at(&self) -> DateTime {
		// SAFETY: ExternCContext::new asserts ctx.ctx is non-null, and the host keeps the
		// ExternCContextRaw alive and aligned for at least the lifetime of the borrow this State was created
		// from.
		DateTime::from_nanos(unsafe { (*self.ctx.ctx).written_at_nanos })
	}
}

#[inline]
fn framed(key: EncodedKey) -> Result<GroupStateKey> {
	match GroupStateKey::from_framed(key) {
		Some(key) => Ok(key),
		None => Err(SdkError::Serialization("host returned a state key that is not framed".to_string())),
	}
}

#[inline]
pub fn encode_payload<T: OperatorState>(value: &T, now: DateTime) -> Result<EncodedOperatorRow> {
	Ok(value.encode_state(now).map_err(ValueError::from)?)
}

#[inline]
pub fn decode_payload<T: OperatorState>(row: &EncodedOperatorRow) -> Result<T> {
	Ok(decode(row).map_err(ValueError::from)?)
}

pub trait GuestRawOperator {
	fn state_get<T: OperatorState>(&self, ctx: &mut impl GuestContext, key: &GroupStateKey) -> Result<Option<T>> {
		ctx.state().get(key)
	}

	fn state_set<T: OperatorState>(
		&self,
		ctx: &mut impl GuestContext,
		key: &GroupStateKey,
		value: &T,
	) -> Result<()> {
		ctx.state().set(key, value)
	}

	fn state_remove(&self, ctx: &mut impl GuestContext, key: &GroupStateKey) -> Result<()> {
		ctx.state().remove(key)
	}

	fn state_clear(&self, ctx: &mut impl GuestContext) -> Result<()> {
		ctx.state().clear()
	}
}