// Copyright 2026 Hugo Osvaldo Barrera
//
// SPDX-License-Identifier: ISC

//! Connector with read and write timeouts on established connections.
//!
//! Wraps any hyper-compatible connector, adding read/write timeouts to its streams.
//! Returns [`std::io::ErrorKind::TimedOut`] on timeout.

use std::future::Future;
use std::io;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Duration;

use hyper::Uri;
use hyper_util::client::legacy::connect::{Connected, Connection};
use tokio::time::{Instant, Sleep, sleep};
use tower_service::Service;

/// Connector that wraps streams from an inner connector with I/O timeouts.
#[derive(Clone)]
pub struct TimeoutConnector<C> {
    inner: C,
    timeout: Duration,
}

impl<C> TimeoutConnector<C> {
    pub fn new(inner: C, timeout: Duration) -> Self {
        Self { inner, timeout }
    }
}

impl<C> Service<Uri> for TimeoutConnector<C>
where
    C: Service<Uri>,
    C::Response: hyper::rt::Read + hyper::rt::Write + Connection + Unpin + Send + 'static,
    C::Error: Send + 'static,
    C::Future: Send + 'static,
{
    type Response = TimeoutStream<C::Response>;
    type Error = C::Error;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, dst: Uri) -> Self::Future {
        let fut = self.inner.call(dst);
        let timeout = self.timeout;
        Box::pin(async move {
            let stream = fut.await?;
            Ok(TimeoutStream::new(stream, timeout))
        })
    }
}

/// Stream wrapper that returns [`io::ErrorKind::TimedOut`] if a read or write
/// makes no progress within the configured duration.
pub struct TimeoutStream<S> {
    inner: S,
    timeout: Duration,
    read_deadline: Option<Pin<Box<Sleep>>>,
    write_deadline: Option<Pin<Box<Sleep>>>,
}

impl<S> TimeoutStream<S> {
    fn new(inner: S, timeout: Duration) -> Self {
        Self {
            inner,
            timeout,
            read_deadline: None,
            write_deadline: None,
        }
    }
}

impl<S: Connection> Connection for TimeoutStream<S> {
    fn connected(&self) -> Connected {
        self.inner.connected()
    }
}

impl<S: hyper::rt::Read + Unpin> hyper::rt::Read for TimeoutStream<S> {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: hyper::rt::ReadBufCursor<'_>,
    ) -> Poll<Result<(), io::Error>> {
        let this = self.get_mut();
        let timeout = this.timeout;

        let deadline = this
            .read_deadline
            .get_or_insert_with(|| Box::pin(sleep(timeout)));

        if deadline.as_mut().poll(cx).is_ready() {
            return Poll::Ready(Err(io::Error::new(
                io::ErrorKind::TimedOut,
                "read timed out",
            )));
        }

        let result = Pin::new(&mut this.inner).poll_read(cx, buf);

        if let Poll::Ready(Ok(())) = &result {
            deadline.as_mut().reset(Instant::now() + timeout);
        }

        result
    }
}

impl<S: hyper::rt::Write + Unpin> hyper::rt::Write for TimeoutStream<S> {
    fn poll_write(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<Result<usize, io::Error>> {
        let this = self.get_mut();
        let timeout = this.timeout;

        let deadline = this
            .write_deadline
            .get_or_insert_with(|| Box::pin(sleep(timeout)));

        if deadline.as_mut().poll(cx).is_ready() {
            return Poll::Ready(Err(io::Error::new(
                io::ErrorKind::TimedOut,
                "write timed out",
            )));
        }

        let result = Pin::new(&mut this.inner).poll_write(cx, buf);

        if let Poll::Ready(Ok(n)) = &result
            && *n > 0
        {
            deadline.as_mut().reset(Instant::now() + timeout);
        }

        result
    }

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

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

#[cfg(test)]
mod tests {
    use std::future::Future;
    use std::io;
    use std::pin::Pin;
    use std::task::{Context, Poll};
    use std::time::Duration;

    use hyper::Uri;
    use hyper_util::{
        client::legacy::connect::{Connected, Connection},
        rt::tokio::TokioIo,
    };
    use tokio::{
        io::{AsyncReadExt, AsyncWriteExt, duplex},
        time::{sleep, timeout},
    };
    use tower_service::Service;

    use crate::timeout::{TimeoutConnector, TimeoutStream};

    /// Dummy stream for testing.
    ///
    /// Implements Connection for `TokioIo<DuplexStream>`.
    struct TestStream(TokioIo<tokio::io::DuplexStream>);

    impl hyper::rt::Read for TestStream {
        fn poll_read(
            mut self: Pin<&mut Self>,
            cx: &mut Context<'_>,
            buf: hyper::rt::ReadBufCursor<'_>,
        ) -> Poll<Result<(), io::Error>> {
            Pin::new(&mut self.0).poll_read(cx, buf)
        }
    }

    impl hyper::rt::Write for TestStream {
        fn poll_write(
            mut self: Pin<&mut Self>,
            cx: &mut Context<'_>,
            buf: &[u8],
        ) -> Poll<Result<usize, io::Error>> {
            Pin::new(&mut self.0).poll_write(cx, buf)
        }

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

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

    impl Connection for TestStream {
        fn connected(&self) -> Connected {
            Connected::new()
        }
    }

    // These tests could be greatly simplified if `cargo test` had a `--timeout`.
    // See: https://github.com/rust-lang/cargo/issues/9776

    #[tokio::test]
    async fn read_succeeds_immediately_when_data_available() {
        let (mut tx, rx) = duplex(64);
        tx.write_all(b"hello").await.unwrap();

        let stream = TestStream(TokioIo::new(rx));
        let stream = TimeoutStream::new(stream, Duration::from_secs(1));
        let mut stream = TokioIo::new(stream);

        let mut buf = [0u8; 5];
        let result = timeout(Duration::from_millis(50), stream.read(&mut buf)).await;

        assert!(result.is_ok(), "read should not timeout");
        assert_eq!(&buf[..5], b"hello");
    }

    #[tokio::test]
    async fn read_times_out_when_no_data() {
        let (_tx, rx) = duplex(64);

        let stream = TestStream(TokioIo::new(rx));
        let stream = TimeoutStream::new(stream, Duration::from_millis(50));
        let mut stream = TokioIo::new(stream);

        let mut buf = [0u8; 1];
        let result = timeout(Duration::from_millis(200), stream.read(&mut buf)).await;

        let inner_result = result.expect("outer timeout should not fire");
        let err = inner_result.expect_err("read should fail with timeout");
        assert_eq!(err.kind(), io::ErrorKind::TimedOut);
    }

    #[tokio::test]
    async fn write_succeeds_immediately_when_buffer_accepted() {
        let (tx, _rx) = duplex(64);

        let stream = TestStream(TokioIo::new(tx));
        let stream = TimeoutStream::new(stream, Duration::from_secs(1));
        let mut stream = TokioIo::new(stream);

        let result = timeout(Duration::from_millis(50), stream.write(b"hello")).await;

        assert!(result.is_ok(), "write should not timeout");
        assert_eq!(result.unwrap().unwrap(), 5);
    }

    #[tokio::test]
    async fn write_times_out_when_buffer_full() {
        let (tx, _rx) = duplex(4);

        let stream = TestStream(TokioIo::new(tx));
        let stream = TimeoutStream::new(stream, Duration::from_millis(50));
        let mut stream = TokioIo::new(stream);

        let result = timeout(Duration::from_millis(200), stream.write_all(b"hello")).await;

        let inner_result = result.expect("outer timeout should not fire");
        let err: io::Error = inner_result.expect_err("write should fail with timeout");
        assert_eq!(err.kind(), io::ErrorKind::TimedOut);
    }

    #[tokio::test]
    async fn read_deadline_resets_after_success() {
        let (mut tx, rx) = duplex(64);

        let stream = TestStream(TokioIo::new(rx));
        let stream = TimeoutStream::new(stream, Duration::from_millis(100));
        let mut stream = TokioIo::new(stream);

        let mut buf = [0u8; 3];

        sleep(Duration::from_millis(65)).await;
        tx.write_all(b"one").await.unwrap();
        stream.read_exact(&mut buf).await.expect("read 1");
        assert_eq!(&buf, b"one");

        sleep(Duration::from_millis(65)).await;
        tx.write_all(b"two").await.unwrap();
        stream.read_exact(&mut buf).await.expect("read 2");
        assert_eq!(&buf, b"two");
    }

    #[tokio::test]
    async fn write_deadline_resets_after_success() {
        let (tx, mut rx) = duplex(64);

        let stream = TestStream(TokioIo::new(tx));
        let stream = TimeoutStream::new(stream, Duration::from_millis(100));
        let mut stream = TokioIo::new(stream);

        let mut buf = [0u8; 3];

        sleep(Duration::from_millis(65)).await;
        stream.write_all(b"one").await.expect("write 1");
        rx.read_exact(&mut buf).await.unwrap();

        sleep(Duration::from_millis(65)).await;
        stream.write_all(b"two").await.expect("write 2");
        rx.read_exact(&mut buf).await.unwrap();
    }

    struct MockConnector(Option<TestStream>);

    impl Service<Uri> for MockConnector {
        type Response = TestStream;
        type Error = io::Error;
        type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;

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

        fn call(&mut self, _dst: Uri) -> Self::Future {
            let stream = self.0.take();
            Box::pin(async move { stream.ok_or_else(|| io::Error::other("no stream")) })
        }
    }
    #[tokio::test]
    async fn connector_applies_timeout_to_stream() {
        let (_tx, rx) = duplex(64);

        let connector = MockConnector(Some(TestStream(TokioIo::new(rx))));

        let mut connector = TimeoutConnector::new(connector, Duration::from_millis(50));
        let uri: Uri = "http://example.com".parse().unwrap();
        let stream = connector.call(uri).await.expect("call should succeed");
        let mut stream = TokioIo::new(stream);

        // Try to read without any data being sent shall time out.
        let mut buf = [0u8; 1];
        let result = timeout(Duration::from_millis(100), stream.read(&mut buf)).await;

        let inner = result.expect("outer timeout should not fire");
        let err = inner.expect_err("read should timeout");
        assert_eq!(err.kind(), io::ErrorKind::TimedOut);
    }
}
