rhailib/examples/end_to_end/main.rs

use rhai::{Engine, EvalAltResult};
use rhai_client::RhaiClient;
use rhailib_worker::spawn_rhai_worker;
use std::{fs, path::Path, time::Duration};
use tokio::sync::mpsc;
use uuid::Uuid;

// Custom Rhai function for authorization
// It takes the caller's public key as an argument.
fn check_permission(caller_pk: String) -> Result<String, Box<EvalAltResult>> {
    log::info!("check_permission called with PK: {}", caller_pk);
    if caller_pk == "admin_pk" {
        Ok("Access Granted: Welcome Admin!".to_string())
    } else if caller_pk == "user_pk" {
        Ok("Limited Access: Welcome User!".to_string())
    } else {
        Ok(format!("Access Denied: Unknown public key '{}'", caller_pk))
    }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();

    let redis_url = "redis://127.0.0.1/";
    let worker_circle_pk = "auth_worker_circle".to_string();

    // 1. Create a Rhai engine and register custom functionality
    let mut engine = Engine::new();
    engine.register_fn("check_permission", check_permission);
    log::info!("Custom 'check_permission' function registered with Rhai engine.");

    // 2. Spawn the Rhai worker
    let (shutdown_tx, shutdown_rx) = mpsc::channel(1);
    let worker_handle = tokio::spawn(spawn_rhai_worker(
        0, // worker_id
        worker_circle_pk.clone(),
        engine,
        redis_url.to_string(),
        shutdown_rx,
        false, // use_sentinel
    ));
    log::info!("Rhai worker spawned for circle: {}", worker_circle_pk);

    // Give the worker a moment to start up
    tokio::time::sleep(Duration::from_secs(1)).await;

    // 3. Create a Rhai client
    let client = RhaiClient::new(redis_url)?;
    log::info!("Rhai client created.");

    // 4. Load the Rhai script content
    let script_path_str = "examples/end_to_end/auth_script.rhai"; // Relative to Cargo.toml / rhailib root
    let script_content = match fs::read_to_string(script_path_str) {
        Ok(content) => content,
        Err(e) => {
            log::error!("Failed to read script file '{}': {}", script_path_str, e);
            // Attempt to read from an alternative path if run via `cargo run --example`
            // where current dir might be the crate root.
            let alt_script_path = Path::new(file!()).parent().unwrap().join("auth_script.rhai");
            log::info!("Attempting alternative script path: {:?}", alt_script_path);
            fs::read_to_string(&alt_script_path)?
        }
    };
    log::info!("Loaded script content from '{}'", script_path_str);

    // Define different caller public keys
    let admin_caller_pk = "admin_pk".to_string();
    let user_caller_pk = "user_pk".to_string();
    let unknown_caller_pk = "unknown_pk".to_string();

    let callers = vec![
        ("Admin", admin_caller_pk),
        ("User", user_caller_pk),
        ("Unknown", unknown_caller_pk),
    ];

    for (caller_name, caller_pk) in callers {
        let task_id = Uuid::new_v4().to_string();
        log::info!(
            "Submitting script for caller '{}' (PK: {}) with task_id: {}",
            caller_name,
            caller_pk,
            task_id
        );

        match client
            .submit_script_and_await_result(
                &worker_circle_pk,
                task_id.clone(),        // task_id (UUID) first
                script_content.clone(), // script_content second
                Duration::from_secs(10),
                Some(caller_pk.clone()), // This is the CALLER_PUBLIC_KEY
            )
            .await
        {
            Ok(details) => {
                log::info!(
                    "Task {} for caller '{}' (PK: {}) completed. Status: {}, Output: {:?}, Error: {:?}",
                    task_id,
                    caller_name,
                    caller_pk,
                    details.status,
                    details.output,
                    details.error
                );
                // Basic assertion for expected output
                if caller_pk == "admin_pk" {
                    assert_eq!(details.output, Some("Access Granted: Welcome Admin!".to_string()));
                } else if caller_pk == "user_pk" {
                    assert_eq!(details.output, Some("Limited Access: Welcome User!".to_string()));
                }
            }
            Err(e) => {
                log::error!(
                    "Task {} for caller '{}' (PK: {}) failed: {}",
                    task_id,
                    caller_name,
                    caller_pk,
                    e
                );
            }
        }
        tokio::time::sleep(Duration::from_millis(100)).await; // Small delay between submissions
    }

    // 5. Shutdown the worker (optional, could also let it run until program exits)
    log::info!("Signaling worker to shutdown...");
    let _ = shutdown_tx.send(()).await;
    if let Err(e) = worker_handle.await {
        log::error!("Worker task panicked or encountered an error: {:?}", e);
    }
    log::info!("Worker shutdown complete.");

    Ok(())
}