//! 仅定位成功收发的同步成本，不作为完整资源合同的 Rust 参照。

use futures_util::{Sink, Stream};
use std::{
    future::{Future, poll_fn},
    pin::Pin,
    sync::Mutex,
    task::{Context, Poll, ready},
};
use tokio::io::{AsyncRead, AsyncWrite};
use tokio_tungstenite::{
    WebSocketStream,
    tungstenite::{Error, Message},
};
use tokio_util::sync::CancellationToken;

pub trait Transport: AsyncRead + AsyncWrite + Unpin + Send {}
impl<Value: AsyncRead + AsyncWrite + Unpin + Send> Transport for Value {}
pub type Boxed = Box<dyn Transport>;
type Wire = WebSocketStream<Boxed>;

pub struct Shared<const CANCEL: bool> {
    wire: Mutex<Wire>,
    read: tokio::sync::Mutex<()>,
    write: tokio::sync::Mutex<()>,
    closed: CancellationToken,
    closing: CancellationToken,
}

impl<const CANCEL: bool> Shared<CANCEL> {
    pub fn new(wire: Wire) -> Self {
        Self {
            wire: Mutex::new(wire),
            read: tokio::sync::Mutex::new(()),
            write: tokio::sync::Mutex::new(()),
            closed: CancellationToken::new(),
            closing: CancellationToken::new(),
        }
    }

    fn poll<Value>(
        &self,
        context: &mut Context<'_>,
        call: impl FnOnce(Pin<&mut Wire>, &mut Context<'_>) -> Poll<Value>,
    ) -> Poll<Value> {
        let mut wire = self.wire.lock().unwrap();
        call(Pin::new(&mut *wire), context)
    }

    async fn active<Value>(
        &self,
        work: impl Future<Output = Result<Value, Error>>,
    ) -> Result<Value, Error> {
        if CANCEL {
            tokio::select! {
                biased;
                () = self.closed.cancelled() => Err(Error::ConnectionClosed),
                () = self.closing.cancelled() => Err(Error::ConnectionClosed),
                result = work => result,
            }
        } else {
            work.await
        }
    }

    pub async fn send(&self, frame: Message) -> Result<(), Error> {
        self.active(async {
            let _direction = self.write.lock().await;
            let mut frame = Some(frame);
            poll_fn(|context| {
                self.poll(context, |mut wire, context| {
                    ready!(wire.as_mut().poll_ready(context))?;
                    Poll::Ready(wire.start_send(frame.take().unwrap()))
                })
            })
            .await?;
            poll_fn(|context| self.poll(context, |wire, context| wire.poll_flush(context))).await
        })
        .await
    }

    pub async fn receive(&self) -> Result<Option<Message>, Error> {
        self.active(async {
            let _direction = self.read.lock().await;
            poll_fn(|context| self.poll(context, |wire, context| wire.poll_next(context)))
                .await
                .transpose()
        })
        .await
    }

    pub async fn close(&mut self) {
        self.wire.get_mut().unwrap().close(None).await.unwrap();
    }
}