reifydb-engine 0.10.1

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

use std::sync::Arc;

use arrow_array::{Array, ArrayRef, RecordBatch};
use arrow_schema::FieldRef;
use reifydb_core::{
	expression::Expression,
	value::{batch::batch, column::headers::ColumnHeaders},
};
use reifydb_evaluate::expression::{context::EvalContext, eval::evaluate};
use reifydb_routine_abi::{Function, Procedure, context::FunctionContext};
use reifydb_transaction::transaction::Transaction;
use reifydb_value::{
	fragment::Fragment,
	params::Params,
	value::{Value, column_view::ColumnView},
};
use tracing::instrument;

use crate::{
	Result,
	error::EngineError,
	vm::{
		callable::{CallSite, ProcedureCall, invoke_procedure_routine},
		volcano::query::{QueryContext, QueryNode, eval_context_from_query},
	},
};

enum GeneratorImpl {
	Function(Arc<dyn Function>),
	Procedure(Arc<dyn Procedure>),
}

pub(crate) struct GeneratorNode {
	function_name: Fragment,
	expressions: Vec<Expression>,
	context: Option<Arc<QueryContext>>,
	exhausted: bool,
	generator: Option<GeneratorImpl>,
}

impl GeneratorNode {
	pub fn new(function_name: Fragment, parameter_expressions: Vec<Expression>) -> Self {
		Self {
			function_name,
			expressions: parameter_expressions,
			context: None,
			exhausted: false,
			generator: None,
		}
	}

	#[instrument(level = "trace", skip_all, name = "volcano::generator::eval_params")]
	fn eval_params(&self, session: &EvalContext<'_>) -> Result<Vec<(FieldRef, ArrayRef)>> {
		let evaluation_ctx = session.with_eval_empty();

		let mut evaluated_columns = Vec::new();
		for expr in &self.expressions {
			let column = evaluate(&evaluation_ctx, expr)?;
			evaluated_columns.push(column);
		}
		Ok(evaluated_columns)
	}

	#[instrument(level = "trace", skip_all, name = "volcano::generator::invoke")]
	fn invoke<'a>(
		&self,
		txn: &mut Transaction<'a>,
		stored_ctx: &QueryContext,
		evaluated_columns: Vec<(FieldRef, ArrayRef)>,
	) -> Result<RecordBatch> {
		match self.generator.as_ref().unwrap() {
			GeneratorImpl::Function(generator) => {
				let mut fn_ctx = FunctionContext {
					fragment: self.function_name.clone(),
					identity: stored_ctx.identity,
					row_count: evaluated_columns.first().map_or(0, |(_, array)| array.len()),
					runtime_context: &stored_ctx.services.runtime_context,
				};
				batch(vec![generator.call(&mut fn_ctx, &evaluated_columns)?])
			}
			GeneratorImpl::Procedure(procedure) => {
				let values: Vec<Value> = evaluated_columns
					.iter()
					.map(|col| Ok(ColumnView::try_from(col)?.get_value(0)))
					.collect::<Result<_>>()?;
				let params = Params::Positional(Arc::new(values));
				invoke_procedure_routine(
					&stored_ctx.services,
					&stored_ctx.symbols,
					txn,
					ProcedureCall {
						routine: procedure,
						fragment: &self.function_name,
						target: self.function_name.text(),
						params: &params,
					},
					CallSite::Named,
				)
			}
		}
	}
}

impl QueryNode for GeneratorNode {
	#[instrument(level = "trace", skip_all, name = "volcano::generator::initialize")]
	fn initialize<'a>(&mut self, _txn: &mut Transaction<'a>, ctx: &QueryContext) -> Result<()> {
		self.context = Some(Arc::new(ctx.clone()));

		let name = self.function_name.text();
		if let Some(func) = ctx.services.routines.get_generator_function(name) {
			func.arity().check(&self.function_name, self.expressions.len())?;
			self.generator = Some(GeneratorImpl::Function(func));
		} else if let Some(proc) = ctx.services.routines.get_procedure(name) {
			self.generator = Some(GeneratorImpl::Procedure(proc));
		} else {
			return Err(EngineError::GeneratorNotFound {
				name: name.to_string(),
				fragment: self.function_name.clone(),
			}
			.into());
		}

		self.exhausted = false;
		Ok(())
	}

	#[instrument(level = "trace", skip_all, name = "volcano::generator::next")]
	fn next<'a>(&mut self, txn: &mut Transaction<'a>, _ctx: &mut QueryContext) -> Result<Option<RecordBatch>> {
		if self.exhausted {
			return Ok(None);
		}

		let stored_ctx = self.context.as_ref().unwrap().clone();

		let session = eval_context_from_query(&stored_ctx);
		let evaluated_columns = self.eval_params(&session)?;

		let columns = self.invoke(txn, &stored_ctx, evaluated_columns)?;

		self.exhausted = true;

		Ok(Some(columns))
	}

	fn headers(&self) -> Option<ColumnHeaders> {
		None
	}
}