use color_eyre::Result;
use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
use tracing::debug;

/// Bidirectional relay between two async streams.
/// Returns once *both* directions have finished.
///
/// The two halves are deliberately independent: `copy_bidirectional` aborts the
/// whole relay as soon as one direction errors, which loses an HTTP response
/// that is already in flight whenever the backend answers early and closes
/// without draining the request body (e.g. a 401 on a large upload). Each
/// direction here runs to completion and shuts its own writer down, so the
/// response half still drains — and the QUIC send stream is finished cleanly
/// instead of being reset by drop.
pub async fn relay<A, B>(a: A, b: B) -> Result<()>
where
    A: AsyncRead + AsyncWrite + Unpin,
    B: AsyncRead + AsyncWrite + Unpin,
{
    let (mut ar, mut aw) = tokio::io::split(a);
    let (mut br, mut bw) = tokio::io::split(b);

    let a_to_b = async {
        let r = tokio::io::copy(&mut ar, &mut bw).await;
        let _ = bw.shutdown().await;
        r
    };
    let b_to_a = async {
        let r = tokio::io::copy(&mut br, &mut aw).await;
        let _ = aw.shutdown().await;
        r
    };

    let (ab, ba) = tokio::join!(a_to_b, b_to_a);
    if let Err(e) = ab {
        debug!("relay a->b ended: {e}");
    }
    if let Err(e) = ba {
        debug!("relay b->a ended: {e}");
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use tokio::io::{AsyncReadExt, AsyncWriteExt};

    /// The upload-502 regression: the backend answers and hangs up without
    /// draining the request body, so the relay's request half fails on its
    /// first write. The response half must still deliver.
    #[tokio::test]
    async fn response_survives_backend_early_close() {
        const RESPONSE: &[u8] = b"HTTP/1.1 401 Unauthorized\r\n\r\n";

        let (a, a_peer) = tokio::io::duplex(1024);
        let (b, mut b_peer) = tokio::io::duplex(1024);

        // Backend: reply, then hang up with the request body still unread.
        b_peer.write_all(RESPONSE).await.unwrap();
        drop(b_peer);

        // Downstream: request bytes already buffered, so the request half has
        // something to copy on its very first poll.
        let (mut a_rd, mut a_wr) = tokio::io::split(a_peer);
        a_wr.write_all(&[0u8; 512]).await.unwrap();

        relay(a, b).await.unwrap();

        let mut got = Vec::new();
        a_rd.read_to_end(&mut got).await.unwrap();
        assert_eq!(got, RESPONSE);
    }
}

/// Wrapper to combine a QUIC send+recv into a single AsyncRead+AsyncWrite.
pub struct QuicBiStream {
    pub send: quinn::SendStream,
    pub recv: quinn::RecvStream,
}

impl tokio::io::AsyncRead for QuicBiStream {
    fn poll_read(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &mut tokio::io::ReadBuf<'_>,
    ) -> std::task::Poll<std::io::Result<()>> {
        std::pin::Pin::new(&mut self.recv).poll_read(cx, buf)
    }
}

impl tokio::io::AsyncWrite for QuicBiStream {
    fn poll_write(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &[u8],
    ) -> std::task::Poll<std::io::Result<usize>> {
        std::pin::Pin::new(&mut self.send)
            .poll_write(cx, buf)
            .map(|r| r.map_err(std::io::Error::other))
    }

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

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