db/heromodels/src/models/flow/rhai.rs

use rhai::plugin::*;
use rhai::{Array, Dynamic, Engine, EvalAltResult, INT, Module, Position};
use std::mem;
use std::sync::Arc;

use super::flow::Flow;
use super::flow_step::FlowStep;
use super::signature_requirement::SignatureRequirement;
type RhaiFlow = Flow;
type RhaiFlowStep = FlowStep;
type RhaiSignatureRequirement = SignatureRequirement;
use crate::db::Collection;
use crate::db::Db;
use crate::db::hero::OurDB;

// Helper to convert i64 from Rhai to u32 for IDs
fn id_from_i64_to_u32(id_i64: i64) -> Result<u32, Box<EvalAltResult>> {
    u32::try_from(id_i64).map_err(|_| {
        Box::new(EvalAltResult::ErrorArithmetic(
            format!("Failed to convert ID '{}' to u32", id_i64).into(),
            Position::NONE,
        ))
    })
}

#[export_module]
mod rhai_flow_module {
    // --- Flow Functions ---
    #[rhai_fn(name = "new_flow")]
    pub fn new_flow(flow_uuid: String) -> RhaiFlow {
        Flow::new(flow_uuid)
    }

    /// Sets the flow name
    #[rhai_fn(name = "name", return_raw, global, pure)]
    pub fn flow_name(flow: &mut RhaiFlow, name: String) -> Result<RhaiFlow, Box<EvalAltResult>> {
        let owned_flow = mem::replace(flow, Flow::new("")); // Dummy for replacement
        *flow = owned_flow.name(name);
        Ok(flow.clone())
    }

    /// Sets the flow status
    #[rhai_fn(name = "status", return_raw, global, pure)]
    pub fn flow_status(
        flow: &mut RhaiFlow,
        status: String,
    ) -> Result<RhaiFlow, Box<EvalAltResult>> {
        let owned_flow = mem::replace(flow, Flow::new("")); // Dummy for replacement
        *flow = owned_flow.status(status);
        Ok(flow.clone())
    }

    /// Adds a step to the flow
    #[rhai_fn(name = "add_step", return_raw, global, pure)]
    pub fn flow_add_step(
        flow: &mut RhaiFlow,
        step: RhaiFlowStep,
    ) -> Result<RhaiFlow, Box<EvalAltResult>> {
        let owned_flow = mem::replace(flow, Flow::new("")); // Dummy for replacement
        *flow = owned_flow.add_step(step);
        Ok(flow.clone())
    }

    // Flow Getters
    #[rhai_fn(get = "id", pure)]
    pub fn get_id(flow: &mut RhaiFlow) -> i64 {
        flow.base_data.id as i64
    }
    #[rhai_fn(get = "created_at", pure)]
    pub fn get_created_at(flow: &mut RhaiFlow) -> i64 {
        flow.base_data.created_at
    }
    #[rhai_fn(get = "modified_at", pure)]
    pub fn get_modified_at(flow: &mut RhaiFlow) -> i64 {
        flow.base_data.modified_at
    }
    #[rhai_fn(get = "flow_uuid", pure)]
    pub fn get_flow_uuid(flow: &mut RhaiFlow) -> String {
        flow.flow_uuid.clone()
    }
    #[rhai_fn(get = "name", pure)]
    pub fn get_name(flow: &mut RhaiFlow) -> String {
        flow.name.clone()
    }
    #[rhai_fn(get = "status", pure)]
    pub fn get_status(flow: &mut RhaiFlow) -> String {
        flow.status.clone()
    }
    #[rhai_fn(get = "steps", pure)]
    pub fn get_steps(flow: &mut RhaiFlow) -> Array {
        flow.steps
            .iter()
            .cloned()
            .map(Dynamic::from)
            .collect::<Array>()
    }

    // --- FlowStep Functions ---
    #[rhai_fn(global)]
    pub fn new_flow_step(step_order_i64: i64) -> Dynamic {
        match id_from_i64_to_u32(step_order_i64) {
            Ok(step_order) => {
                let mut flow_step = FlowStep::default();
                flow_step.step_order = step_order;
                Dynamic::from(flow_step)
            }
            Err(err) => Dynamic::from(err.to_string()),
        }
    }

    /// Sets the flow step description
    #[rhai_fn(name = "description", return_raw, global, pure)]
    pub fn flow_step_description(
        step: &mut RhaiFlowStep,
        description: String,
    ) -> Result<RhaiFlowStep, Box<EvalAltResult>> {
        let owned_step = mem::replace(step, FlowStep::default()); // Use Default trait
        *step = owned_step.description(description);
        Ok(step.clone())
    }

    /// Sets the flow step status
    #[rhai_fn(name = "status", return_raw, global, pure)]
    pub fn flow_step_status(
        step: &mut RhaiFlowStep,
        status: String,
    ) -> Result<RhaiFlowStep, Box<EvalAltResult>> {
        let owned_step = mem::replace(step, FlowStep::default()); // Use Default trait
        *step = owned_step.status(status);
        Ok(step.clone())
    }

    // FlowStep Getters
    #[rhai_fn(get = "id", pure)]
    pub fn get_step_id(step: &mut RhaiFlowStep) -> i64 {
        step.base_data.id as i64
    }
    #[rhai_fn(get = "created_at", pure)]
    pub fn get_step_created_at(step: &mut RhaiFlowStep) -> i64 {
        step.base_data.created_at
    }
    #[rhai_fn(get = "modified_at", pure)]
    pub fn get_step_modified_at(step: &mut RhaiFlowStep) -> i64 {
        step.base_data.modified_at
    }
    #[rhai_fn(get = "description", pure)]
    pub fn get_step_description(step: &mut RhaiFlowStep) -> Dynamic {
        match &step.description {
            Some(desc) => Dynamic::from(desc.clone()),
            None => Dynamic::UNIT,
        }
    }
    #[rhai_fn(get = "step_order", pure)]
    pub fn get_step_order(step: &mut RhaiFlowStep) -> i64 {
        step.step_order as i64
    }
    #[rhai_fn(get = "status", pure)]
    pub fn get_step_status(step: &mut RhaiFlowStep) -> String {
        step.status.clone()
    }

    // --- SignatureRequirement Functions ---
    /// Create a new signature requirement
    #[rhai_fn(global)]
    pub fn new_signature_requirement(
        flow_step_id_i64: i64,
        public_key: String,
        message: String,
    ) -> Dynamic {
        match id_from_i64_to_u32(flow_step_id_i64) {
            Ok(flow_step_id) => {
                let mut signature_requirement = SignatureRequirement::default();
                signature_requirement.flow_step_id = flow_step_id;
                signature_requirement.public_key = public_key;
                signature_requirement.message = message;
                Dynamic::from(signature_requirement)
            }
            Err(err) => Dynamic::from(err.to_string()),
        }
    }

    /// Sets the signed_by field
    #[rhai_fn(name = "signed_by", return_raw, global, pure)]
    pub fn signature_requirement_signed_by(
        sr: &mut RhaiSignatureRequirement,
        signed_by: String,
    ) -> Result<RhaiSignatureRequirement, Box<EvalAltResult>> {
        let owned_sr = mem::replace(sr, SignatureRequirement::default()); // Use Default trait
        *sr = owned_sr.signed_by(signed_by);
        Ok(sr.clone())
    }

    /// Sets the signature field
    #[rhai_fn(name = "signature", return_raw, global, pure)]
    pub fn signature_requirement_signature(
        sr: &mut RhaiSignatureRequirement,
        signature: String,
    ) -> Result<RhaiSignatureRequirement, Box<EvalAltResult>> {
        let owned_sr = mem::replace(sr, SignatureRequirement::default()); // Use Default trait
        *sr = owned_sr.signature(signature);
        Ok(sr.clone())
    }

    /// Sets the status field
    #[rhai_fn(name = "status", return_raw, global, pure)]
    pub fn signature_requirement_status(
        sr: &mut RhaiSignatureRequirement,
        status: String,
    ) -> Result<RhaiSignatureRequirement, Box<EvalAltResult>> {
        let owned_sr = mem::replace(sr, SignatureRequirement::default()); // Use Default trait
        *sr = owned_sr.status(status);
        Ok(sr.clone())
    }

    // SignatureRequirement Getters
    #[rhai_fn(get = "id", pure)]
    pub fn get_sr_id(sr: &mut RhaiSignatureRequirement) -> i64 {
        sr.base_data.id as i64
    }
    #[rhai_fn(get = "created_at", pure)]
    pub fn get_sr_created_at(sr: &mut RhaiSignatureRequirement) -> i64 {
        sr.base_data.created_at
    }
    #[rhai_fn(get = "modified_at", pure)]
    pub fn get_sr_modified_at(sr: &mut RhaiSignatureRequirement) -> i64 {
        sr.base_data.modified_at
    }
    #[rhai_fn(get = "flow_step_id", pure)]
    pub fn get_sr_flow_step_id(sr: &mut RhaiSignatureRequirement) -> i64 {
        sr.flow_step_id as i64
    }
    #[rhai_fn(get = "public_key", pure)]
    pub fn get_sr_public_key(sr: &mut RhaiSignatureRequirement) -> String {
        sr.public_key.clone()
    }
    #[rhai_fn(get = "message", pure)]
    pub fn get_sr_message(sr: &mut RhaiSignatureRequirement) -> String {
        sr.message.clone()
    }
    #[rhai_fn(get = "signed_by", pure)]
    pub fn get_sr_signed_by(sr: &mut RhaiSignatureRequirement) -> Dynamic {
        match &sr.signed_by {
            Some(signed_by) => Dynamic::from(signed_by.clone()),
            None => Dynamic::UNIT,
        }
    }
    #[rhai_fn(get = "signature", pure)]
    pub fn get_sr_signature(sr: &mut RhaiSignatureRequirement) -> Dynamic {
        match &sr.signature {
            Some(signature) => Dynamic::from(signature.clone()),
            None => Dynamic::UNIT,
        }
    }
    #[rhai_fn(get = "status", pure)]
    pub fn get_sr_status(sr: &mut RhaiSignatureRequirement) -> String {
        sr.status.clone()
    }
}

/// Register the flow module with the Rhai engine
pub fn register_flow_rhai_module(engine: &mut Engine, db: Arc<OurDB>) {
    // Create a module for database functions
    let mut db_module = Module::new();

    // Flow database functions
    let db_clone = Arc::clone(&db);
    db_module.set_native_fn(
        "save_flow",
        move |flow: Flow| -> Result<Flow, Box<EvalAltResult>> {
            // Use the Collection trait method directly
            let result = db_clone.set(&flow).map_err(|e| {
                Box::new(EvalAltResult::ErrorRuntime(
                    format!("DB Error save_flow: {:?}", e).into(),
                    Position::NONE,
                ))
            })?;

            // Return the updated flow with the correct ID
            Ok(result.1)
        },
    );

    let db_clone = Arc::clone(&db);
    db_module.set_native_fn(
        "get_flow_by_id",
        move |id_i64: INT| -> Result<Flow, Box<EvalAltResult>> {
            let id_u32 = id_from_i64_to_u32(id_i64)?;
            // Use the Collection trait method directly
            db_clone
                .get_by_id(id_u32)
                .map_err(|e| {
                    Box::new(EvalAltResult::ErrorRuntime(
                        format!("DB Error get_flow_by_id: {:?}", e).into(),
                        Position::NONE,
                    ))
                })?
                .ok_or_else(|| {
                    Box::new(EvalAltResult::ErrorRuntime(
                        format!("Flow with ID {} not found", id_u32).into(),
                        Position::NONE,
                    ))
                })
        },
    );

    let db_clone = Arc::clone(&db);
    db_module.set_native_fn(
        "delete_flow",
        move |id_i64: INT| -> Result<(), Box<EvalAltResult>> {
            let id_u32 = id_from_i64_to_u32(id_i64)?;
            // Use the Collection trait method directly
            let collection = db_clone.collection::<Flow>().map_err(|e| {
                Box::new(EvalAltResult::ErrorRuntime(
                    format!("Failed to get flow collection: {:?}", e).into(),
                    Position::NONE,
                ))
            })?;
            collection.delete_by_id(id_u32).map_err(|e| {
                Box::new(EvalAltResult::ErrorRuntime(
                    format!("Failed to delete Flow (ID: {}): {:?}", id_u32, e).into(),
                    Position::NONE,
                ))
            })
        },
    );

    let db_clone = Arc::clone(&db);
    db_module.set_native_fn(
        "list_flows",
        move || -> Result<Dynamic, Box<EvalAltResult>> {
            let collection = db_clone.collection::<Flow>().map_err(|e| {
                Box::new(EvalAltResult::ErrorRuntime(
                    format!("Failed to get flow collection: {:?}", e).into(),
                    Position::NONE,
                ))
            })?;
            let flows = collection.get_all().map_err(|e| {
                Box::new(EvalAltResult::ErrorRuntime(
                    format!("Failed to get all flows: {:?}", e).into(),
                    Position::NONE,
                ))
            })?;
            let mut array = Array::new();
            for flow in flows {
                array.push(Dynamic::from(flow));
            }
            Ok(Dynamic::from(array))
        },
    );

    // FlowStep database functions
    let db_clone = Arc::clone(&db);
    db_module.set_native_fn(
        "save_flow_step",
        move |step: FlowStep| -> Result<FlowStep, Box<EvalAltResult>> {
            // Use the Collection trait method directly
            let result = db_clone.set(&step).map_err(|e| {
                Box::new(EvalAltResult::ErrorRuntime(
                    format!("DB Error save_flow_step: {:?}", e).into(),
                    Position::NONE,
                ))
            })?;

            // Return the updated flow step with the correct ID
            Ok(result.1)
        },
    );

    let db_clone = Arc::clone(&db);
    db_module.set_native_fn(
        "get_flow_step_by_id",
        move |id_i64: INT| -> Result<FlowStep, Box<EvalAltResult>> {
            let id_u32 = id_from_i64_to_u32(id_i64)?;
            // Use the Collection trait method directly
            db_clone
                .get_by_id(id_u32)
                .map_err(|e| {
                    Box::new(EvalAltResult::ErrorRuntime(
                        format!("DB Error get_flow_step_by_id: {:?}", e).into(),
                        Position::NONE,
                    ))
                })?
                .ok_or_else(|| {
                    Box::new(EvalAltResult::ErrorRuntime(
                        format!("FlowStep with ID {} not found", id_u32).into(),
                        Position::NONE,
                    ))
                })
        },
    );

    let db_clone = Arc::clone(&db);
    db_module.set_native_fn(
        "delete_flow_step",
        move |id_i64: INT| -> Result<(), Box<EvalAltResult>> {
            let id_u32 = id_from_i64_to_u32(id_i64)?;
            // Use the Collection trait method directly
            let collection = db_clone.collection::<FlowStep>().map_err(|e| {
                Box::new(EvalAltResult::ErrorRuntime(
                    format!("Failed to get flow step collection: {:?}", e).into(),
                    Position::NONE,
                ))
            })?;
            collection.delete_by_id(id_u32).map_err(|e| {
                Box::new(EvalAltResult::ErrorRuntime(
                    format!("Failed to delete FlowStep (ID: {}): {:?}", id_u32, e).into(),
                    Position::NONE,
                ))
            })
        },
    );

    let db_clone = Arc::clone(&db);
    db_module.set_native_fn(
        "list_flow_steps",
        move || -> Result<Dynamic, Box<EvalAltResult>> {
            let collection = db_clone.collection::<FlowStep>().map_err(|e| {
                Box::new(EvalAltResult::ErrorRuntime(
                    format!("Failed to get flow step collection: {:?}", e).into(),
                    Position::NONE,
                ))
            })?;
            let steps = collection.get_all().map_err(|e| {
                Box::new(EvalAltResult::ErrorRuntime(
                    format!("Failed to get all flow steps: {:?}", e).into(),
                    Position::NONE,
                ))
            })?;
            let mut array = Array::new();
            for step in steps {
                array.push(Dynamic::from(step));
            }
            Ok(Dynamic::from(array))
        },
    );

    // SignatureRequirement database functions
    let db_clone = Arc::clone(&db);
    db_module.set_native_fn(
        "save_signature_requirement",
        move |sr: SignatureRequirement| -> Result<SignatureRequirement, Box<EvalAltResult>> {
            // Use the Collection trait method directly
            let result = db_clone.set(&sr).map_err(|e| {
                Box::new(EvalAltResult::ErrorRuntime(
                    format!("DB Error save_signature_requirement: {:?}", e).into(),
                    Position::NONE,
                ))
            })?;

            // Return the updated signature requirement with the correct ID
            Ok(result.1)
        },
    );

    let db_clone = Arc::clone(&db);
    db_module.set_native_fn(
        "get_signature_requirement_by_id",
        move |id_i64: INT| -> Result<SignatureRequirement, Box<EvalAltResult>> {
            let id_u32 = id_from_i64_to_u32(id_i64)?;
            // Use the Collection trait method directly
            db_clone
                .get_by_id(id_u32)
                .map_err(|e| {
                    Box::new(EvalAltResult::ErrorRuntime(
                        format!("DB Error get_signature_requirement_by_id: {:?}", e).into(),
                        Position::NONE,
                    ))
                })?
                .ok_or_else(|| {
                    Box::new(EvalAltResult::ErrorRuntime(
                        format!("SignatureRequirement with ID {} not found", id_u32).into(),
                        Position::NONE,
                    ))
                })
        },
    );

    let db_clone = Arc::clone(&db);
    db_module.set_native_fn(
        "delete_signature_requirement",
        move |id_i64: INT| -> Result<(), Box<EvalAltResult>> {
            let id_u32 = id_from_i64_to_u32(id_i64)?;
            // Use the Collection trait method directly
            let collection = db_clone.collection::<SignatureRequirement>().map_err(|e| {
                Box::new(EvalAltResult::ErrorRuntime(
                    format!("Failed to get signature requirement collection: {:?}", e).into(),
                    Position::NONE,
                ))
            })?;
            collection.delete_by_id(id_u32).map_err(|e| {
                Box::new(EvalAltResult::ErrorRuntime(
                    format!(
                        "Failed to delete SignatureRequirement (ID: {}): {:?}",
                        id_u32, e
                    )
                    .into(),
                    Position::NONE,
                ))
            })
        },
    );

    let db_clone = Arc::clone(&db);
    db_module.set_native_fn(
        "list_signature_requirements",
        move || -> Result<Dynamic, Box<EvalAltResult>> {
            let collection = db_clone.collection::<SignatureRequirement>().map_err(|e| {
                Box::new(EvalAltResult::ErrorRuntime(
                    format!("Failed to get signature requirement collection: {:?}", e).into(),
                    Position::NONE,
                ))
            })?;
            let srs = collection.get_all().map_err(|e| {
                Box::new(EvalAltResult::ErrorRuntime(
                    format!("Failed to get all signature requirements: {:?}", e).into(),
                    Position::NONE,
                ))
            })?;
            let mut array = Array::new();
            for sr in srs {
                array.push(Dynamic::from(sr));
            }
            Ok(Dynamic::from(array))
        },
    );

    // Register the database module globally
    engine.register_static_module("db", db_module.into());

    // Register the flow module using exported_module! macro
    let module = exported_module!(rhai_flow_module);
    engine.register_global_module(module.into());

    println!("Flow Rhai module registered.");
}