fabro/lib/components/fabro-workflow/src/test_support.rs

use std::collections::{BTreeMap, HashMap};
use std::hash::{Hash, Hasher};
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;

use fabro_agent::Sandbox;
use fabro_auth::{CredentialSource, test_support as auth_test_support};
use fabro_graphviz::graph::Graph as GvGraph;
use fabro_interview::AutoApproveInterviewer;
use fabro_model::Catalog;
#[cfg(feature = "test-support")]
use fabro_model::ProviderId;
use fabro_store::{ArtifactStore, RunProjection, test_support as store_test_support};
use object_store::local::LocalFileSystem;

use crate::artifact_upload::ArtifactSink;
use crate::error::{Error, Result};
use crate::event::{Emitter, Event, StoreProgressLogger, append_event};
use crate::handler::HandlerRegistry;
use crate::outcome::Outcome;
use crate::pipeline;
use crate::pipeline::types::{Executed, Initialized};
use crate::pipeline::{billing_from_projection, build_terminal_event};
use crate::records::Checkpoint;
use crate::run_metadata::RunMetadataRuntime;
use crate::run_options::RunOptions;
use crate::sandbox_git_runtime::SandboxGitRuntime;
use crate::services::{EngineServices, RunLocations, RunServices};
use crate::stage_execution::StageExecutionTracker;

#[cfg(feature = "test-support")]
pub(crate) fn test_configured_provider_ids(
    catalog: &Catalog,
    configured_provider_ids: Vec<ProviderId>,
    assume_ready: bool,
) -> Vec<ProviderId> {
    if assume_ready {
        catalog.all_provider_ids().into_iter().collect()
    } else {
        configured_provider_ids
    }
}

/// These helpers stop at EXECUTE, so they emit the terminal event here to
/// keep test consumers seeing the same end-of-run signal as production
/// (FINALIZE).
///
/// The first flush is needed because `StoreProgressLogger` forwards events
/// through an mpsc channel — without it, billing would read from a stale
/// checkpoint. The second flush ensures the just-emitted terminal event is
/// persisted before tests reopen the run store.
async fn execute_and_emit_terminal(initialized: InitializedState) -> Executed {
    let executed = Box::pin(pipeline::execute(initialized.initialized)).await;
    initialized
        .store_logger
        .flush()
        .await
        .expect("test run events should persist");
    let state = executed.engine.run.run_store.state().await.ok();
    let billing = state.as_ref().and_then(billing_from_projection);
    let event = build_terminal_event(
        &executed.outcome,
        fabro_types::RunTiming::wall_only(executed.wall_time_ms),
        0,
        None,
        None,
        None,
        billing,
    );
    executed.engine.run.emitter.emit(&event);
    initialized
        .store_logger
        .flush()
        .await
        .expect("test run events should persist");
    executed
}

/// Construct a fully-populated `BilledModelUsage` for tests. Centralised so
/// callers don't keep rebuilding the same JSON skeleton.
#[must_use]
pub fn test_usage(
    model_id: &str,
    input_tokens: i64,
    output_tokens: i64,
) -> fabro_types::BilledModelUsage {
    serde_json::from_value(serde_json::json!({
        "input": {
            "usage": {
                "model": {
                    "provider": "openai",
                    "model_id": model_id
                },
                "tokens": {
                    "input_tokens": input_tokens,
                    "output_tokens": output_tokens
                }
            },
            "facts": { "algorithm": "openai" }
        },
        "total_usd_micros": input_tokens + output_tokens
    }))
    .expect("test_usage JSON must deserialise")
}

/// Append the `RunStartRequested → RunRunnable → RunStarting → RunRunning`
/// sequence so subsequent calls observe the run as live.
pub async fn mark_run_running(run_store: &fabro_store::RunDatabase, run_id: &fabro_types::RunId) {
    append_event(run_store, run_id, &Event::RunStartRequested {
        resume: false,
        actor:  None,
    })
    .await
    .expect("seed run.start_requested");
    append_event(run_store, run_id, &Event::RunRunnable {
        source: fabro_types::RunRunnableSource::StartRequested,
        actor:  None,
    })
    .await
    .expect("seed run.runnable");
    append_event(run_store, run_id, &Event::RunStarting)
        .await
        .expect("seed run.starting");
    append_event(run_store, run_id, &Event::RunRunning)
        .await
        .expect("seed run.running");
}

/// Record every event the emitter publishes, for assertions after a run.
pub fn collect_events(emitter: &Emitter) -> Arc<std::sync::Mutex<Vec<fabro_types::RunEvent>>> {
    let events = Arc::new(std::sync::Mutex::new(Vec::new()));
    let captured = Arc::clone(&events);
    emitter.on_event(move |event| captured.lock().unwrap().push(event.clone()));
    events
}

pub fn test_store_dir(run_dir: &std::path::Path) -> PathBuf {
    let mut hasher = std::collections::hash_map::DefaultHasher::new();
    std::process::id().hash(&mut hasher);
    run_dir.hash(&mut hasher);
    std::env::temp_dir()
        .join("fabro-test-run-stores")
        .join(format!("{:016x}", hasher.finish()))
}

struct InitializedOptions {
    hook_runner: Option<Arc<fabro_hooks::HookRunner>>,
    env:         HashMap<String, String>,
    checkpoint:  Option<Checkpoint>,
    llm_source:  Option<Arc<dyn CredentialSource>>,
}

struct InitializedState {
    initialized:  Initialized,
    store_logger: StoreProgressLogger,
}

fn bound_emitter(run_id: fabro_types::RunId, observer: &Arc<Emitter>) -> Arc<Emitter> {
    let emitter = Arc::new(Emitter::new(run_id));
    let observer_clone = Arc::clone(observer);
    emitter.on_event(move |event| observer_clone.dispatch_run_event(event));
    emitter
}

async fn initialized(
    registry: HandlerRegistry,
    emitter: Arc<Emitter>,
    sandbox: Arc<dyn Sandbox>,
    graph: &GvGraph,
    run_options: &RunOptions,
    options: InitializedOptions,
) -> InitializedState {
    std::fs::create_dir_all(&run_options.run_dir).expect("failed to create run dir");
    let store_dir = test_store_dir(&run_options.run_dir);
    let _ = std::fs::remove_dir_all(&store_dir);
    for database_path in [
        store_test_support::test_blob_store_path(&store_dir),
        store_test_support::test_run_summary_store_path(&store_dir),
    ] {
        for suffix in ["", "-wal", "-shm"] {
            let mut sibling = database_path.clone().into_os_string();
            sibling.push(suffix);
            let _ = std::fs::remove_file(sibling);
        }
    }
    std::fs::create_dir_all(&store_dir).expect("failed to create local test run store dir");
    let store = Arc::new(store_test_support::test_database_at(
        Arc::new(
            LocalFileSystem::new_with_prefix(&store_dir)
                .expect("failed to create local test run store"),
        ),
        "",
        Duration::from_millis(1),
        None,
        &store_dir,
    ));
    let inner_store = store
        .create_run(&run_options.run_id)
        .await
        .expect("failed to create slate-backed test run store");
    let run_store = inner_store;
    append_event(&run_store, &run_options.run_id, &Event::RunCreated {
        run_id:              run_options.run_id,
        title:               None,
        settings:            serde_json::to_value(&run_options.settings)
            .expect("failed to serialize settings"),
        graph:               serde_json::to_value(graph).expect("failed to serialize graph"),
        workflow_source:     None,
        labels:              run_options
            .labels
            .clone()
            .into_iter()
            .collect::<BTreeMap<_, _>>(),
        source_directory:    Some(sandbox.working_directory().to_string()),
        workflow_slug:       run_options.workflow_slug.clone(),
        workflow_version_id: None,
        target:              None,
        automation:          None,
        provenance:          fabro_types::RunProvenance {
            server:  None,
            client:  None,
            subject: fabro_types::Principal::System {
                system_kind: fabro_types::SystemActorKind::Engine,
            },
        },
        manifest_blob:       None,
        spec_blob:           None,
        git:                 run_options.pre_run_git.clone(),
        fork_source_ref:     run_options.fork_source_ref.clone(),
        retried_from:        None,
        parent_id:           None,
        web_url:             None,
    })
    .await
    .expect("failed to seed run.created event in run store");
    append_event(&run_store, &run_options.run_id, &Event::RunRunnable {
        source: fabro_types::RunRunnableSource::StartRequested,
        actor:  None,
    })
    .await
    .expect("failed to seed run.runnable event in run store");
    append_event(&run_store, &run_options.run_id, &Event::RunStarting)
        .await
        .expect("failed to seed run.starting event in run store");
    let emitter = bound_emitter(run_options.run_id, &emitter);
    let store_logger = StoreProgressLogger::new(run_store.clone());
    store_logger.register(emitter.as_ref());
    let artifact_store = ArtifactStore::new(
        Arc::new(
            LocalFileSystem::new_with_prefix(&store_dir)
                .expect("failed to create local test artifact store"),
        ),
        "artifacts",
    );
    let locations = RunLocations::for_sandbox(None, sandbox.as_ref(), run_options.run_dir.clone());
    InitializedState {
        initialized: Initialized {
            graph:         graph.clone(),
            source:        String::new(),
            run_options:   run_options.clone(),
            checkpoint:    options.checkpoint,
            seed_context:  None,
            on_node:       None,
            artifact_sink: Some(ArtifactSink::Store(artifact_store)),
            run_control:   None,
            engine:        Arc::new(EngineServices {
                run:             RunServices::new(
                    run_store.into(),
                    emitter,
                    sandbox,
                    options.hook_runner,
                    locations,
                    run_options.cancel_token.clone(),
                    fabro_model::ProviderId::anthropic(),
                    "claude-sonnet-4-6".to_string(),
                    options
                        .llm_source
                        .unwrap_or_else(auth_test_support::vault_only_credential_source),
                    Arc::new(Catalog::from_builtin().expect("default catalog should build")),
                    Arc::new(SandboxGitRuntime::new()),
                    Arc::new(RunMetadataRuntime::new()),
                    None,
                    StageExecutionTracker::default(),
                ),
                registry:        Arc::new(registry),
                interviewer:     Arc::new(AutoApproveInterviewer::engine()),
                base_env:        options.env,
                github_token:    None,
                inputs:          run_options.settings.run.inputs.clone(),
                dry_run:         run_options.dry_run_enabled(),
                workflow_path:   None,
                workflow_bundle: None,
            }),
            model:         String::new(),
        },
        store_logger,
    }
}

pub async fn run_graph(
    registry: HandlerRegistry,
    emitter: Arc<Emitter>,
    sandbox: Arc<dyn Sandbox>,
    graph: &GvGraph,
    run_options: &RunOptions,
) -> Result<Outcome> {
    let initialized = initialized(
        registry,
        emitter,
        sandbox,
        graph,
        run_options,
        InitializedOptions {
            hook_runner: None,
            env:         HashMap::new(),
            checkpoint:  None,
            llm_source:  None,
        },
    )
    .await;
    let executed = execute_and_emit_terminal(initialized).await;
    executed.outcome
}

pub async fn run_graph_with_state(
    registry: HandlerRegistry,
    emitter: Arc<Emitter>,
    sandbox: Arc<dyn Sandbox>,
    graph: &GvGraph,
    run_options: &RunOptions,
) -> Result<(Outcome, RunProjection)> {
    let initialized = initialized(
        registry,
        emitter,
        sandbox,
        graph,
        run_options,
        InitializedOptions {
            hook_runner: None,
            env:         HashMap::new(),
            checkpoint:  None,
            llm_source:  None,
        },
    )
    .await;
    let executed = execute_and_emit_terminal(initialized).await;
    let outcome = executed.outcome?;
    let state = executed
        .engine
        .run
        .run_store
        .state()
        .await
        .map_err(|err| Error::engine(err.to_string()))?;
    Ok((outcome, state))
}

pub async fn run_graph_with_hooks(
    registry: HandlerRegistry,
    emitter: Arc<Emitter>,
    sandbox: Arc<dyn Sandbox>,
    graph: &GvGraph,
    run_options: &RunOptions,
    hook_runner: Arc<fabro_hooks::HookRunner>,
    env: Option<HashMap<String, String>>,
) -> Result<Outcome> {
    let initialized = initialized(
        registry,
        emitter,
        sandbox,
        graph,
        run_options,
        InitializedOptions {
            hook_runner: Some(hook_runner),
            env:         env.unwrap_or_default(),
            checkpoint:  None,
            llm_source:  None,
        },
    )
    .await;
    let executed = execute_and_emit_terminal(initialized).await;
    executed.outcome
}

pub async fn run_graph_with_hooks_and_state(
    registry: HandlerRegistry,
    emitter: Arc<Emitter>,
    sandbox: Arc<dyn Sandbox>,
    graph: &GvGraph,
    run_options: &RunOptions,
    hook_runner: Arc<fabro_hooks::HookRunner>,
    env: Option<HashMap<String, String>>,
) -> Result<(Outcome, RunProjection)> {
    let initialized = initialized(
        registry,
        emitter,
        sandbox,
        graph,
        run_options,
        InitializedOptions {
            hook_runner: Some(hook_runner),
            env:         env.unwrap_or_default(),
            checkpoint:  None,
            llm_source:  None,
        },
    )
    .await;
    let executed = execute_and_emit_terminal(initialized).await;
    let outcome = executed.outcome?;
    let state = executed
        .engine
        .run
        .run_store
        .state()
        .await
        .map_err(|err| Error::engine(err.to_string()))?;
    Ok((outcome, state))
}

pub async fn run_graph_from_checkpoint(
    registry: HandlerRegistry,
    emitter: Arc<Emitter>,
    sandbox: Arc<dyn Sandbox>,
    graph: &GvGraph,
    run_options: &RunOptions,
    checkpoint: &Checkpoint,
) -> Result<Outcome> {
    let initialized = initialized(
        registry,
        emitter,
        sandbox,
        graph,
        run_options,
        InitializedOptions {
            hook_runner: None,
            env:         HashMap::new(),
            checkpoint:  Some(checkpoint.clone()),
            llm_source:  None,
        },
    )
    .await;
    let executed = execute_and_emit_terminal(initialized).await;
    executed.outcome
}

pub async fn run_graph_from_checkpoint_with_state(
    registry: HandlerRegistry,
    emitter: Arc<Emitter>,
    sandbox: Arc<dyn Sandbox>,
    graph: &GvGraph,
    run_options: &RunOptions,
    checkpoint: &Checkpoint,
) -> Result<(Outcome, RunProjection)> {
    let initialized = initialized(
        registry,
        emitter,
        sandbox,
        graph,
        run_options,
        InitializedOptions {
            hook_runner: None,
            env:         HashMap::new(),
            checkpoint:  Some(checkpoint.clone()),
            llm_source:  None,
        },
    )
    .await;
    let executed = execute_and_emit_terminal(initialized).await;
    let outcome = executed.outcome?;
    let state = executed
        .engine
        .run
        .run_store
        .state()
        .await
        .map_err(|err| Error::engine(err.to_string()))?;
    Ok((outcome, state))
}

pub async fn run_graph_with_state_and_llm_source(
    registry: HandlerRegistry,
    emitter: Arc<Emitter>,
    sandbox: Arc<dyn Sandbox>,
    graph: &GvGraph,
    run_options: &RunOptions,
    llm_source: Arc<dyn CredentialSource>,
) -> Result<(Outcome, RunProjection)> {
    let initialized = initialized(
        registry,
        emitter,
        sandbox,
        graph,
        run_options,
        InitializedOptions {
            hook_runner: None,
            env:         HashMap::new(),
            checkpoint:  None,
            llm_source:  Some(llm_source),
        },
    )
    .await;
    let executed = pipeline::execute(initialized.initialized).await;
    initialized
        .store_logger
        .flush()
        .await
        .expect("test run events should persist");
    let outcome = executed.outcome?;
    let state = executed
        .engine
        .run
        .run_store
        .state()
        .await
        .map_err(|err| Error::engine(err.to_string()))?;
    Ok((outcome, state))
}

pub struct WorkflowRunner {
    registry: std::sync::Mutex<Option<HandlerRegistry>>,
    emitter:  Arc<Emitter>,
    sandbox:  Arc<dyn Sandbox>,
}

impl WorkflowRunner {
    #[must_use]
    pub fn new(
        registry: HandlerRegistry,
        emitter: Arc<Emitter>,
        sandbox: Arc<dyn Sandbox>,
    ) -> Self {
        Self {
            registry: std::sync::Mutex::new(Some(registry)),
            emitter,
            sandbox,
        }
    }

    pub async fn run(&self, graph: &GvGraph, run_options: &RunOptions) -> Result<Outcome> {
        let registry = self
            .registry
            .lock()
            .unwrap()
            .take()
            .expect("WorkflowRunner may only be used once");
        Box::pin(run_graph(
            registry,
            Arc::clone(&self.emitter),
            Arc::clone(&self.sandbox),
            graph,
            run_options,
        ))
        .await
    }

    pub async fn run_with_state(
        &self,
        graph: &GvGraph,
        run_options: &RunOptions,
    ) -> Result<(Outcome, RunProjection)> {
        let registry = self
            .registry
            .lock()
            .unwrap()
            .take()
            .expect("WorkflowRunner may only be used once");
        Box::pin(run_graph_with_state(
            registry,
            Arc::clone(&self.emitter),
            Arc::clone(&self.sandbox),
            graph,
            run_options,
        ))
        .await
    }

    pub async fn run_with_state_and_llm_source(
        &self,
        graph: &GvGraph,
        run_options: &RunOptions,
        llm_source: Arc<dyn CredentialSource>,
    ) -> Result<(Outcome, RunProjection)> {
        let registry = self
            .registry
            .lock()
            .unwrap()
            .take()
            .expect("WorkflowRunner may only be used once");
        Box::pin(run_graph_with_state_and_llm_source(
            registry,
            Arc::clone(&self.emitter),
            Arc::clone(&self.sandbox),
            graph,
            run_options,
            llm_source,
        ))
        .await
    }

    pub async fn run_from_checkpoint(
        &self,
        graph: &GvGraph,
        run_options: &RunOptions,
        checkpoint: &Checkpoint,
    ) -> Result<Outcome> {
        let registry = self
            .registry
            .lock()
            .unwrap()
            .take()
            .expect("WorkflowRunner may only be used once");
        Box::pin(run_graph_from_checkpoint(
            registry,
            Arc::clone(&self.emitter),
            Arc::clone(&self.sandbox),
            graph,
            run_options,
            checkpoint,
        ))
        .await
    }

    pub async fn run_from_checkpoint_with_state(
        &self,
        graph: &GvGraph,
        run_options: &RunOptions,
        checkpoint: &Checkpoint,
    ) -> Result<(Outcome, RunProjection)> {
        let registry = self
            .registry
            .lock()
            .unwrap()
            .take()
            .expect("WorkflowRunner may only be used once");
        Box::pin(run_graph_from_checkpoint_with_state(
            registry,
            Arc::clone(&self.emitter),
            Arc::clone(&self.sandbox),
            graph,
            run_options,
            checkpoint,
        ))
        .await
    }
}