curfew 0.0.5

use crate::outcome::{
    PanicInfo, TaskId, TaskOutcome, TaskResult, TimedOutTask,
};
use crossbeam_channel::{unbounded, Receiver, RecvTimeoutError, Sender};
use log::trace;
use std::collections::HashMap;
use std::thread;
use std::time::Instant;

/// Coordinates a set of background tasks and collects the results that finish before a deadline.
#[derive(Debug)]
pub struct DeadlineJoiner<T: Send + 'static> {
    deadline: Instant,
    sender: Sender<TaskResult<T>>,
    receiver: Receiver<TaskResult<T>>,
    pending: HashMap<TaskId, Instant>,
    next_id: TaskId,
}

impl<T: Send + 'static> DeadlineJoiner<T> {
    /// Creates a new joiner that stops collecting results at `deadline`.
    pub fn new(deadline: Instant) -> Self {
        let (sender, receiver) = unbounded();
        Self {
            deadline,
            sender,
            receiver,
            pending: HashMap::new(),
            next_id: 0,
        }
    }

    /// Returns the absolute deadline used by the joiner.
    pub fn deadline(&self) -> Instant {
        self.deadline
    }

    /// Number of tasks that have been spawned and not yet reported a result.
    pub fn pending_count(&self) -> usize {
        self.pending.len()
    }

    /// Spawns a task onto its own thread and registers it with the joiner.
    pub fn spawn<F>(&mut self, task: F) -> TaskId
    where
        F: FnOnce() -> T + Send + 'static,
    {
        let task_id = self.next_id;
        self.next_id = self.next_id.wrapping_add(1);
        let started_at = Instant::now();
        self.pending.insert(task_id, started_at);

        let sender = self.sender.clone();
        trace!("spawning task {task_id}");
        thread::spawn(move || {
            let outcome =
                std::panic::catch_unwind(std::panic::AssertUnwindSafe(task))
                    .map(TaskOutcome::Completed)
                    .unwrap_or_else(|payload| {
                        TaskOutcome::Panicked(PanicInfo::from_payload(payload))
                    });

            let result = TaskResult {
                task_id,
                started_at: Some(started_at),
                finished_at: Instant::now(),
                outcome,
            };

            if sender.send(result).is_err() {
                trace!("receiver dropped before task {task_id} could report its status");
            }
        });

        task_id
    }

    /// Waits for registered tasks to finish until either they have all completed or the deadline is reached.
    pub fn join_ready(mut self) -> DeadlineSummary<T> {
        let mut completed = Vec::new();
        let mut timed_out = Vec::new();
        let deadline = self.deadline;

        while !self.pending.is_empty() {
            let now = Instant::now();
            if now >= deadline {
                trace!(
                    "deadline reached with {} tasks still pending",
                    self.pending.len()
                );
                break;
            }

            match self.receiver.recv_deadline(deadline) {
                Ok(result) => {
                    trace!(
                        "task {} finished at {:?}",
                        result.task_id,
                        result.finished_at
                    );
                    self.pending.remove(&result.task_id);

                    if result.finished_at <= deadline {
                        completed.push(result);
                    } else {
                        timed_out.push(TimedOutTask {
                            task_id: result.task_id,
                            started_at: result.started_at,
                        });
                    }
                }
                Err(RecvTimeoutError::Timeout) => {
                    trace!(
                        "recv timed out with {} tasks still pending",
                        self.pending.len()
                    );
                    break;
                }
                Err(RecvTimeoutError::Disconnected) => {
                    trace!("all senders dropped; exiting join loop");
                    break;
                }
            }
        }

        for (task_id, started_at) in self.pending.into_iter() {
            timed_out
                .push(TimedOutTask { task_id, started_at: Some(started_at) });
        }

        DeadlineSummary { deadline, completed, timed_out }
    }
}

/// Summary produced by [`DeadlineJoiner::join_ready`].
#[derive(Debug)]
pub struct DeadlineSummary<T> {
    /// Deadline that bounded execution.
    pub deadline: Instant,
    /// Results completed before the deadline.
    pub completed: Vec<TaskResult<T>>,
    /// Metadata for tasks that did not produce a result before the deadline.
    pub timed_out: Vec<TimedOutTask>,
}

impl<T> DeadlineSummary<T> {
    /// Returns the number of tasks that completed before the deadline.
    pub fn completed_count(&self) -> usize {
        self.completed.len()
    }

    /// Returns the number of tasks that missed the deadline.
    pub fn timed_out_count(&self) -> usize {
        self.timed_out.len()
    }

    /// Returns the total number of tasks registered with the joiner.
    pub fn total_tasks(&self) -> usize {
        self.completed.len() + self.timed_out.len()
    }
}