use super::*;

/// A builder that lets you prepare a connection before [`Self::finish`]ing it to attempt to
/// open the connection.
#[must_use = "A ConnectionBuilder must be `finish`ed to actually establish a connection"]
pub struct ConnectionBuilder<
    'a,
    Arguments = NoArguments,
    OnReceiveFactory = (),
    OnReceiveOverload = (),
> {
    pub(crate) endpoint: &'a Endpoint,
    pub(crate) peer: PeerPubkey,
    pub(crate) topic: Option<datastreams_ffi::Topic>,
    pub(crate) reliability: datastreams_ffi::Reliability,
    pub(crate) on_receive: OnReceiveFactory,
    pub(crate) enable_compression: bool,
    pub(crate) timeout_ms: u32,
    pub(crate) marker: core::marker::PhantomData<OnReceiveOverload>,
    pub(crate) arguments: Arguments,
}

impl<'a, OnReceiveFactory, OnReceiveOverload>
    ConnectionBuilder<'a, NoArguments, OnReceiveFactory, OnReceiveOverload>
{
    /// Allows setting arguments to be sent along with the connection request.
    ///
    /// These arguments will be made available to the receiving side's [`StreamCandidate`] through
    /// [`StreamCandidate::arguments`].
    ///
    /// Arguments may only be set once.
    ///
    /// # Reliability
    ///
    /// Arguments are supported for both reliable and unreliable streams. You may expect them to
    /// piggy-back on the stream establishment handshake (letting them be delivered before the
    /// remote accepts the stream); and the remote is guaranteed to be able to access their full
    /// contents before making a decision on whether or not to accept the candidate.
    ///
    /// # Example
    ///
    /// Let's say you're building a file transfer protocol, where the client may request files from
    /// the server by their ID, and may state what compression schemes it supports:
    /// ```
    /// # mod dittolive_ditto {
    /// #   pub mod preview {
    /// #       pub use dittolive_ditto_base::*;
    /// #   }
    /// # }
    /// # use dittolive_ditto::preview::{datastreams::{Endpoint, Payload}, peer_pubkey::PeerPubkey};
    /// # use safer_ffi::bytes::Bytes;
    /// # async fn usage(endpoint: Endpoint, target: PeerPubkey) {
    /// let candidate = endpoint
    ///     .connect(target, "file_transfer")
    ///     .with_arguments(Payload::from(
    ///         br#"{"id": "my_file", "compression": ["lz4", "xz"]}"#,
    ///     ))
    ///     .finish_async()
    ///     .await
    ///     .unwrap();
    ///
    /// match candidate.arguments().as_deref() {
    ///     None => { /* No compression */ }
    ///     Some(b"lz4") => { /* Server picked LZ4 */ }
    ///     Some(b"xz") => { /* Server picked XZ */ }
    ///     Some(_) => { /* Unsupported compression */ }
    /// }
    /// # }
    /// ```
    pub fn with_arguments<Args: IArguments>(
        self,
        arguments: Args,
    ) -> ConnectionBuilder<'a, Args, OnReceiveFactory, OnReceiveOverload> {
        let Self {
            endpoint,
            peer,
            topic,
            reliability,
            enable_compression,
            timeout_ms,
            on_receive,
            marker: _,
            arguments: _,
        } = self;
        ConnectionBuilder {
            endpoint,
            peer,
            topic,
            reliability,
            enable_compression,
            timeout_ms,
            on_receive,
            marker: core::marker::PhantomData,
            arguments,
        }
    }
}

impl<'a, Args> ConnectionBuilder<'a, Args> {
    /// Allows setting a [`SingleCandidateHandler`] factory that will be used to handle the incoming
    /// [`StreamCandidate`] once it is constructed.
    ///
    /// You may still
    /// finish<sup>[(0)](ConnectionBuilder::finish)[(1)](ConnectionBuilder::finish_with)[(2)](ConnectionBuilder::finish_async)
    /// </sup> your [`ConnectionBuilder`] without setting such a factory. If you do so, your handler
    /// will receive a [`StreamCandidate`] instead of the already open [`Stream`] you'd be provided
    /// by if a factory has been set.
    pub fn on_receive_factory<OnReceiveOverload, OnReceiveFactory>(
        self,
        on_recv_factory: OnReceiveFactory,
    ) -> ConnectionBuilder<'a, Args, OnReceiveFactory, OnReceiveOverload> {
        let Self {
            endpoint,
            peer,
            topic,
            reliability,
            enable_compression,
            timeout_ms,
            on_receive: _,
            marker: _,
            arguments,
        } = self;
        ConnectionBuilder {
            endpoint,
            peer,
            topic,
            reliability,
            enable_compression,
            timeout_ms,
            arguments,
            on_receive: on_recv_factory,
            marker: core::marker::PhantomData,
        }
    }
}

impl<
        OnReceiveFactory: SingleCandidateHandler<OnReceiveOverload, ArgumentsDisallowed>,
        OnReceiveOverload,
        Args: IArguments,
    > ConnectionBuilder<'_, Args, OnReceiveFactory, OnReceiveOverload>
{
    /// Sets the reliability mode for the connection.
    ///
    /// The default is [`Reliability::Reliable`].
    pub fn reliability(mut self, mode: Reliability) -> Self {
        self.reliability = mode;
        self
    }

    /// Enables compression for the stream, letting the peers pick a compression scheme they share.
    ///
    /// By default, compression is disabled.
    pub fn enable_compression(mut self) -> Self {
        self.enable_compression = true;
        self
    }

    /// Sets a timeout in milliseconds for the connection handshake.
    ///
    /// If the handshake has not completed by then, it will be cancelled, and
    /// [`ConnectionError::HandshakeTimedOut`] will be provided to the continuation.
    pub fn timeout_ms(mut self, timeout_ms: u32) -> Self {
        self.timeout_ms = timeout_ms;
        self
    }

    /// Attempts to open the connection, calling the `handler`'s continuation once opening has
    /// either succeeded of failed.
    ///
    /// # [`IntoChannel`]
    /// The signature of this method is built to let you pass various handlers for the events
    /// that will follow:
    /// - You may pass closures directly.
    /// - You may also pass channel tuples, such as `std::sync::mpsc::channel()`. The returned value
    ///   will wrap your channel's receiver end into the handle that keeps your binding alive.
    pub fn finish<OverloadId, Handler>(
        self,
        handler: Handler,
    ) -> ConnectionHandle<Handler::Receiver>
    where
        Handler: IntoChannel<OverloadId, Result<OnReceiveFactory::Output, ConnectionError>>,
        Handler::Sender:
            OneShotChannel<OverloadId, Result<OnReceiveFactory::Output, ConnectionError>>,
    {
        let Self {
            endpoint,
            peer,
            topic,
            reliability,
            enable_compression,
            on_receive: on_recv,
            arguments,
            timeout_ms,
            marker: _,
        } = self;
        let (channel, receiver) = handler.into_channel();
        let Some(topic) = topic else {
            channel.send_infallibly_once(Err(ConnectionError::InvalidTopic));
            return ConnectionHandle::from_failure(receiver);
        };

        let mut channel = Some((channel, on_recv));
        let closure = Box::new(move |value: datastreams_ffi::ConnectionResult| {
            let Some((channel, on_recv)) = channel.take() else {
                unreachable!("Connections only resolve once.")
            };
            channel.send_infallibly_once(match value {
                datastreams_ffi::ConnectionResult::Accepted(candidate) => {
                    Ok(on_recv.open_once(StreamCandidate {
                        inner: candidate,
                        arguments: ArgumentsDisallowed,
                    }))
                }
                datastreams_ffi::ConnectionResult::Rejected(err) => Err(err),
            })
        })
        .into();
        let handle = endpoint.inner.connect(
            peer,
            (&*topic).into(),
            reliability,
            enable_compression,
            closure,
            timeout_ms,
            arguments.into_arguments(),
        );
        ConnectionHandle {
            inner: handle,
            receiver,
        }
    }

    /// Attempts to open the connection, calling the `on_connect` once opening has
    /// either succeeded of failed.
    ///
    /// # Relation  to [`ConnectionBuilder::finish`]
    /// This method compensates for a weakness in Rust's type inference: when you pass a closure to
    /// [`ConnectionBuilder::finish`], Rust will generally require that you explicitly type its
    /// arguments, despite being able to compute the only valid type for it's argument (as
    /// highlighted if you purposely mis-type the argument).
    ///
    /// Hence, it essentially exists to highlight this paper-cut and to let you _still_ not worry
    /// about typing your closures.
    pub fn finish_with<
        F: FnOnce(Result<OnReceiveFactory::Output, ConnectionError>) + Send + Sync + 'static,
    >(
        self,
        on_connect: F,
    ) -> ConnectionHandle<()> {
        self.finish(on_connect)
    }

    /// Returns a future that resolves once the connection has been established or determined to
    /// fail.
    pub fn finish_async(self) -> ConnectionFuture<OnReceiveFactory::Output>
    where
        OnReceiveFactory::Output: Send + 'static,
    {
        ConnectionFuture {
            rx: self.finish(tokio::sync::oneshot::channel()).into_future(),
        }
    }
}

/// A future that resolves once a connection is established or has failed.
///
/// # Cancel-safety
/// Cancelling this future will cancel the connection attempt if it has not yet completed.
#[must_use = "Dropping a ConnectionFuture may cancel the connection attempt."]
pub struct ConnectionFuture<T> {
    rx: <ConnectionHandle<tokio::sync::oneshot::Receiver<Result<T, ConnectionError>>> as core::future::IntoFuture>::IntoFuture
}
impl<T> core::future::Future for ConnectionFuture<T> {
    type Output = Result<T, ConnectionError>;
    fn poll(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Self::Output> {
        match core::future::Future::poll(std::pin::Pin::new(&mut self.get_mut().rx), cx) {
            std::task::Poll::Ready(r) => {
                std::task::Poll::Ready(r.unwrap_or(Err(ConnectionError::PeerNotFound)))
            }
            std::task::Poll::Pending => std::task::Poll::Pending,
        }
    }
}

/// When a connection is initiated, a handle is returned that will cancel that connection attempt if
/// dropped.
#[must_use = "Dropping a ConnectionHandle may cancel the connection attempt."]
pub struct ConnectionHandle<Receiver> {
    pub(crate) inner: datastreams_ffi::ConnectionHandle,
    pub(crate) receiver: Receiver,
}

impl<Receiver> ConnectionHandle<Receiver> {
    /// Attempts to cancel the connection request.
    ///
    /// Note that this wil not cancel the connection if it has already been established.
    pub fn cancel(self) -> Receiver {
        self.receiver
    }

    /// Detaches the handle, preventing it from cancelling the connection attempt when dropped.
    pub fn detach(self) -> Receiver {
        let Self { inner, receiver } = self;
        inner.detach();
        receiver
    }

    /// Splits the handle into its underlying FFI handle and the receiver.
    pub fn split(self) -> (datastreams_ffi::ConnectionHandle, Receiver) {
        (self.inner, self.receiver)
    }

    fn from_failure(receiver: Receiver) -> Self {
        Self {
            inner: datastreams_ffi::ConnectionHandle::noop(),
            receiver,
        }
    }
}

impl<R> core::ops::Deref for ConnectionHandle<R> {
    type Target = R;
    fn deref(&self) -> &Self::Target {
        &self.receiver
    }
}

impl<R> core::ops::DerefMut for ConnectionHandle<R> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.receiver
    }
}

impl<T> core::future::Future for ConnectionHandle<T>
where
    T: core::future::Future,
{
    type Output = T::Output;
    fn poll(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Self::Output> {
        match core::future::Future::poll(
            // SAFETY: We never move the receiver out of the pin, just pass a pinned mutable
            unsafe { self.as_mut().map_unchecked_mut(|this| &mut this.receiver) },
            cx,
        ) {
            std::task::Poll::Ready(r) => {
                core::mem::replace(
                    unsafe { &mut self.get_unchecked_mut().inner },
                    datastreams_ffi::ConnectionHandle::noop(),
                )
                .detach();
                std::task::Poll::Ready(r)
            }
            std::task::Poll::Pending => std::task::Poll::Pending,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn connection_future() {
        for value in [
            Ok(()),
            Err(ConnectionError::HandshakeFailed),
            Err(ConnectionError::HandshakeTimedOut),
            Err(ConnectionError::ConnectionRejected),
            Err(ConnectionError::PeerNotFound),
        ] {
            let (tx, rx) = tokio::sync::oneshot::channel();
            _ = tx.send(value);
            assert_eq!(
                value,
                ConnectionFuture {
                    rx: ConnectionHandle::from_failure(rx)
                }
                .await
            )
        }
        let (_, rx) = tokio::sync::oneshot::channel();
        assert_eq!(
            Err::<(), _>(ConnectionError::PeerNotFound),
            ConnectionFuture {
                rx: ConnectionHandle::from_failure(rx)
            }
            .await
        );
    }
}
