reifydb-flow 0.9.1

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

use std::sync::Arc;

use reifydb_core::{
	interface::{
		catalog::flow::OperatorId,
		change::{Change, Diff},
		flow::OperatorCapability,
	},
	value::column::{ColumnWithName, columns::Columns},
};
use reifydb_evaluate::expression::{
	compile::{CompiledExpr, compile_expression},
	context::{CompileContext, EvalContext},
};
use reifydb_routine_abi::registry::Routines;
use reifydb_rql::expression::{Expression, name::display_label};
use reifydb_runtime::context::RuntimeContext;
use reifydb_value::{Result, fragment::Fragment, value::system_columns::SystemColumns};
use tracing::instrument;

use crate::{
	context::FlowContext,
	operator::{HostOperator, host::HostContext},
};

pub struct ExtendOperator {
	parent_schema: Option<Columns>,
	operator: OperatorId,
	expressions: Vec<Expression>,
	compiled_expressions: Vec<CompiledExpr>,
	routines: Routines,
	runtime_context: RuntimeContext,
	ctx: Arc<FlowContext>,
}

impl ExtendOperator {
	pub fn new(
		parent_schema: Option<Columns>,
		operator: OperatorId,
		expressions: Vec<Expression>,
		routines: Routines,
		runtime_context: RuntimeContext,
		ctx: Arc<FlowContext>,
	) -> Self {
		let compile_ctx = CompileContext {
			symbols: &ctx.symbols,
		};
		let compiled_expressions: Vec<CompiledExpr> = expressions
			.iter()
			.map(|e| compile_expression(&compile_ctx, e))
			.collect::<Result<Vec<_>>>()
			.expect("Failed to compile expressions");

		Self {
			parent_schema,
			operator,
			expressions,
			compiled_expressions,
			routines,
			runtime_context,
			ctx,
		}
	}

	pub(crate) fn output_schema(&self) -> Option<Columns> {
		self.parent_schema.clone()
	}

	#[instrument(name = "flow::operator::extend::extend", level = "trace", skip_all, fields(rows = columns.row_count()))]
	fn extend(&self, columns: &Columns) -> Result<Columns> {
		let row_count = columns.row_count();
		if row_count == 0 {
			return Ok(Columns::empty());
		}

		let session = EvalContext {
			params: &self.ctx.params,
			symbols: &self.ctx.symbols,
			routines: &self.routines,
			runtime_context: &self.runtime_context,
			identity: self.ctx.identity,
			is_aggregate_context: false,
			columns: Columns::empty(),
			row_count: 1,
			target: None,
			take: None,
		};
		let exec_ctx = session.with_eval(columns.clone(), row_count);

		let mut result_columns: Vec<ColumnWithName> =
			columns.iter().map(|col| ColumnWithName::new(col.name().clone(), col.data().clone())).collect();

		for (i, compiled_expr) in self.compiled_expressions.iter().enumerate() {
			let evaluated_col = compiled_expr.execute(&exec_ctx)?;

			let expr = &self.expressions[i];
			let field_name = display_label(expr).text().to_string();

			result_columns.push(ColumnWithName::new(
				Fragment::internal(field_name),
				evaluated_col.data().clone(),
			));
		}

		let row_numbers = if columns.row_numbers().is_empty() {
			Vec::new()
		} else {
			columns.row_numbers().to_vec()
		};

		Ok(Columns::with_system(
			result_columns,
			SystemColumns::new(
				row_numbers,
				Vec::new(),
				columns.created_at().to_vec(),
				columns.updated_at().to_vec(),
				columns.time().to_vec(),
			),
		))
	}
}

impl HostOperator for ExtendOperator {
	fn id(&self) -> OperatorId {
		self.operator
	}

	fn capabilities(&self) -> &[OperatorCapability] {
		OperatorCapability::STANDARD
	}

	fn apply(&mut self, _host: &mut dyn HostContext, change: Change) -> Result<Change> {
		let mut result = Vec::new();

		for diff in change.diffs.into_iter() {
			match diff {
				Diff::Insert {
					post,
					..
				} => {
					let extended = match self.extend(&post) {
						Ok(extended) => extended,
						Err(err) => {
							panic!("{:#?}", err)
						}
					};

					if !extended.is_empty() {
						result.push(Diff::insert(extended));
					}
				}
				Diff::Update {
					pre,
					post,
					..
				} => {
					let extended_post = self.extend(&post)?;
					let extended_pre = self.extend(&pre)?;

					if !extended_post.is_empty() {
						result.push(Diff::update(extended_pre, extended_post));
					}
				}
				Diff::Remove {
					pre,
					..
				} => {
					let extended_pre = self.extend(&pre)?;
					if !extended_pre.is_empty() {
						result.push(Diff::remove(extended_pre));
					}
				}
			}
		}

		Ok(Change::from_flow(self.operator, change.version, result, change.changed_at))
	}

	fn output_schema(&self) -> Option<Columns> {
		self.output_schema()
	}
}