#![cfg_attr(coverage_nightly, feature(coverage_attribute))]

use crate::SpeedLimitSession;
use pin_project::pin_project;
use std::future::Future;
use std::io::Error;
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};

enum ReadState {
    Idle,
    Waiting(Option<(Pin<Box<dyn Future<Output = usize> + 'static>>, usize)>),
    Reading(Option<(usize, usize)>),
}

enum WriteState {
    Idle,
    Waiting(Option<(Pin<Box<dyn Future<Output = usize> + 'static>>, usize)>),
    Writing(Option<(usize, usize)>),
}

#[pin_project]
pub struct LocalLimitStream<S: AsyncRead + AsyncWrite + Unpin> {
    #[pin]
    read: LocalLimitRead<sfo_split::ReadHalf<S>>,
    #[pin]
    write: LocalLimitWrite<sfo_split::WriteHalf<S>>,
}

impl<S: AsyncRead + AsyncWrite + Unpin> LocalLimitStream<S> {
    pub fn new(stream: S, read_limit: SpeedLimitSession, write_limit: SpeedLimitSession) -> Self {
        let (read, write) = sfo_split::split(stream);
        let limit_read = LocalLimitRead::new(read, read_limit);
        let limit_write = LocalLimitWrite::new(write, write_limit);
        LocalLimitStream {
            read: limit_read,
            write: limit_write,
        }
    }

    pub fn with_lock_raw_stream<R>(&mut self, f: impl FnOnce(Pin<&mut S>) -> R) -> R {
        self.read.raw_read().with_lock(f)
    }
}

impl<S: AsyncRead + AsyncWrite + Unpin> AsyncWrite for LocalLimitStream<S> {
    fn poll_write(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<Result<usize, Error>> {
        self.project().write.poll_write(cx, buf)
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
        self.project().write.poll_flush(cx)
    }

    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
        self.project().write.poll_shutdown(cx)
    }
}

impl<S: AsyncRead + AsyncWrite + Unpin> AsyncRead for LocalLimitStream<S> {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<std::io::Result<()>> {
        self.project().read.poll_read(cx, buf)
    }
}

#[pin_project]
pub struct LocalLimitRead<S: AsyncRead + Unpin> {
    #[pin]
    read: S,
    read_limit: SpeedLimitSession,
    read_state: ReadState,
}

impl<S: AsyncRead + Unpin> LocalLimitRead<S> {
    pub fn new(read: S, read_limit: SpeedLimitSession) -> Self {
        LocalLimitRead {
            read,
            read_limit,
            read_state: ReadState::Idle,
        }
    }

    pub fn raw_read_mut(&mut self) -> &mut S {
        &mut self.read
    }

    pub fn raw_read(&self) -> &S {
        &self.read
    }

    pub fn into_raw_read(self) -> S {
        self.read
    }
}

impl<S: AsyncRead + Unpin> AsyncRead for LocalLimitRead<S> {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<std::io::Result<()>> {
        let this = self.project();
        buf.initialize_unfilled();
        match this.read_state {
            ReadState::Idle => {
                let mut readded_len = 0;
                let read_limit: &'static mut SpeedLimitSession =
                    unsafe { std::mem::transmute(this.read_limit) };
                let mut waiting_future = Box::pin(read_limit.until_ready());
                match Pin::new(&mut waiting_future).poll(cx) {
                    Poll::Ready(read_len) => {
                        let mut read_buf = if read_len <= buf.remaining() {
                            buf.take(read_len)
                        } else {
                            buf.take(buf.remaining())
                        };
                        match this.read.poll_read(cx, &mut read_buf) {
                            Poll::Ready(Ok(())) => {
                                let len = read_buf.filled().len();
                                readded_len += len;
                                buf.advance(len);
                                if readded_len >= read_len {
                                    *this.read_state = ReadState::Idle;
                                } else {
                                    *this.read_state =
                                        ReadState::Reading(Some((read_len, readded_len)));
                                }
                                Poll::Ready(Ok(()))
                            }
                            Poll::Ready(Err(e)) => {
                                *this.read_state = ReadState::Idle;
                                Poll::Ready(Err(e))
                            }
                            Poll::Pending => {
                                *this.read_state =
                                    ReadState::Reading(Some((read_len, readded_len)));
                                Poll::Pending
                            }
                        }
                    }
                    Poll::Pending => {
                        *this.read_state = ReadState::Waiting(Some((waiting_future, readded_len)));
                        Poll::Pending
                    }
                }
            }
            ReadState::Waiting(state) => {
                let (mut rx, mut readded_len) = state.take().unwrap();
                match Pin::new(&mut rx).poll(cx) {
                    Poll::Ready(read_len) => {
                        let mut read_buf = if (read_len - readded_len) <= buf.remaining() {
                            buf.take(read_len - readded_len)
                        } else {
                            buf.take(buf.remaining())
                        };
                        match this.read.poll_read(cx, &mut read_buf) {
                            Poll::Ready(Ok(())) => {
                                let len = read_buf.filled().len();
                                readded_len += len;
                                buf.advance(len);
                                if readded_len >= read_len {
                                    *this.read_state = ReadState::Idle;
                                } else {
                                    *this.read_state =
                                        ReadState::Reading(Some((read_len, readded_len)));
                                }
                                Poll::Ready(Ok(()))
                            }
                            Poll::Ready(Err(e)) => {
                                *this.read_state = ReadState::Idle;
                                Poll::Ready(Err(e))
                            }
                            Poll::Pending => {
                                *this.read_state =
                                    ReadState::Reading(Some((read_len, readded_len)));
                                Poll::Pending
                            }
                        }
                    }
                    Poll::Pending => {
                        *this.read_state = ReadState::Waiting(Some((rx, readded_len)));
                        Poll::Pending
                    }
                }
            }
            ReadState::Reading(state) => match state.take() {
                Some((read_len, mut readded_len)) => {
                    let mut read_buf = if (read_len - readded_len) <= buf.remaining() {
                        buf.take(read_len - readded_len)
                    } else {
                        buf.take(buf.remaining())
                    };
                    match this.read.poll_read(cx, &mut read_buf) {
                        Poll::Ready(Ok(())) => {
                            let len = read_buf.filled().len();
                            readded_len += len;
                            buf.advance(len);
                            if readded_len >= read_len {
                                *this.read_state = ReadState::Idle;
                            } else {
                                *this.read_state =
                                    ReadState::Reading(Some((read_len, readded_len)));
                            }
                            Poll::Ready(Ok(()))
                        }
                        Poll::Ready(Err(e)) => {
                            *this.read_state = ReadState::Idle;
                            Poll::Ready(Err(e))
                        }
                        Poll::Pending => {
                            *this.read_state = ReadState::Reading(Some((read_len, readded_len)));
                            Poll::Pending
                        }
                    }
                }
                None => match this.read.poll_read(cx, buf) {
                    Poll::Ready(Ok(())) => {
                        *this.read_state = ReadState::Idle;
                        Poll::Ready(Ok(()))
                    }
                    Poll::Ready(Err(e)) => {
                        *this.read_state = ReadState::Idle;
                        Poll::Ready(Err(e))
                    }
                    Poll::Pending => {
                        *this.read_state = ReadState::Reading(None);
                        Poll::Pending
                    }
                },
            },
        }
    }
}

#[pin_project]
pub struct LocalLimitWrite<S: AsyncWrite + Unpin> {
    #[pin]
    write: S,
    write_limit: SpeedLimitSession,
    write_state: WriteState,
}

impl<S: AsyncWrite + Unpin> LocalLimitWrite<S> {
    pub fn new(write: S, write_limit: SpeedLimitSession) -> Self {
        LocalLimitWrite {
            write,
            write_limit,
            write_state: WriteState::Idle,
        }
    }

    pub fn raw_write_mut(&mut self) -> &mut S {
        &mut self.write
    }

    pub fn raw_write(&self) -> &S {
        &self.write
    }

    pub fn into_raw_write(self) -> S {
        self.write
    }
}

impl<S: AsyncWrite + Unpin> AsyncWrite for LocalLimitWrite<S> {
    fn poll_write(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<Result<usize, Error>> {
        let this = self.project();
        match this.write_state {
            WriteState::Idle => {
                let mut written_len = 0;
                let write_limiter: &'static mut SpeedLimitSession =
                    unsafe { std::mem::transmute(this.write_limit) };
                let mut waiting_future = Box::pin(write_limiter.until_ready());
                match Pin::new(&mut waiting_future).poll(cx) {
                    Poll::Ready(write_len) => {
                        let write_buf = if write_len <= buf.len() {
                            &buf[..write_len]
                        } else {
                            buf
                        };
                        match this.write.poll_write(cx, write_buf) {
                            Poll::Ready(Ok(len)) => {
                                written_len += len;
                                if written_len >= write_len {
                                    *this.write_state = WriteState::Idle;
                                } else {
                                    *this.write_state =
                                        WriteState::Writing(Some((write_len, written_len)));
                                }
                                Poll::Ready(Ok(written_len))
                            }
                            Poll::Ready(Err(e)) => {
                                *this.write_state = WriteState::Idle;
                                Poll::Ready(Err(e))
                            }
                            Poll::Pending => {
                                *this.write_state =
                                    WriteState::Writing(Some((write_len, written_len)));
                                Poll::Pending
                            }
                        }
                    }
                    Poll::Pending => {
                        *this.write_state =
                            WriteState::Waiting(Some((waiting_future, written_len)));
                        Poll::Pending
                    }
                }
            }
            WriteState::Waiting(state) => {
                let (mut waiting_future, mut written_len) = state.take().unwrap();
                match Pin::new(&mut waiting_future).poll(cx) {
                    Poll::Ready(write_len) => {
                        let write_buf = if write_len - written_len <= buf.len() {
                            &buf[..(write_len - written_len)]
                        } else {
                            buf
                        };
                        match this.write.poll_write(cx, write_buf) {
                            Poll::Ready(Ok(len)) => {
                                written_len += len;
                                if written_len >= write_len {
                                    *this.write_state = WriteState::Idle;
                                } else {
                                    *this.write_state =
                                        WriteState::Writing(Some((write_len, written_len)));
                                }
                                Poll::Ready(Ok(len))
                            }
                            Poll::Ready(Err(e)) => {
                                *this.write_state = WriteState::Idle;
                                Poll::Ready(Err(e))
                            }
                            Poll::Pending => {
                                *this.write_state =
                                    WriteState::Writing(Some((write_len, written_len)));
                                Poll::Pending
                            }
                        }
                    }
                    Poll::Pending => {
                        *this.write_state =
                            WriteState::Waiting(Some((waiting_future, written_len)));
                        Poll::Pending
                    }
                }
            }
            WriteState::Writing(state) => match state.take() {
                Some((write_len, mut written_len)) => {
                    let write_buf = if write_len - written_len <= buf.len() {
                        &buf[..(write_len - written_len)]
                    } else {
                        buf
                    };
                    match this.write.poll_write(cx, write_buf) {
                        Poll::Ready(Ok(len)) => {
                            written_len += len;
                            if written_len >= write_len {
                                *this.write_state = WriteState::Idle;
                            } else {
                                *this.write_state =
                                    WriteState::Writing(Some((write_len, written_len)));
                            }
                            Poll::Ready(Ok(len))
                        }
                        Poll::Ready(Err(e)) => {
                            *this.write_state = WriteState::Idle;
                            Poll::Ready(Err(e))
                        }
                        Poll::Pending => {
                            *this.write_state = WriteState::Writing(Some((write_len, written_len)));
                            Poll::Pending
                        }
                    }
                }
                None => match this.write.poll_write(cx, buf) {
                    Poll::Ready(Ok(len)) => {
                        *this.write_state = WriteState::Idle;
                        Poll::Ready(Ok(len))
                    }
                    Poll::Ready(Err(e)) => {
                        *this.write_state = WriteState::Idle;
                        Poll::Ready(Err(e))
                    }
                    Poll::Pending => {
                        *this.write_state = WriteState::Writing(None);
                        Poll::Pending
                    }
                },
            },
        }
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
        self.project().write.poll_flush(cx)
    }

    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
        self.project().write.poll_shutdown(cx)
    }
}

#[cfg_attr(coverage_nightly, coverage(off))]
#[cfg(test)]
mod tests {
    use super::*;
    use crate::SpeedLimiter;
    use std::cell::RefCell;
    use std::io;
    use std::num::NonZeroU32;
    use std::rc::Rc;
    use tokio::io::{AsyncReadExt, AsyncWriteExt};

    struct LocalMockStream {
        read_data: Rc<RefCell<Vec<u8>>>,
        written_data: Rc<RefCell<Vec<u8>>>,
    }

    impl AsyncRead for LocalMockStream {
        fn poll_read(
            self: Pin<&mut Self>,
            _cx: &mut Context<'_>,
            buf: &mut ReadBuf<'_>,
        ) -> Poll<io::Result<()>> {
            let mut read_data = self.read_data.borrow_mut();
            let len = read_data.len().min(buf.remaining());
            buf.put_slice(&read_data[..len]);
            read_data.drain(..len);
            Poll::Ready(Ok(()))
        }
    }

    impl AsyncWrite for LocalMockStream {
        fn poll_write(
            self: Pin<&mut Self>,
            _cx: &mut Context<'_>,
            buf: &[u8],
        ) -> Poll<io::Result<usize>> {
            self.written_data.borrow_mut().extend_from_slice(buf);
            Poll::Ready(Ok(buf.len()))
        }

        fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
            Poll::Ready(Ok(()))
        }

        fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
            Poll::Ready(Ok(()))
        }
    }

    #[tokio::test(flavor = "current_thread")]
    async fn local_limit_stream_accepts_non_send_inner() {
        let read_data = Rc::new(RefCell::new(vec![1, 2, 3]));
        let written_data = Rc::new(RefCell::new(Vec::new()));
        let mock = LocalMockStream {
            read_data,
            written_data: written_data.clone(),
        };
        let read_limit = SpeedLimiter::new(None, None, NonZeroU32::new(64)).new_limit_session();
        let write_limit = SpeedLimiter::new(None, None, NonZeroU32::new(64)).new_limit_session();
        let mut stream = LocalLimitStream::new(mock, read_limit, write_limit);

        let mut read_buf = [0; 3];
        stream.read_exact(&mut read_buf).await.unwrap();
        stream.write_all(&[4, 5, 6]).await.unwrap();

        assert_eq!(read_buf, [1, 2, 3]);
        assert_eq!(&*written_data.borrow(), &[4, 5, 6]);
    }
}
