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

use anyhow::bail;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf};
use tokio::time::Sleep;

/// Read a single newline-terminated line, byte by byte.
///
/// Byte-by-byte reading is intentional: after the handshake line,
/// the stream carries raw SSH/TCP bytes. A buffered reader would
/// consume bytes past the newline into its internal buffer, losing
/// them when we switch to the relay phase.
pub async fn read_line<R: AsyncRead + Unpin>(reader: &mut R) -> anyhow::Result<String> {
    let mut buf = Vec::new();
    let mut byte = [0u8; 1];
    loop {
        let n = reader.read(&mut byte).await?;
        if n == 0 {
            bail!("unexpected EOF");
        }
        if byte[0] == b'\n' {
            return Ok(String::from_utf8(buf)?);
        }
        buf.push(byte[0]);
        if buf.len() > 1024 {
            bail!("line too long");
        }
    }
}

/// Write a newline-terminated line in a single write to avoid extra TCP segments.
pub async fn write_line<W: AsyncWrite + Unpin>(writer: &mut W, line: &str) -> anyhow::Result<()> {
    let mut buf = Vec::with_capacity(line.len() + 1);
    buf.extend_from_slice(line.as_bytes());
    buf.push(b'\n');
    writer.write_all(&buf).await?;
    writer.flush().await?;
    Ok(())
}

/// Increase TCP socket buffers for higher throughput. Best-effort — silently
/// falls back to OS defaults if the requested size exceeds system limits.
pub fn tune_socket(stream: &tokio::net::TcpStream) {
    use socket2::SockRef;
    let sock = SockRef::from(stream);
    let _ = sock.set_recv_buffer_size(524_288); // 512 KB
    let _ = sock.set_send_buffer_size(524_288);
}

// ─── Bidirectional relay ───

const BUF_SIZE: usize = 65536;

/// Internal buffer for one direction of a bidirectional copy.
struct CopyBuf {
    buf: Box<[u8]>,
    pos: usize,
    cap: usize,
    total: u64,
    read_done: bool,
}

impl CopyBuf {
    fn new() -> Self {
        Self {
            buf: vec![0u8; BUF_SIZE].into_boxed_slice(),
            pos: 0,
            cap: 0,
            total: 0,
            read_done: false,
        }
    }

    /// Drive one direction of the copy. Returns Ready(bytes_copied) when done.
    fn poll_copy<R: AsyncRead + ?Sized, W: AsyncWrite + ?Sized>(
        &mut self,
        cx: &mut Context<'_>,
        mut reader: Pin<&mut R>,
        mut writer: Pin<&mut W>,
    ) -> Poll<u64> {
        loop {
            // Drain buffered data to writer — only count bytes actually written
            if self.pos < self.cap {
                match writer.as_mut().poll_write(cx, &self.buf[self.pos..self.cap]) {
                    Poll::Ready(Ok(0)) | Poll::Ready(Err(_)) => return Poll::Ready(self.total),
                    Poll::Ready(Ok(n)) => {
                        self.pos += n;
                        self.total += n as u64;
                        continue;
                    }
                    Poll::Pending => return Poll::Pending,
                }
            }

            // Reader finished — shut down writer and return
            if self.read_done {
                return match writer.as_mut().poll_shutdown(cx) {
                    Poll::Ready(_) => Poll::Ready(self.total),
                    Poll::Pending => Poll::Pending,
                };
            }

            // Fill buffer from reader — coalesce multiple reads (e.g. TLS 16KB
            // records) into a single larger write to reduce syscall count.
            let mut read_buf = ReadBuf::new(&mut self.buf);
            match reader.as_mut().poll_read(cx, &mut read_buf) {
                Poll::Ready(Ok(())) => {
                    let n = read_buf.filled().len();
                    if n == 0 {
                        self.read_done = true;
                    } else {
                        self.pos = 0;
                        self.cap = n;
                        // Try to fill more of the buffer while data is ready.
                        // TLS returns one 16KB record per read — batching 4
                        // into a single 64KB write cuts write syscalls by 4x.
                        while self.cap < self.buf.len() {
                            let mut more = ReadBuf::new(&mut self.buf[self.cap..]);
                            match reader.as_mut().poll_read(cx, &mut more) {
                                Poll::Ready(Ok(())) if more.filled().len() > 0 => {
                                    self.cap += more.filled().len();
                                }
                                _ => break,
                            }
                        }
                    }
                }
                Poll::Ready(Err(_)) => return Poll::Ready(self.total),
                Poll::Pending => return Poll::Pending,
            }
        }
    }
}

/// How long the surviving direction may drain after the first direction closes.
const DRAIN_TIMEOUT: Duration = Duration::from_secs(10);

/// Bidirectional byte relay between two streams. Returns (a→b, b→a) byte counts.
/// Always returns byte counts, even if the relay ends with an error (e.g. connection reset).
///
/// Both directions are driven from a single poll-based future, avoiding the overhead
/// and head-of-line blocking of `select!`-based relays. Each direction gets polled
/// every wakeup, so neither can starve the other.
///
/// When one direction finishes (EOF or error), its writer is shut down (FIN) and the
/// other direction gets up to [`DRAIN_TIMEOUT`] to deliver remaining in-flight data
/// before the relay is torn down.
pub async fn relay<A, B>(a: &mut A, b: &mut B) -> (u64, u64)
where
    A: AsyncRead + AsyncWrite + Unpin,
    B: AsyncRead + AsyncWrite + Unpin,
{
    let mut a_to_b = CopyBuf::new();
    let mut b_to_a = CopyBuf::new();
    let mut a_done = false;
    let mut b_done = false;
    let mut drain_deadline: Option<Pin<Box<Sleep>>> = None;

    poll_fn(|cx| {
        if !a_done {
            if let Poll::Ready(_) = a_to_b.poll_copy(
                cx,
                Pin::new(&mut *a),
                Pin::new(&mut *b),
            ) {
                a_done = true;
            }
        }
        if !b_done {
            if let Poll::Ready(_) = b_to_a.poll_copy(
                cx,
                Pin::new(&mut *b),
                Pin::new(&mut *a),
            ) {
                b_done = true;
            }
        }

        if a_done && b_done {
            return Poll::Ready((a_to_b.total, b_to_a.total));
        }

        // One direction finished — give the other a bounded window to drain
        if a_done || b_done {
            let deadline = drain_deadline
                .get_or_insert_with(|| Box::pin(tokio::time::sleep(DRAIN_TIMEOUT)));
            if deadline.as_mut().poll(cx).is_ready() {
                return Poll::Ready((a_to_b.total, b_to_a.total));
            }
        }

        Poll::Pending
    })
    .await
}
