tinyflow-framework 0.1.1

use std::path::PathBuf;

use crate::stream::data_stream::DataStream;
use std::sync::{Arc, Mutex};
use tinyflow_api::functions::*;

use crate::runtime::chain_job::ChainJob;

type PendingInitialState = Vec<(String, String, Vec<u8>)>;

/// Job environment: DSL chain construction panics on programmer errors.
/// Runtime entry: [`crate::stream::chain_execute::execute`].
#[derive(Clone)]
pub struct StreamEnv {
    pub(crate) default_parallelism: usize,
    pub(crate) checkpoint_interval_ms: u64,
    pub chain_job: Arc<Mutex<Option<ChainJob>>>,
    /// Root work directory for `job/state.db`; required when checkpoint or chain_state is enabled.
    pub(crate) state_work_dir: Option<PathBuf>,
    /// When true, each chain subtask opens `chains/{i}/chain.db` (requires a state work dir at run).
    /// Default: **true** — call [`Self::set_enable_chain_state(false)`] to skip per-subtask DB I/O.
    pub(crate) enable_chain_state: bool,
    pub(crate) pending_job_state: Arc<Mutex<PendingInitialState>>,
}

impl Default for StreamEnv {
    fn default() -> Self {
        Self::new()
    }
}

impl StreamEnv {
    pub fn new() -> Self {
        Self {
            default_parallelism: 1,
            checkpoint_interval_ms: 0,
            chain_job: Arc::new(Mutex::new(None)),
            state_work_dir: None,
            enable_chain_state: true,
            pending_job_state: Arc::new(Mutex::new(Vec::new())),
        }
    }

    pub fn set_state_work_dir(&mut self, path: impl Into<PathBuf>) {
        self.state_work_dir = Some(path.into());
    }

    /// Enable per-subtask `chain_state` (`chains/{subtask}/chain.db`) at task spawn.
    pub fn set_enable_chain_state(&mut self, enable: bool) {
        self.enable_chain_state = enable;
    }

    /// Buffer an initial job-state entry that will be flushed into `job_state` before the stream starts.
    pub fn add_initial_job_state(&self, namespace: &str, key: &str, value: &[u8]) {
        self.pending_job_state.lock().unwrap().push((
            namespace.to_string(),
            key.to_string(),
            value.to_vec(),
        ));
    }

    pub fn set_parallelism(&mut self, p: usize) {
        self.default_parallelism = p;
    }

    pub fn set_checkpoint_interval(&mut self, ms: u64) {
        self.checkpoint_interval_ms = ms;
    }

    /// Run the built chain until Ctrl+C, task failure, or checkpoint manager exit.
    pub async fn execute(self) -> tinyflow_api::error::StreamResult<()> {
        crate::stream::chain_execute::execute(self).await
    }

    /// Test-only: whether per-subtask `chain_state` is enabled.
    #[doc(hidden)]
    pub fn test_enable_chain_state(&self) -> bool {
        self.enable_chain_state
    }

    /// Test-only: inspect built chain step count.
    #[doc(hidden)]
    pub fn test_chain_step_count(&self) -> usize {
        self.chain_job
            .lock()
            .unwrap()
            .as_ref()
            .map(|j| j.steps.len())
            .unwrap_or(0)
    }

    #[doc(hidden)]
    pub fn test_chain_frozen(&self) -> bool {
        self.chain_job
            .lock()
            .unwrap()
            .as_ref()
            .map(|j| j.frozen)
            .unwrap_or(false)
    }

    /// Test-only: simulate terminal sink freezing the chain.
    #[doc(hidden)]
    pub fn test_freeze_chain(&self) {
        if let Some(job) = self.chain_job.lock().unwrap().as_mut() {
            job.frozen = true;
        }
    }

    pub fn add_source<T: Send + 'static, S: SourceFunction<T> + Clone + Send + 'static>(
        &self,
        source: S,
    ) -> DataStream<T> {
        self.add_source_with_factory(move || source.clone())
    }

    pub fn add_source_with_factory<T: Send + 'static, S: SourceFunction<T> + Send + 'static>(
        &self,
        factory: impl Fn() -> S + Send + 'static,
    ) -> DataStream<T> {
        let mut job_guard = self.chain_job.lock().unwrap();
        if job_guard.is_some() {
            panic!("StreamEnv supports only one ChainJob; add_source already called");
        }
        let par = self.default_parallelism.max(1);
        let factory = std::sync::Arc::new(std::sync::Mutex::new(factory));
        let chain_factory =
            crate::runtime::chain_segment::source_chain_factory(std::sync::Arc::clone(&factory));
        let mut job = ChainJob::new(par);
        job.push_step_with_types(
            crate::runtime::chain_job::ChainStep {
                kind: crate::runtime::chain_job::ChainOperatorKind::Source,
                user_name: None,
                factory: chain_factory,
            },
            None,
            std::any::TypeId::of::<T>(),
        );
        *job_guard = Some(job);
        DataStream::new(self.clone(), 0)
    }
}