tinyflow-framework 0.1.1

use std::path::Path;
use std::sync::Arc;

use crate::env::stream_env::StreamEnv;
use crate::runtime::chain_job::ChainStep;
use crate::runtime::chain_segment::NamedStep;
use crate::runtime::chained_stream_task::ChainedStreamTask;
use crate::runtime::control_msg::ControlMsg;
use crate::runtime::job_runtime::ChainJobRuntime;
use crate::runtime::pending_task::PendingTask;
use crate::runtime::task_lifecycle::spawn_monitored_task;
use crate::state::open_chain_state_db;
use tinyflow_api::error::{StreamError, StreamResult};
use tinyflow_api::state::StateStore;

use super::validate::{resolve_step_name, validate_chain_job};

/// Build one pending chain task per subtask from the env's [`ChainJob`](crate::runtime::chain_job::ChainJob).
#[allow(clippy::too_many_arguments)]
pub async fn spawn_chained_tasks(
    env: &StreamEnv,
    job_state: Arc<dyn StateStore>,
    state_work_dir: Option<&Path>,
    job_runtime: &ChainJobRuntime,
    ack_tx: &tokio::sync::mpsc::Sender<ControlMsg>,
    error_tx: &tokio::sync::mpsc::Sender<StreamError>,
    chain_id: u64,
    task_id_counter: &mut u32,
) -> StreamResult<Vec<PendingTask>> {
    let (par, step_count, step_names, steps) = {
        let mut guard = env.chain_job.lock().unwrap();
        let job = guard.as_mut().ok_or_else(|| StreamError::TaskFailed {
            task_id: 0,
            reason: "no ChainJob; call add_source first".into(),
        })?;

        validate_chain_job(env, job)?;

        let par = job.parallelism.max(1);
        let step_count = job.steps.len();
        let step_names: Vec<String> = job
            .steps
            .iter()
            .map(|s| resolve_step_name(s.kind, &s.user_name))
            .collect();
        let steps = std::mem::take(&mut job.steps);
        (par, step_count, step_names, steps)
    };

    let mut pending = Vec::new();
    for pi in 0..par {
        let task_id = *task_id_counter;
        *task_id_counter += 1;

        let (ctrl_tx, ctrl_rx) = tokio::sync::mpsc::channel(64);
        job_runtime.register_control(task_id, ctrl_tx);

        let chain_label = if step_names.len() == 1 {
            step_names[0].clone()
        } else {
            format!(
                "{}..{}",
                step_names.first().unwrap(),
                step_names.last().unwrap()
            )
        };
        let task_name = format!("{}-{}", chain_label, pi);

        assert_eq!(steps.len(), step_count);
        let mut substeps: Vec<NamedStep> =
            steps.iter().map(|s: &ChainStep| (s.factory)()).collect();
        for (i, step_info) in steps.iter().enumerate() {
            substeps[i].name = resolve_step_name(step_info.kind, &step_info.user_name);
        }

        let chain_state = if env.enable_chain_state {
            let work_dir = state_work_dir.ok_or_else(|| StreamError::TaskFailed {
                task_id: 0,
                reason: "enable_chain_state requires state_work_dir at spawn".into(),
            })?;
            Some(open_chain_state_db(work_dir, pi as u32).await?)
        } else {
            None
        };

        let task = ChainedStreamTask::from_steps(
            task_id,
            task_name.clone(),
            pi as u32,
            chain_id,
            par as u32,
            substeps,
            ctrl_rx,
            ack_tx.clone(),
            Arc::clone(&job_state),
            chain_state,
        )?;

        let error_tx = error_tx.clone();
        pending.push(PendingTask::new(task_id, move || {
            spawn_monitored_task(task_id, task_name, error_tx, move || async move {
                task.run().await
            })
        }));
    }

    Ok(pending)
}