foundation_core 0.0.3

#![allow(clippy::type_complexity)]

use std::sync::{Arc, Mutex};

use crate::valtron::{
    split_thread_count, task::TaskIterator, BackgroundJobRegistry, ExecutionAction, GenericResult,
    TaskReadyResolver, TaskStatusMapper,
};

use crate::synca::{LockSignal, OnSignal};
pub use crate::valtron::multi::threads::*;
use crate::valtron::SharedTaskQueue;

use crate::valtron::executors::background::DEFAULT_BG_YIELD_DURATION;

/// Global registry - stores the `ThreadRegistry` Arc.
/// Uses Mutex<Option> to allow resetting between tests.
static REGISTRY: Mutex<Option<Arc<ThreadRegistry>>> = Mutex::new(None);

/// Global background job registry.
/// Uses Mutex<Option> to allow resetting between tests.
static BG_REGISTRY: Mutex<Option<Arc<BackgroundJobRegistry>>> = Mutex::new(None);

/// Handle for spawning tasks into the shared queue.
///
/// This is the new return type of `get_pool()`. It provides the same
/// spawn API surface as the old `&ThreadPool` by constructing
/// `ThreadPoolTaskBuilder` directly with the shared queue and latch.
#[derive(Clone)]
pub struct LocalPoolHandle {
    shared_tasks: SharedTaskQueue,
    latch: Arc<LockSignal>,
    kill_signal: Arc<OnSignal>,
    yielders: SharedThreadYielders,
}

impl LocalPoolHandle {
    fn new(registry: &Arc<ThreadRegistry>) -> Self {
        Self {
            shared_tasks: registry.shared_tasks(),
            latch: registry.latch(),
            kill_signal: registry.kill_signal(),
            yielders: registry.yielders(),
        }
    }

    /// Create a task builder for scheduling work.
    #[must_use]
    pub fn spawn<Task, Action>(
        &self,
    ) -> ThreadPoolTaskBuilder<
        Task::Ready,
        Task::Pending,
        Task::Spawner,
        Box<dyn TaskStatusMapper<Task::Ready, Task::Pending, Task::Spawner> + Send + 'static>,
        Box<dyn TaskReadyResolver<Task::Spawner, Task::Ready, Task::Pending> + Send + 'static>,
        Task,
    >
    where
        Task::Ready: Send + 'static,
        Task::Pending: Send + 'static,
        Task: TaskIterator<Spawner = Action> + Send + 'static,
        Action: ExecutionAction + Send + 'static,
    {
        ThreadPoolTaskBuilder::new(self.shared_tasks.clone(), self.latch.clone())
            .with_yielders(self.yielders.clone())
    }

    /// Create a task builder with explicit Mapper and Resolver types.
    #[must_use]
    pub fn spawn2<Task, Action, Mapper, Resolver>(
        &self,
    ) -> ThreadPoolTaskBuilder<Task::Ready, Task::Pending, Action, Mapper, Resolver, Task>
    where
        Task::Ready: Send + 'static,
        Task::Pending: Send + 'static,
        Task: TaskIterator<Spawner = Action> + Send + 'static,
        Action: ExecutionAction + Send + 'static,
        Mapper: TaskStatusMapper<Task::Ready, Task::Pending, Action> + Send + 'static,
        Resolver: TaskReadyResolver<Action, Task::Ready, Task::Pending> + Send + 'static,
    {
        ThreadPoolTaskBuilder::new(self.shared_tasks.clone(), self.latch.clone())
            .with_yielders(self.yielders.clone())
    }

    /// Signal all threads to die.
    pub fn kill(&self) {
        self.kill_signal.turn_on();
        self.latch.signal_all();
    }
}

/// [`get_pool`] returns a handle for spawning tasks.
///
/// # Panics
///
/// Panics if the thread pool has not been initialized (ensure to call `block_on` or
/// `initialize_pool` first), or if the internal registry mutex is poisoned.
pub fn get_pool() -> LocalPoolHandle {
    let registry =
        REGISTRY.lock().unwrap().clone().expect(
            "Thread pool not initialized, ensure to call block_on or initialize_pool first",
        );
    LocalPoolHandle::new(&registry)
}

/// [`block_on`] instantiates the thread registry, spawns worker threads,
/// runs the setup function, then blocks until all tasks complete.
///
/// Returns a `PoolGuard` - when dropped, kills all threads and blocks
/// until they're cleaned up. No Ctrl-C handler is registered automatically.
pub fn block_on<F>(seed_from_rng: u64, thread_num: Option<usize>, setup: F) -> PoolGuard
where
    F: FnOnce(LocalPoolHandle),
{
    let guard = initialize_pool(seed_from_rng, thread_num);
    let handle = get_pool();
    setup(handle);
    tracing::debug!("Initialize and call WaitGroup::wait");
    guard.waitgroup().wait();
    guard
}

/// `initialize_pool` creates a `ThreadRegistry` and `BackgroundJobRegistry`,
/// splitting the available threads between them.
///
/// Returns a `PoolGuard` — when dropped, signals all threads to die,
/// waits for them to complete via `WaitGroup`, joins all handles,
/// and clears the global registries to prevent test interference.
///
/// The caller is responsible for signal handling via the `PoolGuard`.
pub fn initialize_pool(seed_for_rng: u64, user_thread_num: Option<usize>) -> PoolGuard {
    let thread_num = match user_thread_num {
        None => get_allocatable_thread_count(),
        // The multi pool splits threads into background + task workers and needs
        // ≥ 2 TASK workers (`ThreadRegistry`), which means ≥ 3 total once the
        // background worker is carved off. Clamp a smaller request up rather than
        // panic — callers (and the `#[valtron]` macros) pass a plain count.
        Some(num) => num.max(3),
    };

    let (task_threads, bg_threads) = split_thread_count(thread_num);

    tracing::debug!(
        "Splitting {thread_num} threads: {task_threads} for tasks, {bg_threads} for background jobs"
    );

    let registry = Arc::new(ThreadRegistry::with_seed_and_threads(
        seed_for_rng,
        task_threads,
    ));

    // Spawn task worker threads
    for _ in 0..task_threads {
        registry
            .spawn_worker()
            .expect("Failed to spawn worker thread");
    }

    // Create background job registry sharing the kill signal
    let bg_registry = BackgroundJobRegistry::new(
        bg_threads,
        registry.kill_signal(),
        DEFAULT_BG_YIELD_DURATION,
    );

    *REGISTRY.lock().unwrap() = Some(registry.clone());
    *BG_REGISTRY.lock().unwrap() = Some(bg_registry.clone());

    PoolGuard::with_bg_registry(registry, bg_registry)
}

/// Submit a blocking closure for execution on the background job pool.
///
/// WHY: Provides managed background thread execution without spawning new OS threads
/// WHAT: Pushes the closure onto the `BackgroundJobRegistry` queue
/// HOW: Acquires the global `BG_REGISTRY` and calls `submit()`
///
/// # Errors
///
/// Returns an error if the background registry is not initialized or the queue is closed.
///
/// # Panics
///
/// Panics if the background job pool has not been initialized.
pub fn run_background_job(job: impl FnOnce() + Send + 'static) -> GenericResult<()> {
    let bg_registry = BG_REGISTRY
        .lock()
        .unwrap()
        .clone()
        .expect("Background job pool not initialized, ensure to call initialize_pool first");
    bg_registry.submit(job)
}

/// [`spawn`] provides a builder which allows you to build out
/// the underlying tasks to be scheduled into the global queue.
///
/// Uses the global registry's shared queue and latch.
#[must_use]
pub fn spawn<Task, Action>() -> ThreadPoolTaskBuilder<
    Task::Ready,
    Task::Pending,
    Task::Spawner,
    Box<dyn TaskStatusMapper<Task::Ready, Task::Pending, Task::Spawner> + Send + 'static>,
    Box<dyn TaskReadyResolver<Task::Spawner, Task::Ready, Task::Pending> + Send + 'static>,
    Task,
>
where
    Task::Ready: Send + 'static,
    Task::Pending: Send + 'static,
    Task: TaskIterator<Spawner = Action> + Send + 'static,
    Action: ExecutionAction + Send + 'static,
{
    get_pool().spawn::<Task, Action>()
}

/// [`spawn2`] provides a builder which allows you to build out
/// the underlying tasks with explicit Mapper and Resolver types.
#[must_use]
pub fn spawn2<Task, Action, Mapper, Resolver>(
) -> ThreadPoolTaskBuilder<Task::Ready, Task::Pending, Action, Mapper, Resolver, Task>
where
    Task::Ready: Send + 'static,
    Task::Pending: Send + 'static,
    Task: TaskIterator<Spawner = Action> + Send + 'static,
    Action: ExecutionAction + Send + 'static,
    Mapper: TaskStatusMapper<Task::Ready, Task::Pending, Action> + Send + 'static,
    Resolver: TaskReadyResolver<Action, Task::Ready, Task::Pending> + Send + 'static,
{
    get_pool().spawn2::<Task, Action, Mapper, Resolver>()
}

#[cfg(test)]
mod multi_threaded_tests {
    use core::time;
    use std::{
        sync::{Arc, Mutex},
        thread,
    };

    use rand::RngCore;
    use serial_test::serial;
    use tracing_test::traced_test;

    use super::{block_on, get_pool, initialize_pool};
    use crate::{
        synca::mpp,
        valtron::{FnReady, NoSpawner, TaskIterator, TaskStatus},
    };

    struct DCounter(usize, Arc<Mutex<Vec<usize>>>);

    impl DCounter {
        pub fn new(val: usize, list: Arc<Mutex<Vec<usize>>>) -> Self {
            Self(val, list)
        }
    }

    impl TaskIterator for DCounter {
        type Pending = ();

        type Ready = usize;

        type Spawner = NoSpawner;

        fn next_status(
            &mut self,
        ) -> Option<crate::valtron::TaskStatus<Self::Ready, Self::Pending, Self::Spawner>> {
            let mut items = self.1.lock().unwrap();
            let item_size = items.len();

            if item_size == self.0 {
                return None;
            }

            items.push(item_size);
            let new_len = items.len();

            Some(crate::valtron::TaskStatus::Ready(new_len))
        }
    }

    struct Counter(usize, Arc<Mutex<Vec<usize>>>, mpp::Sender<()>);

    impl Counter {
        pub fn new(val: usize, list: Arc<Mutex<Vec<usize>>>) -> (Self, mpp::Receiver<()>) {
            let (sender, receiver) = mpp::bounded(10);
            (Counter(val, list, sender), receiver)
        }
    }

    impl TaskIterator for Counter {
        type Pending = ();

        type Ready = usize;

        type Spawner = NoSpawner;

        fn next_status(
            &mut self,
        ) -> Option<crate::valtron::TaskStatus<Self::Ready, Self::Pending, Self::Spawner>> {
            tracing::debug!("Counter Task is running");

            let result = {
                let mut items = self.1.lock().unwrap();
                let item_size = items.len();

                if item_size == self.0 {
                    None // signal "done" — will send after lock release
                } else {
                    items.push(item_size);
                    Some(crate::valtron::TaskStatus::Ready(items.len()))
                }
            }; // lock released here

            match result {
                None => {
                    tracing::debug!("Sending signal with sender");
                    self.2.send(()).expect("send signal");
                    None
                }
                some => some,
            }
        }
    }

    #[derive(Default)]
    struct PanicCounter;

    impl TaskIterator for PanicCounter {
        type Pending = ();

        type Ready = usize;

        type Spawner = NoSpawner;

        fn next_status(
            &mut self,
        ) -> Option<crate::valtron::TaskStatus<Self::Ready, Self::Pending, Self::Spawner>> {
            tracing::debug!("PanicCounter Task is running");
            panic!("Bad stuff");
        }
    }

    mod blocked_on_execution {
        use super::*;

        #[test]
        #[serial]
        #[traced_test]
        fn can_finish_even_when_task_panics() {
            let seed = rand::rng().next_u64();

            let handler_kill = thread::spawn(move || {
                tracing::debug!("Waiting for kill signal");
                thread::sleep(time::Duration::from_secs(5));
                tracing::debug!("Got kill signal");
                get_pool().kill();
                tracing::debug!("Closing thread");
            });

            let (task_sent_sender, task_sent_receiver) = mpp::bounded(1);
            let _guard = block_on(seed, None, |pool| {
                pool.spawn()
                    .with_task(PanicCounter)
                    .with_resolver(Box::new(FnReady::new(|item, _| {
                        tracing::info!("Received next: {item:?}");
                    })))
                    .schedule()
                    .expect("should deliver task");
                task_sent_sender.send(()).expect("deliver message");
            });

            task_sent_receiver
                .recv_timeout(time::Duration::from_secs(2))
                .expect("should have spawned task");

            tracing::info!("Wait for thread to die");
            handler_kill.join().expect("should finish");
            tracing::info!("Wait for thread to die");
        }

        #[test]
        #[serial]
        #[traced_test]
        fn can_queue_and_complete_task() {
            let seed = rand::rng().next_u64();

            let shared_list = Arc::new(Mutex::new(vec![]));
            let (counter, receiver) = Counter::new(5, shared_list.clone());

            let handler_kill = thread::spawn(move || {
                tracing::debug!("Waiting for kill signal");
                receiver
                    .recv_timeout(time::Duration::from_secs(2))
                    .expect("receive signal");
                tracing::debug!("Got kill signal");
                get_pool().kill();
                tracing::debug!("Closing thread");
            });

            let (task_sent_sender, task_sent_receiver) = mpp::bounded(1);
            let _guard = block_on(seed, None, |pool| {
                tracing::debug!("Spawning new task into pool");
                pool.spawn()
                    .with_task(counter)
                    .with_resolver(Box::new(FnReady::new(|item, _| {
                        tracing::info!("Received next: {item:?}");
                    })))
                    .schedule()
                    .expect("should deliver task");
                tracing::debug!("Task spawned");
                task_sent_sender.send(()).expect("deliver message");
            });

            task_sent_receiver
                .recv_timeout(time::Duration::from_secs(2))
                .expect("should have spawned task");

            handler_kill.join().expect("should finish");

            assert_eq!(shared_list.lock().unwrap().clone(), vec![0, 1, 2, 3, 4]);
        }
    }

    mod inline_execution {
        use crate::valtron::NotificationItem;
        use crate::valtron::Stream;

        use super::*;

        #[test]
        #[serial]
        #[traced_test]
        fn can_queue_and_complete_task_with_iterator() {
            let seed = rand::rng().next_u64();

            let shared_list = Arc::new(Mutex::new(Vec::new()));
            let counter = DCounter::new(5, shared_list.clone());

            let _guard = initialize_pool(seed, None);

            let iter = get_pool()
                .spawn()
                .with_task(counter)
                .schedule_iter(std::time::Duration::from_nanos(50))
                .expect("should deliver task");

            let complete: Vec<usize> = iter
                .map(|item| match item {
                    NotificationItem::Ready(inner) => match inner {
                        TaskStatus::Ready(value) => Some(value),
                        _ => None,
                    },
                    NotificationItem::None => None,
                })
                .take_while(|t: &Option<usize>| t.is_some())
                .map(|t: Option<usize>| t.unwrap())
                .collect();

            assert_eq!(complete, vec![1, 2, 3, 4, 5]);
        }

        #[test]
        #[serial]
        #[traced_test]
        fn can_queue_use_stream_iterator_from_task_iterator() {
            let seed = rand::rng().next_u64();

            let shared_list = Arc::new(Mutex::new(Vec::new()));
            let counter = DCounter::new(5, shared_list.clone());

            let _guard = initialize_pool(seed, None);

            let iter = get_pool()
                .spawn()
                .with_task(counter)
                .stream_iter_with_config(
                    std::time::Duration::from_nanos(20), // wait_cycle (used as park_duration)
                    10,                                  // max_turns
                )
                .expect("should deliver task");

            let complete: Vec<usize> = iter
                .map(|item| match item {
                    Stream::Next(value) => Some(value),
                    _ => None,
                })
                .take_while(|t| t.is_some())
                .map(|t| t.unwrap())
                .collect();

            assert_eq!(complete, vec![1, 2, 3, 4, 5]);
        }
    }
}