curfew 0.0.5

use crate::deadline_joiner::DeadlineSummary;
use crate::outcome::{
    PanicInfo, TaskId, TaskOutcome, TaskResult, TimedOutTask,
};
use futures::FutureExt;
use std::collections::HashMap;
use std::future::Future;
use std::panic::AssertUnwindSafe;
use std::time::Instant;
use tokio::sync::mpsc::{self, UnboundedReceiver, UnboundedSender};
use tokio::time;

const SINGLE_TASK_ID: TaskId = TaskId::MAX;

/// Attempts to await a future before an absolute deadline, returning late results via [`LateTaskAsync`].
pub async fn join_with_deadline_async<F, T>(
    future: F,
    deadline: Instant,
) -> JoinOutcomeAsync<T>
where
    F: Future<Output = T> + Send + 'static,
    T: Send + 'static,
{
    let (sender, mut receiver) = mpsc::unbounded_channel();

    tokio::spawn(async move {
        let started_at = Instant::now();
        let finished_outcome = AssertUnwindSafe(future).catch_unwind().await;
        let finished_at = Instant::now();
        let outcome = match finished_outcome {
            Ok(value) => TaskOutcome::Completed(value),
            Err(panic) => {
                TaskOutcome::Panicked(PanicInfo::from_payload(panic))
            }
        };

        let _ = sender.send(TaskResult {
            task_id: SINGLE_TASK_ID,
            started_at: Some(started_at),
            finished_at,
            outcome,
        });
    });

    let now = Instant::now();
    if now >= deadline {
        return JoinOutcomeAsync::TimedOut(LateTaskAsync { receiver });
    }

    let remaining = deadline.saturating_duration_since(now);
    match time::timeout(remaining, receiver.recv()).await {
        Ok(Some(result)) => JoinOutcomeAsync::Completed(result),
        Ok(None) => JoinOutcomeAsync::Completed(TaskResult {
            task_id: SINGLE_TASK_ID,
            started_at: None,
            finished_at: Instant::now(),
            outcome: TaskOutcome::Panicked(PanicInfo::from_payload(Box::new(
                "async task dropped before sending result",
            ))),
        }),
        Err(_) => JoinOutcomeAsync::TimedOut(LateTaskAsync { receiver }),
    }
}

/// Coordinates async tasks and collects the results that complete before the deadline.
pub struct DeadlineJoinerAsync<T: Send + 'static> {
    deadline: Instant,
    sender: UnboundedSender<TaskResult<T>>,
    receiver: UnboundedReceiver<TaskResult<T>>,
    pending: HashMap<TaskId, Instant>,
    next_id: TaskId,
}

impl<T: Send + 'static> DeadlineJoinerAsync<T> {
    /// Creates a new async joiner that stops collecting results at `deadline`.
    pub fn new(deadline: Instant) -> Self {
        let (sender, receiver) = mpsc::unbounded_channel();
        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 awaiting completion.
    pub fn pending_count(&self) -> usize {
        self.pending.len()
    }

    /// Spawns an asynchronous task and registers it with the joiner.
    pub fn spawn<F>(&mut self, future: F) -> TaskId
    where
        F: Future<Output = 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();
        tokio::spawn(async move {
            let finished_outcome =
                AssertUnwindSafe(future).catch_unwind().await;
            let finished_at = Instant::now();
            let outcome = match finished_outcome {
                Ok(value) => TaskOutcome::Completed(value),
                Err(panic) => {
                    TaskOutcome::Panicked(PanicInfo::from_payload(panic))
                }
            };

            let _ = sender.send(TaskResult {
                task_id,
                started_at: Some(started_at),
                finished_at,
                outcome,
            });
        });

        task_id
    }

    /// Waits for registered tasks until either all have finished or the deadline expires.
    pub async 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 {
                break;
            }

            let remaining = deadline.saturating_duration_since(now);
            match time::timeout(remaining, self.receiver.recv()).await {
                Ok(Some(result)) => {
                    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,
                        });
                    }
                }
                Ok(None) => {
                    break;
                }
                Err(_) => {
                    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 }
    }
}

/// Result wrapper returned by [`join_with_deadline_async`].
#[derive(Debug)]
pub enum JoinOutcomeAsync<T> {
    /// The task finished before the deadline.
    Completed(TaskResult<T>),
    /// The task was still running when the deadline expired.
    TimedOut(LateTaskAsync<T>),
}

/// Handle that allows recovering the result of a task that exceeded the deadline.
#[derive(Debug)]
pub struct LateTaskAsync<T> {
    receiver: UnboundedReceiver<TaskResult<T>>,
}

impl<T> LateTaskAsync<T> {
    /// Waits until the underlying task finishes and returns its result.
    pub async fn wait(mut self) -> TaskResult<T> {
        self.receiver
            .recv()
            .await
            .expect("async task dropped before reporting result")
    }

    /// Attempts to fetch a finished result without blocking.
    pub fn try_poll(&mut self) -> Option<TaskResult<T>> {
        self.receiver.try_recv().ok()
    }
}