dcl-rpc 2.3.6

use std::{sync::Arc, time::Duration};

use async_trait::async_trait;
use log::debug;
use tokio::time::interval;

use super::{Transport, TransportError, TransportMessage};

#[cfg(feature = "tungstenite")]
pub mod tungstenite;

#[cfg(feature = "warp")]
pub mod warp;

#[derive(Debug)]
pub enum Error {
    ConnectionClosed,
    AlreadyClosed,
    Other(Box<dyn std::error::Error + Sync + Send>),
}

#[derive(Debug)]
pub enum Message {
    Text(String),
    Binary(Vec<u8>),
    Ping,
    Pong,
    Close,
}

#[async_trait]
pub trait WebSocket: Send + Sync + 'static {
    async fn send(&self, message: Message) -> Result<(), Error>;

    async fn receive(&self) -> Option<Result<Message, Error>>;

    async fn close(&self) -> Result<(), Error>;

    async fn ping_every(self: Arc<Self>, ping_interval: Duration) {
        tokio::spawn(async move {
            let mut ping_interval = interval(ping_interval);
            loop {
                ping_interval.tick().await;
                _ = self.send(Message::Ping).await;
            }
        });
    }
}

pub struct WebSocketTransport<WebSocket, Context> {
    websocket: Arc<WebSocket>,
    pub context: Context,
}

impl<WebSocket> WebSocketTransport<WebSocket, ()> {
    /// Crates a new [`WebSocketTransport`] from a websocket connection generated by [`WebSocketServer`] or [`WebSocketClient`]
    pub fn new(websocket: Arc<WebSocket>) -> Self {
        Self {
            websocket,
            context: (),
        }
    }
}

impl<WebSocket, Context> WebSocketTransport<WebSocket, Context> {
    pub fn with_context(websocket: Arc<WebSocket>, context: Context) -> Self {
        Self { websocket, context }
    }
}

#[async_trait]
impl<W: WebSocket, C: Send + Sync + 'static> Transport for WebSocketTransport<W, C> {
    async fn receive(&self) -> Result<TransportMessage, TransportError> {
        loop {
            match self.websocket.receive().await {
                Some(Ok(message)) => match message {
                    Message::Binary(data) => return Ok(data),
                    Message::Ping | Message::Pong => continue,
                    Message::Close => return Err(TransportError::Closed),
                    _ => return Err(TransportError::NotBinaryMessage),
                },
                Some(Err(err)) => {
                    debug!("> WebSocketTransport > Failed to receive message {:?}", err);
                    match err {
                        Error::ConnectionClosed | Error::AlreadyClosed => {
                            return Err(TransportError::Closed)
                        }
                        Error::Other(error) => return Err(TransportError::Internal(error)),
                    }
                }
                None => {
                    debug!("> WebSocketTransport > None received > Closing...");
                    return Err(TransportError::Closed);
                }
            }
        }
    }

    async fn send(&self, message: Vec<u8>) -> Result<(), TransportError> {
        let message = Message::Binary(message);
        match self.websocket.send(message).await {
            Err(err) => {
                debug!(
                    "> WebSocketTransport > Error on sending in a ws connection {:?}",
                    err
                );

                let error = match err {
                    Error::ConnectionClosed | Error::AlreadyClosed => TransportError::Closed,
                    Error::Other(error) => TransportError::Internal(error),
                };

                Err(error)
            }
            Ok(_) => Ok(()),
        }
    }

    async fn close(&self) {
        match self.websocket.close().await {
            Ok(_) => {
                debug!("> WebSocketTransport > Closed successfully")
            }
            Err(err) => {
                debug!("> WebSocketTransport > Error: Couldn't close tranport: {err:?}")
            }
        }
    }
}

pub(crate) fn convert<M, E>(value: Result<M, E>) -> Result<Message, Error>
where
    M: Into<Message>,
    E: Into<Error>,
{
    value.map(|m| m.into()).map_err(|e| e.into())
}