// Copyright 2023-Present Jonas Kruckenberg
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

use core::fmt;
use core::future::Future;
use core::marker::PhantomData;
use core::panic::{RefUnwindSafe, UnwindSafe};
use core::pin::Pin;
use core::task::{Context, Poll};

use crate::task::{Id, TaskRef};

pub struct JoinHandle<T, M> {
    state: JoinHandleState,
    id: Id,
    _p: PhantomData<(T, M)>,
}
static_assertions::assert_impl_all!(JoinHandle<(), ()>: Send);

#[derive(Debug)]
enum JoinHandleState {
    Task(TaskRef),
    Empty,
}

pub struct JoinError<T> {
    kind: JoinErrorKind,
    id: Id,
    output: Option<T>,
}

#[derive(Debug)]
#[non_exhaustive]
pub enum JoinErrorKind {
    Cancelled { completed: bool },
    Panic,
}

// === impl JoinHandle ===

impl<T, M> UnwindSafe for JoinHandle<T, M> {}

impl<T, M> RefUnwindSafe for JoinHandle<T, M> {}

impl<T, M> Unpin for JoinHandle<T, M> {}

impl<T, M> Drop for JoinHandle<T, M> {
    fn drop(&mut self) {
        // if the JoinHandle has not already been consumed, clear the join
        // handle flag on the task.
        if let JoinHandleState::Task(ref task) = self.state {
            tracing::trace!(
                state=?self.state,
                task.id=?task.id(),
                consumed=false,
                "drop JoinHandle"
            );

            task.state().drop_join_handle();
        } else {
            tracing::trace!(
                state=?self.state,
                consumed=false,
                "drop JoinHandle"
            );
        }
    }
}

impl<T, M> fmt::Debug for JoinHandle<T, M>
where
    T: fmt::Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("JoinHandle")
            .field("output", &core::any::type_name::<T>())
            .field("task", &self.state)
            .field("id", &self.id)
            .finish()
    }
}

impl<T, M> Future for JoinHandle<T, M> {
    type Output = Result<T, JoinError<T>>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.get_mut();
        let task = match core::mem::replace(&mut this.state, JoinHandleState::Empty) {
            JoinHandleState::Task(task) => task,
            JoinHandleState::Empty => {
                panic!("`TaskRef` only taken while polling a `JoinHandle`; this is a bug")
            }
        };

        // // Keep track of task budget
        // TODO let coop = ready!(crate::runtime::coop::poll_proceed(cx));

        // Safety: the `JoinHandle` must have been constructed with the
        // task's actual output type!
        let poll = unsafe { task.poll_join::<T>(cx) };

        if poll.is_pending() {
            this.state = JoinHandleState::Task(task);
        } else {
            // TODO coop.made_progress();

            // clear join interest
            task.state().drop_join_handle();
        }
        poll
    }
}

// ==== PartialEq impls for JoinHandle/TaskRef ====

impl<T, M> PartialEq<TaskRef> for JoinHandle<T, M> {
    fn eq(&self, other: &TaskRef) -> bool {
        match self.state {
            JoinHandleState::Task(ref task) => task == other,
            _ => false,
        }
    }
}

impl<T, M> PartialEq<&'_ TaskRef> for JoinHandle<T, M> {
    fn eq(&self, other: &&TaskRef) -> bool {
        match self.state {
            JoinHandleState::Task(ref task) => task == *other,
            _ => false,
        }
    }
}

impl<T, M> PartialEq<JoinHandle<T, M>> for TaskRef {
    fn eq(&self, other: &JoinHandle<T, M>) -> bool {
        match other.state {
            JoinHandleState::Task(ref task) => self == task,
            _ => false,
        }
    }
}

impl<T, M> PartialEq<&'_ JoinHandle<T, M>> for TaskRef {
    fn eq(&self, other: &&JoinHandle<T, M>) -> bool {
        match other.state {
            JoinHandleState::Task(ref task) => self == task,
            _ => false,
        }
    }
}

// ==== PartialEq impls for JoinHandle/Id ====

impl<T, M> PartialEq<Id> for JoinHandle<T, M> {
    #[inline]
    fn eq(&self, other: &Id) -> bool {
        self.id == *other
    }
}

impl<T, M> PartialEq<JoinHandle<T, M>> for Id {
    #[inline]
    fn eq(&self, other: &JoinHandle<T, M>) -> bool {
        *self == other.id
    }
}

impl<T, M> PartialEq<&'_ JoinHandle<T, M>> for Id {
    #[inline]
    fn eq(&self, other: &&JoinHandle<T, M>) -> bool {
        *self == other.id
    }
}

impl<T, M> JoinHandle<T, M> {
    /// Cancels the task associated with the handle.
    ///
    /// Awaiting a cancelled task might complete as usual if the task was
    /// already completed at the time it was cancelled, but most likely it
    /// will fail with a [cancelled] [`JoinError`].
    ///
    /// See [the module level docs] for details.
    ///
    /// [cancelled]: JoinError::is_cancelled
    /// [the module level docs]: super#cancellation
    pub fn cancel(&self) -> bool {
        match self.state {
            JoinHandleState::Task(ref task) => task.cancel(),
            _ => false,
        }
    }

    #[inline]
    #[must_use]
    pub fn is_complete(&self) -> bool {
        match self.state {
            JoinHandleState::Task(ref task) => task.is_complete(),
            // if the `JoinHandle`'s `TaskRef` has been taken, we know the
            // `Future` impl for `JoinHandle` completed, and the task has
            // _definitely_ completed.
            _ => true,
        }
    }

    pub fn metadata(&self) -> Option<&M> {
        if let JoinHandleState::Task(ref task) = self.state {
            // Safety: we know this generic is correct through construction
            Some(unsafe { task.metadata::<M>() })
        } else {
            None
        }
    }

    pub(crate) fn new(task: TaskRef) -> Self {
        task.state().create_join_handle();

        Self {
            id: task.id(),
            state: JoinHandleState::Task(task),
            _p: PhantomData,
        }
    }
}

// === impl JoinError ===

impl JoinError<()> {
    pub(super) fn cancelled(completed: bool, id: Id) -> Self {
        Self {
            kind: JoinErrorKind::Cancelled { completed },
            id,
            output: None,
        }
    }

    pub(super) fn with_output<T>(self, output: Option<T>) -> JoinError<T> {
        JoinError {
            kind: self.kind,
            id: self.id,
            output,
        }
    }
}

impl<T> JoinError<T> {
    pub(super) fn panic(id: Id) -> Self {
        Self {
            kind: JoinErrorKind::Panic,
            id,
            output: None,
        }
    }

    pub fn is_completed(&self) -> bool {
        matches!(&self.kind, JoinErrorKind::Cancelled { completed: true })
    }

    /// Returns true if the error was caused by the task being cancelled.
    ///
    /// See [the module level docs] for more information on cancellation.
    ///
    /// [the module level docs]: crate::task#cancellation
    pub fn is_cancelled(&self) -> bool {
        matches!(&self.kind, JoinErrorKind::Cancelled { .. })
    }

    /// Returns true if the error was caused by the task panicking.
    ///
    /// # Examples
    //     ```
    //     use std::panic;
    //
    //     #[tokio::main]
    //     async fn main() {
    //         let err = tokio::spawn(async {
    //             panic!("boom");
    //         }).await.unwrap_err();
    //
    //         assert!(err.is_panic());
    //     }
    //     ```
    pub fn is_panic(&self) -> bool {
        matches!(&self.kind, JoinErrorKind::Panic)
    }

    /// Returns a [task ID] that identifies the task which errored relative to
    /// other currently spawned tasks.
    ///
    /// [task ID]: Id
    pub fn id(&self) -> Id {
        self.id
    }

    /// Returns the task's output, if the task completed successfully before it
    /// was canceled.
    ///
    /// Otherwise, returns `None`.
    pub fn output(self) -> Option<T> {
        self.output
    }
}

impl<T> fmt::Display for JoinError<T> {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        match &self.kind {
            JoinErrorKind::Cancelled { completed: false } => {
                write!(fmt, "task {} was cancelled before completion", self.id)
            }
            JoinErrorKind::Cancelled { completed: true } => {
                write!(fmt, "task {} was cancelled after completion", self.id)
            }
            JoinErrorKind::Panic => {
                write!(fmt, "task {} panicked", self.id)
            }
        }
    }
}

impl<T> fmt::Debug for JoinError<T> {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        match &self.kind {
            JoinErrorKind::Cancelled { completed } => write!(
                fmt,
                "JoinError::Cancelled({:?}, completed: {completed})",
                self.id
            ),
            JoinErrorKind::Panic => {
                write!(fmt, "JoinError::Panic({:?}, ...)", self.id)
            }
        }
    }
}

impl<T> core::error::Error for JoinError<T> {}
