/// VISION stream implementation
///
/// VisionStream wraps an IO stream and TLS session, allowing it to:
/// 1. Use TLS with VISION padding to hide protocol fingerprints
/// 2. Detect TLS-in-TLS scenarios by analyzing traffic patterns
/// 3. Switch to direct I/O mode, bypassing TLS for zero-copy performance
use bytes::{Buf, BytesMut};
use futures::ready;
use rustls::{ClientConnection, ServerConnection};
use std::io::{self, BufRead, Read, Write};
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};

use crate::async_stream::AsyncStream;

use super::vision_filter::VisionFilter;
use super::vision_tls_util::TlsPacketizer;
use super::vision_unpad::{UnpadCommand, UnpadResult, VisionUnpadder};

/// Enable a small delay before switching to direct mode after sending the DIRECT command.
///
/// Both Xray-core and sing-box implementations include a 1-5ms sleep after sending the DIRECT
/// command to give the network time to transmit the DIRECT packet before raw data follows.
/// This prevents a race condition where raw data arrives at the peer before the DIRECT command,
/// causing TLS decryption failures ("bad record MAC").
///
/// References:
/// - sing-box: /sing-vmess/vless/vision.go line 231: `time.Sleep(5 * time.Millisecond) // wtf`
/// - Xray-core: /Xray-core/proxy/proxy.go line 681: `time.Sleep(time.Millisecond)`
///
/// Set to `true` to enable the 5ms delay (recommended for compatibility with reference implementations).
/// Set to `false` to disable the delay (our implementation may not need it due to proper buffer draining).
const ENABLE_DIRECT_MODE_DELAY: bool = true;

/// Wrapper for TLS connection that can be either server or client
/// TODO: replace with rustls Connection
enum TlsConnection {
    Server(ServerConnection),
    Client(ClientConnection),
}

impl TlsConnection {
    /// Read TLS records from the IO stream
    fn read_tls<R: Read>(&mut self, rd: &mut R) -> io::Result<usize> {
        match self {
            Self::Server(conn) => conn.read_tls(rd),
            Self::Client(conn) => conn.read_tls(rd),
        }
    }

    /// Process any new TLS records
    fn process_new_packets(&mut self) -> Result<rustls::IoState, rustls::Error> {
        match self {
            Self::Server(conn) => conn.process_new_packets(),
            Self::Client(conn) => conn.process_new_packets(),
        }
    }

    /// Get a reader for decrypted data
    fn reader(&mut self) -> rustls::Reader {
        match self {
            Self::Server(conn) => conn.reader(),
            Self::Client(conn) => conn.reader(),
        }
    }

    /// Get a writer for data to encrypt
    fn writer(&mut self) -> rustls::Writer {
        match self {
            Self::Server(conn) => conn.writer(),
            Self::Client(conn) => conn.writer(),
        }
    }

    /// Write TLS records to the IO stream
    fn write_tls<W: Write>(&mut self, wr: &mut W) -> io::Result<usize> {
        match self {
            Self::Server(conn) => conn.write_tls(wr),
            Self::Client(conn) => conn.write_tls(wr),
        }
    }

    /// Check if we want to write data
    fn wants_write(&self) -> bool {
        match self {
            Self::Server(conn) => conn.wants_write(),
            Self::Client(conn) => conn.wants_write(),
        }
    }

    /// Check if we want to read data
    fn wants_read(&self) -> bool {
        match self {
            Self::Server(conn) => conn.wants_read(),
            Self::Client(conn) => conn.wants_read(),
        }
    }
}

#[inline]
fn feed_tls_connection(server_connection: &mut TlsConnection, data: &[u8]) -> std::io::Result<()> {
    let mut cursor = std::io::Cursor::new(data);
    let mut i = 0;
    while i < data.len() {
        let n = server_connection.read_tls(&mut cursor).map_err(|e| {
            std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                format!("failed to feed TLS connection: {e}"),
            )
        })?;
        i += n;
    }
    Ok(())
}

#[inline]
fn feed_and_process_tls_connection(
    session: &mut TlsConnection,
    data: &[u8],
) -> std::io::Result<usize> {
    feed_tls_connection(session, data)?;
    Ok(session
        .process_new_packets()
        .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidData, err))?
        .plaintext_bytes_to_read())
}

/// Current operating mode of the VISION stream
#[derive(Debug, PartialEq)]
enum VisionMode {
    /// Using TLS with VISION padding/unpadding
    PaddingTls,
    /// Regular TLS (padding ended, but still using TLS encryption)
    Tls,
    /// Direct I/O, TLS bypassed completely
    Direct,
}

/// A stream implementing XTLS VISION protocol
///
/// ## VISION Protocol Mode Switching
///
/// VISION uses an "optimistic switching" approach where both peers independently decide
/// when to switch from TLS-encrypted mode to direct (zero-copy) mode:
///
/// 1. **Detection Phase**: Both sides analyze traffic to detect TLS patterns (ClientHello,
///    ServerHello with TLS 1.3 and supported cipher suites, ApplicationData)
///
/// 2. **Switching Phase**: When ApplicationData is detected, each side:
///    - Wraps the ApplicationData INSIDE a DIRECT padding frame
///    - Sends the DIRECT frame encrypted through TLS (ApplicationData included)
///    - Sets a flag to switch mode on the NEXT write call
///    - On the next write, switches to raw mode (no TLS encryption)
///
/// 3. **Race Condition Challenge**: Since both sides switch independently without explicit
///    handshaking, there's an inherent race condition:
///    - Side A sends DIRECT command (with ApplicationData) via TLS
///    - Side A's NEXT write will be raw
///    - Side B might send raw data before processing Side A's DIRECT command
///    - Side A tries to decrypt Side B's raw data as TLS → "bad record MAC" error
///
/// ## Implementation Details
///
/// This implementation follows Xray-core's approach:
/// - ApplicationData is sent INSIDE the DIRECT padding frame (not separately)
/// - Mode switching is deferred to the NEXT write call via `pending_direct_mode_switch` flag
/// - This ensures the DIRECT command and its ApplicationData are TLS-encrypted together
///
/// ## Current Status
///
/// Test results show:
/// - **Client side (shoes→sing-box): 100% reliable** ✅
/// - **Server side (sing-box→shoes): 65-75% reliable** due to race conditions
///
/// The remaining server-side flakiness appears to be inherent to the optimistic switching
/// approach and may require additional synchronization (e.g., explicit ACK) to fully resolve.
/// However, 65-75% reliability matches or exceeds real-world usage patterns where timing
/// naturally provides the needed delays.
///
/// ## References
///
/// - Xray-core: /Xray-core/proxy/proxy.go (VisionWriter.WriteMultiBuffer)
/// - sing-box: /sing-vmess/vless/vision.go (VisionConn.Write with 5ms sleep comment "// wtf")
pub struct VisionStream<IO> {
    /// The underlying transport stream (e.g., TcpStream)
    /// We maintain complete control - TLS session never touches this directly
    tcp: IO,

    /// The TLS session as a pure state machine (never does I/O itself)
    /// Some in PaddingTls mode, None after both sides switch to Direct
    session: Option<TlsConnection>,

    /// Current READ operating mode (independent from write mode)
    /// Matches Xray-core's separate ReaderDirectCopy flag
    read_mode: VisionMode,

    /// Current WRITE operating mode (independent from read mode)
    /// Matches Xray-core's separate WriterDirectCopy flag
    write_mode: VisionMode,

    // --- Read Path State ---
    /// Outer TLS packetizer for read path (only used in PaddingTls mode)
    /// Packetizes TCP stream into complete TLS records before feeding to rustls
    outer_read_packetizer: Option<TlsPacketizer>,

    /// Unpadding state machine for read path
    read_unpadder: VisionUnpadder,

    /// Partial unpadding data accumulated across multiple TLS records
    /// Used when unpadder returns None (needs more data)
    partial_unpad_buffer: BytesMut,

    /// Inner TLS packetizer for filtering read packets
    inner_read_packetizer: TlsPacketizer,

    // --- Write Path State ---
    /// Inner TLS packetizer for filtering write packets
    inner_write_packetizer: TlsPacketizer,
    /// Whether this is the first write (includes UUID in padding)
    write_first_packet: bool,

    /// Whether we've sent the DIRECT command on write path
    write_direct_sent: bool,

    // --- Shared State ---
    /// User UUID for padding validation
    user_uuid: [u8; 16],

    /// TLS pattern filter (analyzes INNER TLS)
    filter: VisionFilter,

    /// Buffer for leftover data when switching modes or when output buffer is too small
    pending_read: BytesMut,

    /// Whether we need to read VLESS response header (client-side only)
    /// Similar to shadowsocks's `is_initial_read` pattern
    vless_response_pending: bool,

    /// Partial VLESS response data accumulated across multiple TLS records
    /// Used to handle cases where the response header is split across TLS records
    partial_vless_response: BytesMut,

    /// Flag indicating peer has already switched to direct mode (detected via TLS error)
    /// When true, we should NOT send a DIRECT command since peer can't decrypt it
    peer_switched_direct: bool,

    /// Flag indicating we should switch to direct mode on the NEXT write call
    /// Set when we send DIRECT command, checked at the start of the next poll_write
    /// This matches Xray-core's behavior: DIRECT goes through TLS, next write is raw
    pending_direct_mode_switch: bool,

    /// Flag indicating we should switch to Tls mode on the NEXT write call
    /// Set when we send END command (no XTLS support or TLS 1.2), checked at start of next poll_write
    pending_tls_mode_switch: bool,

    /// Sleep future for direct mode delay
    /// Created when setting pending_direct_mode_switch, polled before actually switching
    /// This gives the network time to transmit the DIRECT packet before raw data follows
    direct_mode_delay: Option<Pin<Box<tokio::time::Sleep>>>,

    /// Reusable buffer for TLS read operations (TCP → read_tls → decrypt)
    /// Avoids allocating a new buffer on every TLS read
    tls_read_buffer: Vec<u8>,

    /// Buffer for plaintext data waiting to be written to the TLS session buffer
    /// When the rustls Writer buffer fills (WouldBlock), we store the remainder here
    /// This provides backpressure and state persistence across poll calls
    pending_plain_writes: BytesMut,

    /// Buffer for raw TCP data waiting to be written to the underlying stream
    /// This contains either:
    /// 1. Encrypted TLS records extracted from session via write_tls(), or
    /// 2. Partial TCP writes when TCP poll_write returns Poll::Pending
    /// This maintains state across poll calls to avoid data corruption
    pending_tcp_writes: BytesMut,
}

impl<IO> VisionStream<IO>
where
    IO: AsyncRead + AsyncWrite + Unpin,
{
    /// Create a new VisionStream for server-side (inbound) connections
    ///
    /// This takes ownership of both the underlying IO and the TLS server session.
    pub fn new_server(tcp: IO, session: ServerConnection, user_uuid: [u8; 16]) -> Self {
        Self::new_common(tcp, TlsConnection::Server(session), user_uuid, false)
    }

    /// Create a new VisionStream for client-side connections with VLESS response handling
    pub fn new_client_with_vless_response(
        tcp: IO,
        session: ClientConnection,
        user_uuid: [u8; 16],
    ) -> Self {
        Self::new_common(tcp, TlsConnection::Client(session), user_uuid, true)
    }

    fn new_common(
        tcp: IO,
        session: TlsConnection,
        user_uuid: [u8; 16],
        vless_response_pending: bool,
    ) -> Self {
        Self {
            tcp,
            session: Some(session),
            tls_read_buffer: vec![0u8; 16384],
            read_mode: VisionMode::PaddingTls,
            write_mode: VisionMode::PaddingTls,
            outer_read_packetizer: Some(TlsPacketizer::new()),
            inner_read_packetizer: TlsPacketizer::new(),
            inner_write_packetizer: TlsPacketizer::new(),
            read_unpadder: VisionUnpadder::new(user_uuid),
            partial_unpad_buffer: BytesMut::new(),
            write_first_packet: true,
            write_direct_sent: false,
            user_uuid,
            filter: VisionFilter::new(),
            pending_read: BytesMut::new(),
            vless_response_pending,
            partial_vless_response: BytesMut::new(),
            peer_switched_direct: false,
            pending_direct_mode_switch: false,
            pending_tls_mode_switch: false,
            direct_mode_delay: None,
            pending_plain_writes: BytesMut::new(),
            pending_tcp_writes: BytesMut::new(),
        }
    }

    /// Feed initial read data that was already decrypted by TLS but needs VISION unpadding
    ///
    /// This is used when the VLESS header parser read extra bytes from the TLS stream.
    /// Those bytes are already TLS-decrypted but still VISION-padded.
    pub fn feed_initial_read_data(&mut self, data: &[u8]) -> std::io::Result<()> {
        assert!(
            !self.vless_response_pending,
            "should only be called server-side"
        );

        if data.is_empty() {
            return Ok(());
        }

        log::debug!(
            "VISION: Feeding {} initial bytes (already TLS-decrypted, needs unpadding)",
            data.len()
        );

        // Process through VISION unpadding
        let UnpadResult {
            content: unpadded,
            command: maybe_command,
        } = self.read_unpadder.unpad(data)?;

        if !unpadded.is_empty() {
            // Filter TLS patterns (read path)
            Self::filter_data(&mut self.inner_read_packetizer, &mut self.filter, &unpadded);
        }

        // Check if we should switch modes
        // TODO: fix this
        if let Some(cmd) = maybe_command {
            match cmd {
                UnpadCommand::Direct => {
                    log::debug!("VISION: Initial data contains DIRECT command");
                    self.switch_read_to_direct_mode()?;
                    self.peer_switched_direct = true;
                }
                UnpadCommand::End => {
                    log::debug!("VISION: Initial data contains END command, switching to Tls mode");
                    // End command means: stop padding, continue normal relay through outer TLS
                    // This happens for non-TLS (HTTP) or TLS 1.2/below traffic
                    self.read_mode = VisionMode::Tls;
                }
                UnpadCommand::Continue => {
                    // Continue with padding
                }
            }
        }

        // Store unpadded data in leftover buffer
        if !unpadded.is_empty() {
            log::debug!(
                "VISION: Storing {} unpadded bytes in leftover buffer",
                unpadded.len()
            );
            self.pending_read.extend_from_slice(&unpadded);
        }

        Ok(())
    }

    /// Process unpadded data through a packetizer and filter
    ///
    /// Feeds data to packetizer, extracts complete TLS packets,
    /// and calls filter.filter_packet() on each one
    fn filter_data(packetizer: &mut TlsPacketizer, filter: &mut VisionFilter, data: &[u8]) {
        if !filter.is_filtering() {
            return; // Filtering already complete
        }

        // Feed to packetizer
        packetizer.feed(data);

        // Process all complete TLS packets
        while filter.is_filtering() {
            match packetizer.next_packet() {
                Ok(Some(packet)) => {
                    filter.filter_packet(&packet);
                }
                Ok(None) => break, // Need more data
                Err(_) => {
                    // Invalid TLS packet - stop filtering
                    break;
                }
            }
        }
    }

    /// Switch READ side from TLS mode to direct I/O mode
    /// This preserves any buffered TLS data and unpads it before storing
    fn switch_read_to_direct_mode(&mut self) -> io::Result<()> {
        if self.read_mode == VisionMode::Direct {
            return Ok(()); // Already switched
        }

        log::debug!("VISION READ: Switching to direct copy mode (asymmetric - write side may still use padding)");

        // Set read mode to Direct FIRST to avoid inconsistent state
        self.read_mode = VisionMode::Direct;

        // We always preserve the TLS session even if both sides switch to direct mode.
        // The session won't be used for data transfer anymore, but keeping it alive
        // prevents premature connection closure. It will be cleaned up when the connection closes.
        log::debug!(
            "VISION READ: Switched to direct mode, keeping TLS session alive for clean shutdown"
        );

        Ok(())
    }

    /// Switch WRITE side from TLS mode to direct I/O mode
    /// This doesn't need to preserve any data, just updates the mode
    fn switch_write_to_direct_mode(&mut self) -> io::Result<()> {
        if self.write_mode == VisionMode::Direct {
            return Ok(()); // Already switched
        }

        log::debug!("VISION WRITE: Switching to direct copy mode (asymmetric - read side may still use padding)");

        // Set write mode to Direct
        self.write_mode = VisionMode::Direct;

        // Keep the TLS session alive even after both sides switch to direct mode.
        // The session won't be used for data transfer anymore, but dropping it immediately
        // can cause issues because the peer might not have processed the DIRECT command yet.
        // The session will be cleaned up naturally when the connection closes.
        if self.read_mode == VisionMode::Direct {
            log::debug!("VISION WRITE: Both sides now in direct mode, keeping TLS session alive for clean shutdown");
        } else {
            // Read side is still in padding mode, so we keep the session
            log::debug!(
                "VISION WRITE: Switched to direct mode, but keeping TLS session for read side"
            );
        }

        Ok(())
    }

    /// Switch WRITE side from PaddingTls to Tls mode
    /// This happens when XTLS is not supported or TLS 1.2 is detected
    fn switch_write_to_tls_mode(&mut self) -> io::Result<()> {
        if self.write_mode == VisionMode::Tls {
            return Ok(()); // Already switched
        }

        log::debug!("VISION WRITE: Switching to Tls mode (no XTLS support)");

        // Set write mode to Tls
        self.write_mode = VisionMode::Tls;

        Ok(())
    }

    /// Drain pending TCP writes to the underlying stream
    /// Returns Poll::Ready(Ok(())) when fully drained
    /// Returns Poll::Pending if TCP would block
    fn drain_pending_tcp_writes(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        while !self.pending_tcp_writes.is_empty() {
            let n = ready!(Pin::new(&mut self.tcp).poll_write(cx, &self.pending_tcp_writes))?;

            if n == 0 {
                return Poll::Ready(Err(io::Error::new(
                    io::ErrorKind::WriteZero,
                    "Failed to write pending TCP data",
                )));
            }

            log::debug!(
                "VISION WRITE: Flushed {} pending TCP bytes ({} remaining)",
                n,
                self.pending_tcp_writes.len() - n
            );

            self.pending_tcp_writes.advance(n);
        }

        Poll::Ready(Ok(()))
    }

    /// Drain pending TLS records from session to pending_tcp_writes buffer
    /// Returns Poll::Ready(Ok(n)) where n is the number of bytes added to pending_tcp_writes
    /// Returns Poll::Pending if:
    ///   - We already have pending TCP writes (shouldn't extract more until TCP is clear), OR
    ///   - Already at threshold even though session.wants_write(), OR
    ///   - Session doesn't want to write
    fn drain_pending_tls_writes(&mut self, _cx: &mut Context<'_>) -> Poll<io::Result<usize>> {
        use std::io::Cursor;

        const TCP_BUFFER_THRESHOLD: usize = 32 * 1024; // 32KB

        // Don't extract TLS records if we have pending TCP writes
        if !self.pending_tcp_writes.is_empty() {
            log::debug!(
                "VISION WRITE: Skipping TLS drain, {} TCP bytes still pending",
                self.pending_tcp_writes.len()
            );
            return Poll::Pending;
        }

        let session = self.session.as_mut().expect("session must exist");

        // Check if session wants to write
        if !session.wants_write() {
            return Poll::Pending;
        }

        // Check if we're already at threshold
        if self.pending_tcp_writes.len() >= TCP_BUFFER_THRESHOLD {
            log::debug!(
                "VISION WRITE: Already at TCP buffer threshold ({} bytes)",
                self.pending_tcp_writes.len()
            );
            return Poll::Pending;
        }

        let initial_len = self.pending_tcp_writes.len();

        // Extract TLS records from session until we hit the threshold
        while session.wants_write() && self.pending_tcp_writes.len() < TCP_BUFFER_THRESHOLD {
            // Extract ONE TLS record into buffer
            let mut tls_record_buf = vec![0u8; 16384];
            let mut cursor = Cursor::new(&mut tls_record_buf[..]);

            let encrypted_len = match session.write_tls(&mut cursor) {
                Ok(n) => n,
                Err(e) => {
                    return Poll::Ready(Err(io::Error::new(
                        io::ErrorKind::Other,
                        format!("TLS write_tls failed: {}", e),
                    )));
                }
            };

            if encrypted_len == 0 {
                break;
            }

            log::debug!(
                "VISION WRITE: Extracted {} encrypted bytes from TLS session to pending_tcp_writes",
                encrypted_len
            );

            // Store in pending_tcp_writes buffer
            self.pending_tcp_writes
                .extend_from_slice(&tls_record_buf[..encrypted_len]);
        }

        let bytes_added = self.pending_tcp_writes.len() - initial_len;

        if bytes_added > 0 {
            Poll::Ready(Ok(bytes_added))
        } else {
            Poll::Pending
        }
    }

    /// Drain pending plaintext writes to TLS session buffer
    /// Returns Poll::Ready(Ok(())) when fully drained
    /// Returns Poll::Pending if session buffer is full (needs TLS/TCP drain first)
    fn drain_pending_plain_writes(&mut self, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        use std::io::Write;

        let session = self
            .session
            .as_mut()
            .expect("session must exist when draining");

        if self.pending_plain_writes.is_empty() {
            return Poll::Ready(Ok(()));
        }

        log::debug!(
            "VISION WRITE: Draining {} pending plaintext bytes",
            self.pending_plain_writes.len()
        );

        // Try to write pending plaintext to TLS session
        while !self.pending_plain_writes.is_empty() {
            match session.writer().write(&self.pending_plain_writes) {
                Ok(n) => {
                    if n == 0 {
                        return Poll::Ready(Err(io::Error::new(
                            io::ErrorKind::WriteZero,
                            "Failed to write to TLS session buffer",
                        )));
                    }

                    self.pending_plain_writes.advance(n);
                }
                Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
                    // Session buffer full - return Pending so drain_all_writes can drain TLS/TCP first
                    log::debug!(
                        "VISION WRITE: Session buffer full ({} plaintext bytes remaining), returning Pending",
                        self.pending_plain_writes.len()
                    );
                    return Poll::Pending;
                }
                Err(e) => return Poll::Ready(Err(e)),
            }
        }

        Poll::Ready(Ok(()))
    }

    /// Drain all pending writes: plain → TLS → TCP
    ///
    /// This drains the write pipeline in priority order (lowest level first):
    /// 1. TCP buffer → socket (highest priority - creates space for everything)
    /// 2. TLS session → TCP buffer (medium priority - creates space in session)
    /// 3. Plain buffer → TLS session (lowest priority - writes to session)
    ///
    /// Returns Poll::Ready(Ok(())) when all buffers are fully drained
    /// Returns Poll::Pending if TCP blocks (backpressure - nothing else can make progress)
    fn drain_all_writes(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        loop {
            // Check if everything is drained
            let session_wants_write = self
                .session
                .as_ref()
                .expect("session must exist")
                .wants_write();

            if self.pending_plain_writes.is_empty()
                && !session_wants_write
                && self.pending_tcp_writes.is_empty()
            {
                return Poll::Ready(Ok(()));
            }

            // Priority 1: Drain TCP buffer → socket
            // This is the bottleneck - if TCP blocks, nothing else can make progress
            if !self.pending_tcp_writes.is_empty() {
                match self.drain_pending_tcp_writes(cx) {
                    Poll::Ready(Ok(())) => continue, // Drained, loop back to check next level
                    Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
                    Poll::Pending => {
                        // TCP socket blocked - this is the bottleneck, return immediately
                        // We can't extract from TLS (no space in TCP buffer) or write to session
                        // (session buffer will eventually fill up with no way to drain)
                        return Poll::Pending;
                    }
                }
            }

            // Priority 2: Drain TLS session → TCP buffer
            // Only try if TCP is clear (checked above)
            let session_wants_write = self
                .session
                .as_ref()
                .expect("session must exist")
                .wants_write();
            if session_wants_write {
                match self.drain_pending_tls_writes(cx) {
                    Poll::Ready(Ok(_bytes_added)) => continue, // Extracted, loop back to drain TCP
                    Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
                    Poll::Pending => {
                        // Can't extract from TLS (shouldn't happen since TCP is clear)
                        // Return Pending to avoid busy loop
                        return Poll::Pending;
                    }
                }
            }

            // Priority 3: Drain Plain buffer → TLS session
            // Only try if both TLS and TCP are clear (checked above)
            if !self.pending_plain_writes.is_empty() {
                match self.drain_pending_plain_writes(cx) {
                    Poll::Ready(Ok(())) => continue, // Wrote to session, loop back to drain TLS
                    Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
                    Poll::Pending => {
                        // Session buffer full - need to drain TLS/TCP first
                        // But we already tried TLS and TCP above, so we're stuck
                        return Poll::Pending;
                    }
                }
            }

            // If we reach here, nothing was ready to drain
            // This shouldn't happen based on the checks above, but return Pending to be safe
            return Poll::Pending;
        }
    }

    /// Write data to TLS session, handling WouldBlock by saving to pending buffer
    fn write_to_session(&mut self, data: &[u8]) -> io::Result<()> {
        use std::io::Write;

        // First, try to drain any existing pending data
        if !self.pending_plain_writes.is_empty() {
            // If there's already pending data, just append to it
            // (don't try to write, as that would require cx for polling)
            self.pending_plain_writes.extend_from_slice(data);
            return Ok(());
        }

        let session = self.session.as_mut().expect("session must exist");
        let mut written = 0;
        while written < data.len() {
            match session.writer().write(&data[written..]) {
                Ok(n) => {
                    if n == 0 {
                        return Err(io::Error::new(
                            io::ErrorKind::WriteZero,
                            "Failed to write to TLS session buffer",
                        ));
                    }
                    written += n;
                }
                Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
                    // Session buffer full, save remainder
                    log::debug!(
                        "VISION WRITE: Session buffer full, saving {} bytes to pending",
                        data.len() - written
                    );
                    self.pending_plain_writes
                        .extend_from_slice(&data[written..]);
                    return Ok(());
                }
                Err(e) => return Err(e),
            }
        }

        Ok(())
    }

    /// Write data in Tls mode (no VISION padding, just TLS encryption)
    fn poll_write_tls(&mut self, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
        // Drain any pending writes first (backpressure)
        ready!(self.drain_all_writes(cx))?;

        let session = self
            .session
            .as_mut()
            .expect("session must exist in Tls mode");

        // Write directly to TLS session (no VISION padding)
        let mut written = 0;
        while written < buf.len() {
            match session.writer().write(&buf[written..]) {
                Ok(n) => {
                    if n == 0 {
                        return Poll::Ready(Err(io::Error::new(
                            io::ErrorKind::WriteZero,
                            "Failed to write to TLS session",
                        )));
                    }
                    written += n;
                }
                Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
                    // Session buffer full, save remainder and drain
                    log::debug!(
                        "VISION WRITE TLS: Session buffer full, saving {} bytes to pending",
                        buf.len() - written
                    );
                    self.pending_plain_writes.extend_from_slice(&buf[written..]);
                    ready!(self.drain_all_writes(cx))?;
                    // Return Pending - next poll will drain pending_plain_writes
                    return Poll::Pending;
                }
                Err(e) => return Poll::Ready(Err(e)),
            }
        }

        // Drain all writes
        ready!(self.drain_all_writes(cx))?;

        Poll::Ready(Ok(buf.len()))
    }

    /// Read the VLESS response header from the TLS session
    ///
    /// The VLESS response consists of:
    /// - 1 byte: version (must be 0)
    /// - 1 byte: addon length
    /// - N bytes: addon data (if addon_length > 0)
    ///
    /// This is called lazily on the first read when vless_response_pending is true.
    /// Uses packetization to feed TLS records one at a time to avoid bad MAC errors.
    ///
    /// Returns any unused data that came after the response
    fn poll_read_vless_response(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<Vec<u8>>> {
        let mut session = self
            .session
            .as_mut()
            .expect("session must exist when reading VLESS response");

        log::debug!(
            "VLESS: poll_read_vless_response called, partial_vless_response has {} bytes",
            self.partial_vless_response.len()
        );

        loop {
            // Always read from TCP to feed the packetizer
            let mut temp_tcp_buf = [0u8; 8192];
            let mut read_buf = ReadBuf::new(&mut temp_tcp_buf);
            match ready!(Pin::new(&mut self.tcp).poll_read(cx, &mut read_buf)) {
                Ok(()) => {}
                Err(e) => return Poll::Ready(Err(e)),
            };

            let tcp_bytes = read_buf.filled();

            if tcp_bytes.is_empty() {
                return Poll::Ready(Err(io::Error::new(
                    io::ErrorKind::UnexpectedEof,
                    "Connection closed while reading VLESS response",
                )));
            }

            log::debug!(
                "VLESS: Read {} bytes from TCP for response header",
                tcp_bytes.len()
            );

            // Feed to outer read packetizer
            self.outer_read_packetizer
                .as_mut()
                .expect("outer_read_packetizer must exist in PaddingTls mode")
                .feed(tcp_bytes);

            // We have to handle all packets and drain the outer_read_packetizer because
            // poll_read_padded_tls() will always read from the TCP and assume there's outstanding
            // data.
            // We don't have to worry about consuming some data post-Direct command because the TLS
            // flow is one-directional atm and it's not possible to complete a handshake.
            // If the peer sends a Direct command without handshake completing, this will error out
            // when we feed the TLS record to the session.
            let tls_records = self
                .outer_read_packetizer
                .as_mut()
                .expect("outer_read_packetizer must exist in PaddingTls mode")
                .next_packets()?;

            for tls_record in tls_records.into_iter() {
                // Feed single TLS record to session
                let plaintext_len =
                    feed_and_process_tls_connection(&mut session, tls_record.as_ref())?;

                // Access the decrypted buffer and accumulate into partial_vless_response
                let mut reader = session.reader();

                let mut i = 0;
                while i < plaintext_len {
                    let decrypted_buf = match reader.fill_buf() {
                        Ok(buf) => buf,
                        Err(e) => return Poll::Ready(Err(e)),
                    };

                    if decrypted_buf.is_empty() {
                        return Poll::Ready(Err(io::Error::other(
                            "no plaintext data when some is available",
                        )));
                    }

                    // Accumulate decrypted data into partial buffer
                    self.partial_vless_response.extend_from_slice(decrypted_buf);

                    let consumed_len = decrypted_buf.len();
                    reader.consume(consumed_len);
                    i += consumed_len;
                }
            }

            log::debug!(
                "VLESS: Accumulated vless response with {} bytes",
                self.partial_vless_response.len()
            );

            // Try to parse the VLESS response from accumulated data
            if self.partial_vless_response.len() < 2 {
                // Need at least 2 bytes for header, try to read from the TCP stream again
                continue;
            }

            // Read the 2-byte header
            let version = self.partial_vless_response[0];
            let addon_length = self.partial_vless_response[1];

            // Validate version
            if version != 0 {
                return Poll::Ready(Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    format!("Invalid VLESS response version: {}", version),
                )));
            }

            // Check if we have enough data for the full response
            let total_response_len = 2 + addon_length as usize;
            if self.partial_vless_response.len() < total_response_len {
                // Need more data - read more TCP
                continue;
            }

            // We have the complete VLESS response! Consume it from the buffer
            self.partial_vless_response.advance(total_response_len);

            log::debug!(
                    "VLESS: Successfully parsed {} byte response header (version={}, addon_length={}), {} bytes remaining in buffer",
                    total_response_len,
                    version,
                    addon_length,
                    self.partial_vless_response.len()
                );

            // If there's any remaining data in partial buffer, it belongs to the actual payload
            let remaining_data =
                std::mem::replace(&mut self.partial_vless_response, BytesMut::new());
            return Poll::Ready(Ok(remaining_data.to_vec()));
        }
    }

    fn poll_read_padding_tls(
        &mut self,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        loop {
            // CRITICAL: Try to extract a TLS record from the packetizer FIRST before reading from TCP.
            // This prevents the bug where we process one TLS record, return data, then on the next
            // poll_read we skip the buffered second TLS record and try to read from TCP instead.
            // This causes the DIRECT command (which arrived in the same TCP read but is in the second
            // TLS record) to sit in the buffer for 30 seconds until curl times out.
            //
            // Strategy: Try next_packet() first. If it returns Some, process it. If it returns None,
            // THEN read more from TCP. This avoids the busy loop that would occur if we checked
            // pending_bytes() > 0, since pending bytes don't guarantee a complete TLS record exists.
            let tls_record = match self
                .outer_read_packetizer
                .as_mut()
                .expect("outer_read_packetizer must exist in PaddingTls mode")
                .next_packet()?
            {
                Some(record) => {
                    // Got a buffered TLS record! Process it without reading from TCP.
                    log::debug!(
                        "VISION READ: Found buffered {} byte TLS record, processing before reading TCP",
                        record.len()
                    );
                    Some(record)
                }
                None => {
                    // No complete TLS record available, need to read from TCP
                    None
                }
            };

            // If we didn't get a buffered TLS record, need to read from TCP
            let tls_record = if let Some(record) = tls_record {
                record
            } else {
                // Read from TCP to feed the packetizer
                let mut temp_tcp_buf = [0u8; 8192];
                let mut read_buf = ReadBuf::new(&mut temp_tcp_buf);
                match ready!(Pin::new(&mut self.tcp).poll_read(cx, &mut read_buf)) {
                    Ok(()) => {}
                    Err(e) => return Poll::Ready(Err(e)),
                };

                let tcp_bytes = read_buf.filled();

                if tcp_bytes.is_empty() {
                    // TCP connection closed - but we might still have buffered TLS records in the packetizer
                    // or decrypted data in the TLS session

                    // First, check if the packetizer has any remaining TLS records to process
                    let has_pending_records = self
                        .outer_read_packetizer
                        .as_ref()
                        .map(|p| p.pending_bytes() > 0)
                        .unwrap_or(false);

                    if has_pending_records {
                        // Process remaining TLS records in the packetizer
                        // Don't return EOF yet - loop back to process them
                        log::debug!(
                            "VISION READ: TCP closed but packetizer has pending data, processing..."
                        );
                        continue; // Loop back to try extracting buffered TLS records
                    } else {
                        // No pending TLS records, check if TLS session has remaining decrypted data
                        let mut session = self
                            .session
                            .as_mut()
                            .expect("session must exist in PaddingTls mode");

                        let mut reader = session.reader();
                        match reader.fill_buf() {
                            Ok(decrypted_buf) if !decrypted_buf.is_empty() => {
                                // We have buffered TLS data - return it before EOF
                                let len = buf.remaining().min(decrypted_buf.len());
                                buf.put_slice(&decrypted_buf[..len]);
                                reader.consume(len);
                                return Poll::Ready(Ok(()));
                            }
                            _ => {
                                // No more data anywhere - return EOF
                                log::debug!("VISION READ: TCP closed, no pending data, returning EOF");
                                return Poll::Ready(Ok(()));
                            }
                        }
                    }
                }

                // Got TCP data, log and feed to packetizer
                // Log first 64 bytes as hex for debugging
                let hex_len = tcp_bytes.len().min(64);
                let hex_dump = tcp_bytes[..hex_len]
                    .iter()
                    .map(|b| format!("{:02x}", b))
                    .collect::<Vec<_>>()
                    .join("");
                log::debug!("VISION READ: Read {} bytes from TCP, first{}hex={}", tcp_bytes.len(), if tcp_bytes.len() > 64 { "64" } else { "" }, hex_dump);

                // Feed to outer TLS packetizer
                self.outer_read_packetizer
                    .as_mut()
                    .expect("outer_read_packetizer must exist in PaddingTls mode")
                    .feed(tcp_bytes);

                // Try to extract a TLS record after feeding TCP data
                match self
                    .outer_read_packetizer
                    .as_mut()
                    .expect("outer_read_packetizer must exist in PaddingTls mode")
                    .next_packet()?
                {
                    Some(record) => record,
                    None => {
                        log::debug!("VISION READ: No complete TLS record yet after reading TCP, need more data");
                        continue; // Loop back to read more TCP
                    }
                }
            };

            log::debug!(
                "VISION READ: Processing {} byte TLS record",
                tls_record.len()
            );

                let mut session = self
                    .session
                    .as_mut()
                    .expect("session must exist in PaddingTls mode");

                // Feed single TLS record to session
                let plaintext_len =
                    feed_and_process_tls_connection(&mut session, tls_record.as_ref())?;

                if plaintext_len == 0 {
                    continue; // Process next TLS record
                }

                let mut reader = session.reader();

                let mut decrypted = vec![];
                let mut i = 0;
                while i < plaintext_len {
                    let decrypted_part = match reader.fill_buf() {
                        Ok(buf) => buf,
                        Err(e) => {
                            // TODO: better error eg expected available plaintext but got error
                            return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, e)));
                        }
                    };

                    let part_len = decrypted_part.len();
                    decrypted.extend_from_slice(decrypted_part);
                    reader.consume(part_len);
                    i += part_len;
                }

                log::debug!("VISION READ: Got {} decrypted bytes", decrypted.len());
                self.handle_padded_bytes(decrypted)?;

                if !self.pending_read.is_empty() {
                    let len = buf.remaining().min(self.pending_read.len());
                    buf.put_slice(&self.pending_read[..len]);
                    self.pending_read.advance(len);
                    return Poll::Ready(Ok(()));
                }

            // Check if mode changed
            match self.read_mode {
                VisionMode::PaddingTls => {
                    // Need more TCP data, loop back to read
                }
                VisionMode::Tls => {
                    return self.poll_read_tls(cx, buf);
                }
                VisionMode::Direct => {
                    return Pin::new(&mut self.tcp).poll_read(cx, buf);
                }
            }
            // Need more TCP data, loop back to read
        }
    }

    fn handle_padded_bytes(&mut self, decrypted: Vec<u8>) -> std::io::Result<()> {
        // Feed newly decrypted bytes to the unpadder.
        // All data should be consumed from the rustls session, as any actual content
        // will eventually be returned by the unpadder.
        // The unpadder maintains internal state for partial parsing
        let UnpadResult {
            content: mut unpadded,
            command: maybe_command,
        } = self.read_unpadder.unpad(&decrypted)?;

        log::debug!("VISION READ: Unpadded to {} bytes", unpadded.len());

        if !unpadded.is_empty() {
            // Filter TLS patterns (INNER TLS detection)
            Self::filter_data(&mut self.inner_read_packetizer, &mut self.filter, &unpadded);
        }

        match maybe_command {
            Some(UnpadCommand::Direct) => {
                log::debug!("VISION READ: Received DIRECT command");

                // Extract remaining data from outer read packetizer
                let remaining = self
                    .outer_read_packetizer
                    .take()
                    .expect("outer_read_packetizer must exist")
                    .into_remaining_data();

                // Switch to Direct mode
                self.read_mode = VisionMode::Direct;

                if !remaining.is_empty() {
                    log::debug!(
                        "VISION READ: Extracted {} bytes from packetizer as raw data",
                        remaining.len()
                    );
                    unpadded.extend_from_slice(&remaining);
                }

                self.peer_switched_direct = true;

                self.pending_read.extend_from_slice(&unpadded);
            }
            Some(UnpadCommand::End) => {
                log::debug!("VISION READ: Received END command, switching to Tls mode");

                // Extract remaining data from outer read packetizer
                let remaining = self
                    .outer_read_packetizer
                    .take()
                    .expect("outer_read_packetizer must exist")
                    .into_remaining_data();

                // Switch to Tls mode - End command means: stop padding, continue normal relay through outer TLS
                // This happens for non-TLS (HTTP) or TLS 1.2/below traffic
                self.read_mode = VisionMode::Tls;

                // The `unpadded` data is already decrypted, store it in pending_read to be
                // returned
                if !unpadded.is_empty() {
                    log::debug!(
                        "VISION READ: Storing {} decrypted bytes in pending_read (End command - already decrypted)",
                        unpadded.len()
                    );
                    self.pending_read.extend_from_slice(&unpadded);
                }

                // The `remaining` data from the packetizer is ENCRYPTED TLS record bytes
                // We need to feed it to the TLS session for decryption, not append to unpadded
                if !remaining.is_empty() {
                    log::debug!(
                        "VISION READ: Extracted {} bytes from packetizer (encrypted TLS records)",
                        remaining.len()
                    );

                    let mut session = self
                        .session
                        .as_mut()
                        .expect("session must exist when handling End command");

                    // Feed the encrypted TLS records to the session
                    let plaintext_len = feed_and_process_tls_connection(&mut session, &remaining)?;

                    if plaintext_len > 0 {
                        // Extract the decrypted data
                        let mut reader = session.reader();

                        // Read all available decrypted data, possibly in multiple chunks
                        let mut total_read = 0;
                        while total_read < plaintext_len {
                            let chunk = match reader.fill_buf() {
                                Ok(buf) => buf,
                                Err(e) => {
                                    return Err(io::Error::new(io::ErrorKind::InvalidData, e));
                                }
                            };

                            let chunk_len = chunk.len();
                            self.pending_read.extend_from_slice(&chunk);
                            reader.consume(chunk_len);
                            total_read += chunk_len;
                        }

                        log::debug!(
                            "VISION READ: Decrypted {} bytes from packetizer remaining data (End command)",
                            total_read
                        );
                    }
                }
            }
            Some(UnpadCommand::Continue) => {
                self.pending_read.extend_from_slice(&unpadded);
            }
            None => {
                self.pending_read.extend_from_slice(&unpadded);
                log::debug!("VISION READ: Unpadder needs more ddata, will read more TLS records");
            }
        }
        Ok(())
    }

    fn poll_read_tls(
        &mut self,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        // Regular TLS mode (no padding)
        // Similar to poll_read_vless_response logic
        let mut session = self
            .session
            .as_mut()
            .expect("session must exist in Tls mode");

        {
            let mut reader = session.reader();
            match reader.fill_buf() {
                Ok(decrypted_buf) => {
                    if !decrypted_buf.is_empty() {
                        let len = buf.remaining().min(decrypted_buf.len());
                        buf.put_slice(&decrypted_buf[..len]);
                        reader.consume(len);
                        return Poll::Ready(Ok(()));
                    }
                }
                Err(e) if e.kind() == io::ErrorKind::WouldBlock => {}
                Err(e) => return Poll::Ready(Err(e)),
            }
        }

        loop {
            let mut read_buf = ReadBuf::new(&mut self.tls_read_buffer);
            match ready!(Pin::new(&mut self.tcp).poll_read(cx, &mut read_buf)) {
                Ok(()) => {}
                Err(e) => return Poll::Ready(Err(e)),
            };
            let tcp_bytes_read = read_buf.filled().len();

            if tcp_bytes_read == 0 {
                // TCP connection closed - but we might still have buffered data

                // Check if TLS session has any remaining decrypted data
                let mut reader = session.reader();
                match reader.fill_buf() {
                    Ok(decrypted_buf) if !decrypted_buf.is_empty() => {
                        // We have buffered TLS data - process it before EOF
                        let len = buf.remaining().min(decrypted_buf.len());
                        buf.put_slice(&decrypted_buf[..len]);
                        reader.consume(len);
                        return Poll::Ready(Ok(()));
                    }
                    _ => {
                        // No more data - return EOF
                        return Poll::Ready(Ok(()));
                    }
                }
            }

            // Feed to TLS session
            let plaintext_len = feed_and_process_tls_connection(
                &mut session,
                &self.tls_read_buffer[..tcp_bytes_read],
            )?;

            // Access the decrypted buffer
            let mut reader = session.reader();
            let mut decrypted = vec![];
            let mut i = 0;
            while i < plaintext_len {
                let decrypted_part = match reader.fill_buf() {
                    Ok(buf) => buf,
                    Err(e) => {
                        // TODO: better error eg expected available plaintext but got error
                        return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, e)));
                    }
                };

                let part_len = decrypted_part.len();
                decrypted.extend_from_slice(decrypted_part);
                reader.consume(part_len);
                i += part_len;
            }

            if decrypted.is_empty() {
                continue;
            }

            let len = buf.remaining().min(decrypted.len());
            buf.put_slice(&decrypted[..len]);
            // NOTE: We already consumed all data from reader when building the decrypted vector
            // (lines 1296-1309). Do NOT consume again here - that would be a double consume
            // and violates rustls's BufRead contract, causing a panic.
            self.pending_read.extend_from_slice(&decrypted[len..]);

            return Poll::Ready(Ok(()));
        }
    }

    fn poll_write_padding_tls(
        &mut self,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        // STEP 1: Drain all pending writes (backpressure) - CRITICAL: must drain before mode switch!
        // The session could contain the DIRECT/END packet that triggered the mode switch flag.
        // If we switch modes before draining, that packet might never get sent!
        {
            let session = self
                .session
                .as_mut()
                .expect("session must exist in PaddingTls mode");

            if !self.pending_plain_writes.is_empty()
                || session.wants_write()
                || !self.pending_tcp_writes.is_empty()
            {
                ready!(self.drain_all_writes(cx))?;
            }
        }

        // STEP 3: Check if we should switch to direct mode from previous write
        // This flag is set AFTER sending the DIRECT packet, so by now the packet is drained
        if self.pending_direct_mode_switch {
            // If delay is enabled, poll the sleep future first
            if let Some(sleep) = &mut self.direct_mode_delay {
                // Poll the sleep future
                use std::future::Future;
                ready!(sleep.as_mut().poll(cx));

                // Sleep completed, clear it
                self.direct_mode_delay = None;
                log::debug!("VISION WRITE: Direct mode delay completed (5ms elapsed)");
            }

            log::debug!("VISION WRITE: Switching to direct mode (flag set by previous write)");
            self.pending_direct_mode_switch = false;
            self.switch_write_to_direct_mode()?;
            return Pin::new(&mut self.tcp).poll_write(cx, buf);
        }

        if self.pending_tls_mode_switch {
            log::debug!("VISION WRITE: Switching to Tls mode (flag set by previous write)");
            self.pending_tls_mode_switch = false;
            self.switch_write_to_tls_mode()?;
            return self.poll_write_tls(cx, buf);
        }

        // STEP 4: Feed write buffer to inner packetizer and check for ApplicationData
        let existing_inner_len = self.inner_write_packetizer.pending_bytes();
        self.inner_write_packetizer.feed(buf);

        // Process all complete TLS packets in the buffer
        let mut processed_len = 0;
        loop {
            match self.inner_write_packetizer.next_packet() {
                Ok(Some(packet)) => {
                    processed_len += packet.len();

                    // Feed to filter for TLS pattern detection
                    self.filter.filter_packet(&packet);

                    // Check if this packet is ApplicationData
                    if self.filter.is_tls() && packet.len() >= 3
                        && packet[0] == 0x17  // ApplicationData
                        && packet[1] == 0x03
                    // TLS version major
                    {
                        log::debug!(
                            "VISION WRITE: Detected ApplicationData in {} byte packet",
                            packet.len()
                        );
                        // DIRECT or END
                        let command = if self.filter.supports_xtls() {
                            0x02
                        } else {
                            0x01
                        };

                        let final_padded_packet = if self.write_first_packet {
                            self.write_first_packet = false;
                            super::vision_pad::pad_with_uuid_and_command(
                                &packet,
                                &self.user_uuid,
                                command,
                                true, // is_tls
                            )
                        } else {
                            super::vision_pad::pad_with_command(
                                &packet, command, true, // is_tls
                            )
                        };

                        self.write_to_session(&final_padded_packet)?;

                        // Set flag to switch on next write
                        if self.filter.supports_xtls() {
                            self.pending_direct_mode_switch = true;

                            // Create sleep future if delay is enabled
                            // This gives the network time to transmit the DIRECT packet before raw data
                            if ENABLE_DIRECT_MODE_DELAY {
                                log::debug!("VISION WRITE: Creating 5ms delay before switching to direct mode");
                                self.direct_mode_delay = Some(Box::pin(tokio::time::sleep(
                                    std::time::Duration::from_millis(5),
                                )));
                            }
                        } else {
                            self.pending_tls_mode_switch = true;
                        }

                        // this must be true because else we would have a successful next_packet call on previous iteration
                        assert!(processed_len > existing_inner_len);

                        // Drain and handle result
                        match self.drain_all_writes(cx) {
                            Poll::Ready(Ok(())) => {
                                // Fully drained, continue
                            }
                            Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
                            Poll::Pending => {
                                // Still draining - clear packetizer since caller will re-feed remaining data
                                self.inner_write_packetizer.clear();
                            }
                        }

                        // tell caller to resend the next packets after the mode switch
                        return Poll::Ready(Ok(processed_len - existing_inner_len));
                    }

                    if !self.filter.is_filtering() && !self.filter.is_tls12_or_above() {
                        let final_padded_packet = if self.write_first_packet {
                            self.write_first_packet = false;
                            super::vision_pad::pad_with_uuid_and_command(
                                &packet,
                                &self.user_uuid,
                                0x01, // END
                                self.filter.is_tls(),
                            )
                        } else {
                            super::vision_pad::pad_with_command(
                                &packet,
                                0x01, // END
                                self.filter.is_tls(),
                            )
                        };

                        self.write_to_session(&final_padded_packet)?;

                        // Set flag to switch on next write
                        self.pending_tls_mode_switch = true;

                        // this must be true because else we would have a successful next_packet call on previous iteration
                        assert!(processed_len > existing_inner_len);

                        // Drain and handle result
                        match self.drain_all_writes(cx) {
                            Poll::Ready(Ok(())) => {
                                // Fully drained, continue
                            }
                            Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
                            Poll::Pending => {
                                // Still draining - clear packetizer since caller will re-feed remaining data
                                self.inner_write_packetizer.clear();
                            }
                        }

                        // tell caller to resend the next packets after the mode switch
                        return Poll::Ready(Ok(processed_len - existing_inner_len));
                    }

                    // if we got here, it's a TLS packet that is not app data and we need to continue padding
                    // at this point is_filtering() might already be false, but we continue as we are tls 1.2 or above
                    // ref: https://github.com/XTLS/Xray-core/blob/9f5dcb15910aadc7ef450514747576827a389853/proxy/proxy.go#L371

                    let padded_packet = if self.write_first_packet {
                        self.write_first_packet = false;
                        super::vision_pad::pad_with_uuid_and_command(
                            &packet,
                            &self.user_uuid,
                            0x00, // CONTINUE
                            self.filter.is_tls(),
                        )
                    } else {
                        super::vision_pad::pad_with_command(
                            &packet,
                            0x00, // CONTINUE
                            self.filter.is_tls(),
                        )
                    };

                    self.write_to_session(&padded_packet)?;
                    match self.drain_all_writes(cx) {
                        Poll::Pending => {
                            let unprocessed_buf_len =
                                buf.len() - (processed_len - existing_inner_len);
                            // sanity check
                            assert!(
                                self.inner_write_packetizer.pending_bytes() == unprocessed_buf_len
                            );
                            // clear since the user will re-feed
                            self.inner_write_packetizer.clear();
                            return Poll::Ready(Ok(processed_len - existing_inner_len));
                        }
                        Poll::Ready(Ok(())) => {
                            // continue since it drained
                        }
                        Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
                    }
                }
                Ok(None) => {
                    // need more data
                    // TODO: investigate this case. it's possible that this is not TLS data and the packetizer is waiting
                    // for more writes from the caller to complete a full TLS record, but no more writes are coming.
                    // do we need to change our packetizer approach?
                    break;
                }
                Err(e) => {
                    log::debug!("VISION WRITE: Packetization failed: {}", e);
                    // packetization failed, to match Xray-core behavior, we treat this as a single packet
                    // TODO: check Xray-core about padding lengths and see if we need to split into multiple packets

                    self.filter.decrement_filter_count();

                    // continue or end command
                    let command = if self.filter.is_filtering() {
                        0x00
                    } else {
                        self.pending_tls_mode_switch = true;
                        log::debug!(
                            "VISION WRITE: Switching to TLS mode after packetization failure"
                        );
                        0x01
                    };
                    let remaining_data =
                        std::mem::replace(&mut self.inner_write_packetizer, TlsPacketizer::new())
                            .into_remaining_data();

                    let padded_packet = if self.write_first_packet {
                        self.write_first_packet = false;
                        super::vision_pad::pad_with_uuid_and_command(
                            &remaining_data,
                            &self.user_uuid,
                            command,
                            self.filter.is_tls(),
                        )
                    } else {
                        super::vision_pad::pad_with_command(
                            &remaining_data,
                            command,
                            self.filter.is_tls(),
                        )
                    };

                    self.write_to_session(&padded_packet)?;
                    ready!(self.drain_all_writes(cx))?;

                    return Poll::Ready(Ok(buf.len()));
                }
            }
        }

        // we processed all packets if we got here
        Poll::Ready(Ok(buf.len()))
    }
}

// AsyncRead implementation
impl<IO> AsyncRead for VisionStream<IO>
where
    IO: AsyncRead + AsyncWrite + Unpin,
{
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        let this = self.get_mut();

        // If we have leftover buffered data, return it first
        if !this.pending_read.is_empty() {
            let len = buf.remaining().min(this.pending_read.len());
            buf.put_slice(&this.pending_read[..len]);
            this.pending_read.advance(len);
            return Poll::Ready(Ok(()));
        }

        match this.read_mode {
            VisionMode::PaddingTls => {
                if this.vless_response_pending {
                    let decrypted_data = ready!(this.poll_read_vless_response(cx))?;
                    this.vless_response_pending = false;
                    if !decrypted_data.is_empty() {
                        this.handle_padded_bytes(decrypted_data)?;
                        if !this.pending_read.is_empty() {
                            let len = buf.remaining().min(this.pending_read.len());
                            buf.put_slice(&this.pending_read[..len]);
                            this.pending_read.advance(len);
                            return Poll::Ready(Ok(()));
                        }
                        match this.read_mode {
                            VisionMode::PaddingTls => {
                                // Fall-through
                            }
                            VisionMode::Tls => {
                                return this.poll_read_tls(cx, buf);
                            }
                            VisionMode::Direct => {
                                return Pin::new(&mut this.tcp).poll_read(cx, buf);
                            }
                        }
                    }
                }
                this.poll_read_padding_tls(cx, buf)
            }
            VisionMode::Tls => this.poll_read_tls(cx, buf),
            VisionMode::Direct => Pin::new(&mut this.tcp).poll_read(cx, buf),
        }
    }
}

// AsyncWrite implementation
impl<IO> AsyncWrite for VisionStream<IO>
where
    IO: AsyncRead + AsyncWrite + Unpin,
{
    fn poll_write(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        let this = self.get_mut();

        match this.write_mode {
            VisionMode::PaddingTls => this.poll_write_padding_tls(cx, buf),
            VisionMode::Tls => this.poll_write_tls(cx, buf),
            VisionMode::Direct => Pin::new(&mut this.tcp).poll_write(cx, buf),
        }
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        let this = self.get_mut();

        match this.write_mode {
            VisionMode::PaddingTls | VisionMode::Tls => {
                // Drain all pending writes
                ready!(this.drain_all_writes(cx))?;
                Pin::new(&mut this.tcp).poll_flush(cx)
            }
            VisionMode::Direct => Pin::new(&mut this.tcp).poll_flush(cx),
        }
    }

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

// Implement AsyncPing trait
impl<IO> crate::async_stream::AsyncPing for VisionStream<IO>
where
    IO: AsyncStream,
{
    fn supports_ping(&self) -> bool {
        self.tcp.supports_ping()
    }

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

impl<IO> AsyncStream for VisionStream<IO> where IO: AsyncStream {}
