dscale 0.5.2

use crate::{
    Pid, ProcessHandle,
    events::PidHandlerEvent,
    global::local_access::{self, setup_local_access},
    process_handle::ProcessHandlePtr,
    random::Seed,
    runners::{
        task::{TaskId, TaskResult},
        threads::Threads,
    },
};
use crossbeam_channel::{Receiver, RecvError, Sender};
use std::thread::{self, JoinHandle};

pub(crate) struct WorkItem {
    pub task_id: TaskId,
    pub pidx: usize,
    pub event: PidHandlerEvent,
}

struct ThreadHandle {
    join: JoinHandle<()>,
    tx: Sender<WorkItem>,
}

pub(crate) struct Workers {
    threads: Vec<ThreadHandle>,
    rx: Receiver<TaskResult>,
    num_procs: usize,
}

impl Workers {
    pub(crate) fn new(procs: Vec<ProcessHandlePtr>, threads: Threads, seed: Seed) -> Self {
        let num_procs = procs.len();
        let threads_number: usize = threads.into();
        log::warn!("Using {threads_number} threads for simulation");

        let (result_tx, result_rx) = crossbeam_channel::bounded::<TaskResult>(4 * num_procs);

        let mut per_thread: Vec<Vec<(usize, ProcessHandlePtr)>> =
            (0..threads_number).map(|_| Vec::new()).collect();

        for (global_id, proc) in procs.into_iter().enumerate() {
            per_thread[global_id % threads_number].push((global_id, proc));
        }

        let thread_handles = per_thread
            .into_iter()
            .map(|mut owned_procs| {
                let (work_tx, work_rx) = crossbeam_channel::bounded::<WorkItem>(100);
                let result_tx = result_tx.clone();

                let join = thread::spawn(move || {
                    setup_local_access(seed, result_tx);

                    while let Ok(WorkItem {
                        task_id,
                        pidx,
                        event,
                    }) = work_rx.recv()
                    {
                        let (global_id, ref mut proc) = owned_procs[pidx];
                        let pid = global_id as Pid;

                        local_access::prepare_task(task_id, pid);

                        match event {
                            PidHandlerEvent::Start { .. } => proc.on_start(),
                            PidHandlerEvent::Network {
                                source, message, ..
                            } => proc.on_message(source, message),
                            PidHandlerEvent::Timer { id, .. } => proc.on_timer(id),
                        }

                        local_access::done();
                    }
                });

                ThreadHandle { join, tx: work_tx }
            })
            .collect();

        Self {
            threads: thread_handles,
            rx: result_rx,
            num_procs,
        }
    }

    pub(crate) fn num_procs(&self) -> usize {
        self.num_procs
    }

    pub(crate) fn spawn_event(&self, task_id: TaskId, event: PidHandlerEvent) {
        let global_pid = event.target_pid();
        let threads_number = self.threads.len();

        let thread_idx = global_pid % threads_number;
        let pidx = global_pid / threads_number;

        self.threads[thread_idx]
            .tx
            .send(WorkItem {
                task_id,
                pidx,
                event,
            })
            .expect("worker thread has exited unexpectedly");
    }

    pub(crate) fn try_next_result(&self) -> Option<TaskResult> {
        self.rx.try_recv().ok()
    }

    pub(crate) fn next_result(&self) -> Result<TaskResult, RecvError> {
        self.rx.recv()
    }
}

impl Drop for Workers {
    fn drop(&mut self) {
        let handles: Vec<JoinHandle<()>> = self
            .threads
            .drain(..)
            .map(|h| {
                drop(h.tx);
                h.join
            })
            .collect();

        for h in handles {
            h.join().expect("worker thread panicked");
        }
    }
}