//! Fake TLS 1.3 stream wrappers
//!
//! This module provides stateful async stream wrappers that handle TLS record
//! framing with proper partial read/write handling.
//!
//! These are "fake" TLS streams:
//! - We wrap raw bytes into syntactically valid TLS 1.3 records (Application Data).
//! - We DO NOT perform real TLS handshake/encryption.
//! - Real crypto for MTProto is handled by the crypto layer underneath.
//!
//! Why do we need this?
//! Telegram MTProto proxy "FakeTLS" mode uses a TLS-looking outer layer for
//! domain fronting / traffic camouflage. iOS Telegram clients are known to
//! produce slightly different TLS record sizing patterns than Android/Desktop,
//! including records that exceed MAX_TLS_PLAINTEXT_SIZE payload bytes by a small overhead.
//!
//! Key design principles:
//! - Explicit state machines for all async operations
//! - Never lose data on partial reads
//! - Atomic TLS record formation for writes

#![allow(dead_code)]
//! - Proper handling of all TLS record types
//!
//! Important nuance (Telegram FakeTLS):
//! - The TLS spec limits "plaintext fragments" to MAX_TLS_PLAINTEXT_SIZE bytes.
//! - However, the on-the-wire record length can exceed MAX_TLS_PLAINTEXT_SIZE because TLS 1.3
//!   uses AEAD and can include tag/overhead/padding.
//! - Telegram FakeTLS clients (notably iOS) may send Application Data records
//!   with length up to MAX_TLS_CIPHERTEXT_SIZE bytes (RFC 8446 §5.2).
//!
//! If you reject those (e.g. validate length <= MAX_TLS_PLAINTEXT_SIZE), you will see errors like:
//!   "TLS record too large: 16408 bytes"
//! and uploads from iOS will break (media/file sending), while small traffic
//! may still work.

use bytes::{Bytes, BytesMut};
use std::io::{self, Error, ErrorKind, Result};
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf};

use super::state::{HeaderBuffer, StreamState, YieldBuffer};
use crate::protocol::constants::{
    MAX_TLS_CIPHERTEXT_SIZE, MAX_TLS_PLAINTEXT_SIZE, TLS_RECORD_ALERT, TLS_RECORD_APPLICATION,
    TLS_RECORD_CHANGE_CIPHER, TLS_RECORD_HANDSHAKE, TLS_VERSION,
};

// ============= Constants =============

/// TLS record header size (type + version + length)
const TLS_HEADER_SIZE: usize = 5;

/// Maximum TLS fragment size we emit for Application Data.
/// Real TLS 1.3 allows up to 16384 + 256 bytes of ciphertext (incl. tag).
const MAX_TLS_PAYLOAD: usize = MAX_TLS_CIPHERTEXT_SIZE;

/// Maximum pending write buffer for one record remainder.
/// Note: we never queue unlimited amount of data here; state holds at most one record.
const MAX_PENDING_WRITE: usize = 64 * 1024;

/// Maximum record buffer capacity retained between writes.
const MAX_RETAINED_RECORD_CAPACITY: usize = 2 * (TLS_HEADER_SIZE + MAX_TLS_PAYLOAD);

// ============= TLS Record Types =============

/// Parsed TLS record header (5 bytes)
#[derive(Debug, Clone, Copy)]
struct TlsRecordHeader {
    /// Record type (0x17 = Application Data, 0x14 = Change Cipher, etc.)
    record_type: u8,
    /// TLS version bytes
    version: [u8; 2],
    /// Payload length
    length: u16,
}

impl TlsRecordHeader {
    /// Parse header from exactly 5 bytes.
    ///
    /// This currently never returns None, but is kept as Option to allow future
    /// stricter parsing rules without changing callers.
    fn parse(header: &[u8; 5]) -> Option<Self> {
        let record_type = header[0];
        let version = [header[1], header[2]];
        let length = u16::from_be_bytes([header[3], header[4]]);
        Some(Self {
            record_type,
            version,
            length,
        })
    }

    /// Validate the header.
    ///
    /// Nuances:
    /// - We accept TLS 1.0 header version for ClientHello-like records (0x03 0x01),
    ///   and TLS 1.2/1.3 style version bytes for the rest (we use TLS_VERSION = 0x03 0x03).
    /// - For Application Data, Telegram FakeTLS may send payload length up to
    ///   MAX_TLS_CIPHERTEXT_SIZE (16384 + 256).
    /// - For other record types we keep stricter bounds to avoid memory abuse.
    fn validate(&self) -> Result<()> {
        // Version precheck: only 0x0301 and 0x0303 are recognized at all.
        if self.version != [0x03, 0x01] && self.version != TLS_VERSION {
            return Err(Error::new(
                ErrorKind::InvalidData,
                format!("Invalid TLS version: {:02x?}", self.version),
            ));
        }

        // Narrow FakeTLS wire profile: TLS 1.0 compatibility header is allowed
        // only on Handshake records (ClientHello compatibility quirk).
        if self.record_type != TLS_RECORD_HANDSHAKE && self.version != TLS_VERSION {
            return Err(Error::new(
                ErrorKind::InvalidData,
                format!(
                    "invalid TLS version for record type 0x{:02x}: {:02x?}",
                    self.record_type, self.version
                ),
            ));
        }

        let len = self.length as usize;

        // Length checks depend on record type.
        // Telegram FakeTLS: ApplicationData may use ciphertext envelope limit,
        // while control records stay structurally strict to reduce probe surface.
        match self.record_type {
            TLS_RECORD_APPLICATION => {
                if len == 0 || len > MAX_TLS_CIPHERTEXT_SIZE {
                    return Err(Error::new(
                        ErrorKind::InvalidData,
                        format!(
                            "invalid TLS application data length: {} (min 1, max {})",
                            len, MAX_TLS_CIPHERTEXT_SIZE
                        ),
                    ));
                }
            }

            TLS_RECORD_CHANGE_CIPHER => {
                if len != 1 {
                    return Err(Error::new(
                        ErrorKind::InvalidData,
                        format!("invalid TLS ChangeCipherSpec length: {} (expected 1)", len),
                    ));
                }
            }

            TLS_RECORD_ALERT => {
                if len != 2 {
                    return Err(Error::new(
                        ErrorKind::InvalidData,
                        format!("invalid TLS alert length: {} (expected 2)", len),
                    ));
                }
            }

            TLS_RECORD_HANDSHAKE => {
                if !(4..=MAX_TLS_PLAINTEXT_SIZE).contains(&len) {
                    return Err(Error::new(
                        ErrorKind::InvalidData,
                        format!(
                            "invalid TLS handshake length: {} (min 4, max {})",
                            len, MAX_TLS_PLAINTEXT_SIZE
                        ),
                    ));
                }
            }

            _ => {
                return Err(Error::new(
                    ErrorKind::InvalidData,
                    format!("unknown TLS record type: 0x{:02x}", self.record_type),
                ));
            }
        }

        Ok(())
    }

    /// Build header bytes
    fn to_bytes(self) -> [u8; 5] {
        [
            self.record_type,
            self.version[0],
            self.version[1],
            (self.length >> 8) as u8,
            self.length as u8,
        ]
    }
}

// ============= FakeTlsReader State =============

/// State machine states for FakeTlsReader
#[derive(Debug)]
enum TlsReaderState {
    /// Ready to read a new TLS record
    Idle,

    /// Reading the 5-byte TLS record header
    ReadingHeader {
        /// Header buffer (5 bytes)
        header: HeaderBuffer<TLS_HEADER_SIZE>,
    },

    /// Reading the TLS record body (payload)
    ReadingBody {
        record_type: u8,
        length: usize,
        buffer: BytesMut,
    },

    /// Have buffered data ready to yield to caller
    Yielding { buffer: YieldBuffer },

    /// Stream encountered an error and cannot be used
    Poisoned { error: Option<io::Error> },
}

impl StreamState for TlsReaderState {
    fn is_terminal(&self) -> bool {
        matches!(self, Self::Poisoned { .. })
    }

    fn is_poisoned(&self) -> bool {
        matches!(self, Self::Poisoned { .. })
    }

    fn state_name(&self) -> &'static str {
        match self {
            Self::Idle => "Idle",
            Self::ReadingHeader { .. } => "ReadingHeader",
            Self::ReadingBody { .. } => "ReadingBody",
            Self::Yielding { .. } => "Yielding",
            Self::Poisoned { .. } => "Poisoned",
        }
    }
}

// ============= FakeTlsReader =============

/// Reader that unwraps TLS records (FakeTLS).
///
/// This wrapper is responsible ONLY for TLS record framing and skipping
/// non-data records (like CCS). It does not decrypt TLS: payload bytes are passed
/// as-is to upper layers (crypto stream).
///
/// State machine overview:
///
/// ┌──────────┐                    ┌───────────────┐
/// │   Idle   │ -----------------> │ ReadingHeader │
/// └──────────┘                    └───────┬───────┘
///      ▲                                  │
///      │                           header complete
///      │                                  │
///      │                                  ▼
///      │                          ┌───────────────┐
///      │        skip record       │  ReadingBody  │
///      │ <-------- (CCS) -------- │               │
///      │                          └───────┬───────┘
///      │                                  │
///      │                              body complete
///      │                                  ▼
///      │                          ┌───────────────┐
///      │                          │   Yielding    │
///      │                          └───────────────┘
///      │
///      │    errors / w any state
///      ▼
/// ┌───────────────────────────────────────────────┐
/// │                    Poisoned                   │
/// └───────────────────────────────────────────────┘
///
/// NOTE: We must correctly handle partial reads from upstream:
/// - do not assume header arrives in one poll
/// - do not assume body arrives in one poll
/// - never lose already-read bytes
pub struct FakeTlsReader<R> {
    upstream: R,
    state: TlsReaderState,
    body_scratch: Vec<u8>,
}

impl<R> FakeTlsReader<R> {
    pub fn new(upstream: R) -> Self {
        Self {
            upstream,
            state: TlsReaderState::Idle,
            body_scratch: Vec::new(),
        }
    }

    pub fn get_ref(&self) -> &R {
        &self.upstream
    }
    pub fn get_mut(&mut self) -> &mut R {
        &mut self.upstream
    }
    pub fn into_inner(self) -> R {
        self.upstream
    }

    pub fn into_inner_with_pending_plaintext(mut self) -> (R, Vec<u8>) {
        let pending = match std::mem::replace(&mut self.state, TlsReaderState::Idle) {
            TlsReaderState::Yielding { buffer } => buffer.as_slice().to_vec(),
            TlsReaderState::ReadingBody {
                record_type,
                buffer,
                ..
            } if record_type == TLS_RECORD_APPLICATION => buffer.to_vec(),
            _ => Vec::new(),
        };
        (self.upstream, pending)
    }

    pub fn is_poisoned(&self) -> bool {
        self.state.is_poisoned()
    }
    pub fn state_name(&self) -> &'static str {
        self.state.state_name()
    }

    fn poison(&mut self, error: io::Error) {
        self.state = TlsReaderState::Poisoned { error: Some(error) };
    }

    fn take_poison_error(&mut self) -> io::Error {
        match &mut self.state {
            TlsReaderState::Poisoned { error } => error
                .take()
                .unwrap_or_else(|| io::Error::other("stream previously poisoned")),
            _ => io::Error::other("stream not poisoned"),
        }
    }
}

enum HeaderPollResult {
    Pending,
    Eof,
    Complete(TlsRecordHeader),
    Error(io::Error),
}

enum BodyPollResult {
    Pending,
    Complete(Bytes),
    Error(io::Error),
}

impl<R: AsyncRead + Unpin> AsyncRead for FakeTlsReader<R> {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<Result<()>> {
        let this = self.get_mut();

        loop {
            // Take ownership of state to avoid borrow conflicts
            let state = std::mem::replace(&mut this.state, TlsReaderState::Idle);

            match state {
                // Poisoned state: always return the stored error
                TlsReaderState::Poisoned { error } => {
                    this.state = TlsReaderState::Poisoned { error: None };
                    let err =
                        error.unwrap_or_else(|| io::Error::other("stream previously poisoned"));
                    return Poll::Ready(Err(err));
                }

                // Yield buffered plaintext to caller
                TlsReaderState::Yielding { mut buffer } => {
                    if buf.remaining() == 0 {
                        this.state = TlsReaderState::Yielding { buffer };
                        return Poll::Ready(Ok(()));
                    }

                    let to_copy = buffer.remaining().min(buf.remaining());
                    let dst = buf.initialize_unfilled_to(to_copy);
                    let copied = buffer.copy_to(dst);
                    buf.advance(copied);

                    if buffer.is_empty() {
                        this.state = TlsReaderState::Idle;
                    } else {
                        this.state = TlsReaderState::Yielding { buffer };
                    }

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

                // Start reading new record
                TlsReaderState::Idle => {
                    if buf.remaining() == 0 {
                        this.state = TlsReaderState::Idle;
                        return Poll::Ready(Ok(()));
                    }

                    this.state = TlsReaderState::ReadingHeader {
                        header: HeaderBuffer::new(),
                    };
                    // loop continues and will handle ReadingHeader
                }

                // Read TLS header (5 bytes)
                TlsReaderState::ReadingHeader { mut header } => {
                    let result = poll_read_header(&mut this.upstream, cx, &mut header);

                    match result {
                        HeaderPollResult::Pending => {
                            this.state = TlsReaderState::ReadingHeader { header };
                            return Poll::Pending;
                        }
                        HeaderPollResult::Eof => {
                            // Clean EOF at record boundary
                            this.state = TlsReaderState::Idle;
                            return Poll::Ready(Ok(()));
                        }
                        HeaderPollResult::Error(e) => {
                            this.poison(Error::new(e.kind(), e.to_string()));
                            return Poll::Ready(Err(e));
                        }
                        HeaderPollResult::Complete(parsed) => {
                            if let Err(e) = parsed.validate() {
                                this.poison(Error::new(e.kind(), e.to_string()));
                                return Poll::Ready(Err(e));
                            }

                            let length = parsed.length as usize;
                            this.state = TlsReaderState::ReadingBody {
                                record_type: parsed.record_type,
                                length,
                                buffer: BytesMut::with_capacity(length),
                            };
                        }
                    }
                }

                // Read TLS payload
                TlsReaderState::ReadingBody {
                    record_type,
                    length,
                    mut buffer,
                } => {
                    let result = poll_read_body(
                        &mut this.upstream,
                        cx,
                        &mut buffer,
                        length,
                        &mut this.body_scratch,
                    );

                    match result {
                        BodyPollResult::Pending => {
                            this.state = TlsReaderState::ReadingBody {
                                record_type,
                                length,
                                buffer,
                            };
                            return Poll::Pending;
                        }
                        BodyPollResult::Error(e) => {
                            this.poison(Error::new(e.kind(), e.to_string()));
                            return Poll::Ready(Err(e));
                        }
                        BodyPollResult::Complete(data) => {
                            match record_type {
                                TLS_RECORD_CHANGE_CIPHER => {
                                    // CCS is expected in some clients, ignore it.
                                    this.state = TlsReaderState::Idle;
                                    continue;
                                }

                                TLS_RECORD_APPLICATION => {
                                    // This is what we actually want.
                                    if data.is_empty() {
                                        this.state = TlsReaderState::Idle;
                                        continue;
                                    }

                                    this.state = TlsReaderState::Yielding {
                                        buffer: YieldBuffer::new(data),
                                    };
                                    // loop continues and will yield immediately
                                }

                                TLS_RECORD_ALERT => {
                                    // Treat TLS alert as EOF-like termination.
                                    this.state = TlsReaderState::Idle;
                                    return Poll::Ready(Ok(()));
                                }

                                TLS_RECORD_HANDSHAKE => {
                                    // After FakeTLS handshake is done, we do not expect any Handshake records.
                                    let err = Error::new(
                                        ErrorKind::InvalidData,
                                        "unexpected TLS handshake record",
                                    );
                                    this.poison(Error::new(err.kind(), err.to_string()));
                                    return Poll::Ready(Err(err));
                                }

                                _ => {
                                    let err = Error::new(
                                        ErrorKind::InvalidData,
                                        format!("unknown TLS record type: 0x{:02x}", record_type),
                                    );
                                    this.poison(Error::new(err.kind(), err.to_string()));
                                    return Poll::Ready(Err(err));
                                }
                            }
                        }
                    }
                }
            }
        }
    }
}

/// Poll to read and fill header buffer (standalone function to avoid borrow issues)
fn poll_read_header<R: AsyncRead + Unpin>(
    upstream: &mut R,
    cx: &mut Context<'_>,
    header: &mut HeaderBuffer<TLS_HEADER_SIZE>,
) -> HeaderPollResult {
    while !header.is_complete() {
        let unfilled = header.unfilled_mut();
        let mut read_buf = ReadBuf::new(unfilled);

        match Pin::new(&mut *upstream).poll_read(cx, &mut read_buf) {
            Poll::Pending => return HeaderPollResult::Pending,
            Poll::Ready(Err(e)) => return HeaderPollResult::Error(e),
            Poll::Ready(Ok(())) => {
                let n = read_buf.filled().len();
                if n == 0 {
                    // EOF
                    if header.as_slice().is_empty() {
                        return HeaderPollResult::Eof;
                    } else {
                        return HeaderPollResult::Error(Error::new(
                            ErrorKind::UnexpectedEof,
                            format!(
                                "unexpected EOF in TLS header (got {} of 5 bytes)",
                                header.as_slice().len()
                            ),
                        ));
                    }
                }
                header.advance(n);
            }
        }
    }

    let header_bytes = *header.as_array();
    match TlsRecordHeader::parse(&header_bytes) {
        Some(h) => HeaderPollResult::Complete(h),
        None => HeaderPollResult::Error(Error::new(
            ErrorKind::InvalidData,
            "failed to parse TLS header",
        )),
    }
}

/// Poll to read record body (standalone function to avoid borrow issues)
fn poll_read_body<R: AsyncRead + Unpin>(
    upstream: &mut R,
    cx: &mut Context<'_>,
    buffer: &mut BytesMut,
    target_len: usize,
    scratch: &mut Vec<u8>,
) -> BodyPollResult {
    while buffer.len() < target_len {
        let remaining = target_len - buffer.len();
        let chunk_len = remaining.min(8192);

        if scratch.len() < chunk_len {
            scratch.resize(chunk_len, 0);
        }

        let n = {
            let mut read_buf = ReadBuf::new(&mut scratch[..chunk_len]);
            match Pin::new(&mut *upstream).poll_read(cx, &mut read_buf) {
                Poll::Pending => return BodyPollResult::Pending,
                Poll::Ready(Err(e)) => return BodyPollResult::Error(e),
                Poll::Ready(Ok(())) => read_buf.filled().len(),
            }
        };

        if n == 0 {
            return BodyPollResult::Error(Error::new(
                ErrorKind::UnexpectedEof,
                format!(
                    "unexpected EOF in TLS body (got {} of {} bytes)",
                    buffer.len(),
                    target_len
                ),
            ));
        }
        buffer.extend_from_slice(&scratch[..n]);
    }

    BodyPollResult::Complete(buffer.split().freeze())
}

impl<R: AsyncRead + Unpin> FakeTlsReader<R> {
    /// Read exactly n bytes through TLS layer.
    ///
    /// This accumulates data across multiple TLS ApplicationData records.
    pub async fn read_exact(&mut self, n: usize) -> Result<Bytes> {
        if self.is_poisoned() {
            return Err(self.take_poison_error());
        }

        let mut result = BytesMut::with_capacity(n);

        while result.len() < n {
            let mut buf = vec![0u8; n - result.len()];
            let read = AsyncReadExt::read(self, &mut buf).await?;

            if read == 0 {
                return Err(Error::new(
                    ErrorKind::UnexpectedEof,
                    format!("expected {} bytes, got {}", n, result.len()),
                ));
            }

            result.extend_from_slice(&buf[..read]);
        }

        Ok(result.freeze())
    }
}

// ============= FakeTlsWriter State =============

#[derive(Debug)]
enum TlsWriterState {
    /// Ready to accept new data
    Idle,

    /// Writing a complete TLS record (header + body), possibly partially
    WritingRecord {
        position: usize,
        payload_size: usize,
    },

    /// Stream encountered an error and cannot be used
    Poisoned { error: Option<io::Error> },
}

impl StreamState for TlsWriterState {
    fn is_terminal(&self) -> bool {
        matches!(self, Self::Poisoned { .. })
    }

    fn is_poisoned(&self) -> bool {
        matches!(self, Self::Poisoned { .. })
    }

    fn state_name(&self) -> &'static str {
        match self {
            Self::Idle => "Idle",
            Self::WritingRecord { .. } => "WritingRecord",
            Self::Poisoned { .. } => "Poisoned",
        }
    }
}

// ============= FakeTlsWriter =============

/// Writer that wraps bytes into TLS 1.3 Application Data records.
///
/// We chunk outgoing data into records of <= MAX_TLS_CIPHERTEXT_SIZE payload bytes.
/// We do not try to mimic AEAD overhead on the wire; Telegram clients accept it.
/// If you want to be more camouflage-accurate later, you could add optional padding
/// to produce records sized closer to MAX_TLS_CIPHERTEXT_SIZE.
pub struct FakeTlsWriter<W> {
    upstream: W,
    state: TlsWriterState,
    record_buffer: BytesMut,
}

impl<W> FakeTlsWriter<W> {
    pub fn new(upstream: W) -> Self {
        Self {
            upstream,
            state: TlsWriterState::Idle,
            record_buffer: BytesMut::new(),
        }
    }

    pub fn get_ref(&self) -> &W {
        &self.upstream
    }
    pub fn get_mut(&mut self) -> &mut W {
        &mut self.upstream
    }
    pub fn into_inner(self) -> W {
        self.upstream
    }

    pub fn is_poisoned(&self) -> bool {
        self.state.is_poisoned()
    }
    pub fn state_name(&self) -> &'static str {
        self.state.state_name()
    }

    pub fn has_pending(&self) -> bool {
        matches!(
            &self.state,
            TlsWriterState::WritingRecord { position, .. }
                if *position < self.record_buffer.len()
        )
    }

    fn poison(&mut self, error: io::Error) {
        self.record_buffer = BytesMut::new();
        self.state = TlsWriterState::Poisoned { error: Some(error) };
    }

    fn take_poison_error(&mut self) -> io::Error {
        match &mut self.state {
            TlsWriterState::Poisoned { error } => error
                .take()
                .unwrap_or_else(|| io::Error::other("stream previously poisoned")),
            _ => io::Error::other("stream not poisoned"),
        }
    }

    fn build_record(&mut self, data: &[u8]) {
        let header = TlsRecordHeader {
            record_type: TLS_RECORD_APPLICATION,
            version: TLS_VERSION,
            length: data.len() as u16,
        };

        self.record_buffer.clear();
        self.record_buffer.reserve(TLS_HEADER_SIZE + data.len());
        self.record_buffer.extend_from_slice(&header.to_bytes());
        self.record_buffer.extend_from_slice(data);
        debug_assert!(self.record_buffer.len() <= MAX_PENDING_WRITE);
    }

    fn recycle_record_buffer(&mut self) {
        if self.record_buffer.capacity() > MAX_RETAINED_RECORD_CAPACITY {
            self.record_buffer = BytesMut::new();
        } else {
            self.record_buffer.clear();
        }
    }
}

enum FlushResult {
    Complete(usize),
    Pending,
    Error(io::Error),
}

impl<W: AsyncWrite + Unpin> FakeTlsWriter<W> {
    fn poll_flush_record_inner(
        upstream: &mut W,
        cx: &mut Context<'_>,
        record: &[u8],
        position: &mut usize,
    ) -> FlushResult {
        while *position < record.len() {
            let data = &record[*position..];
            match Pin::new(&mut *upstream).poll_write(cx, data) {
                Poll::Pending => return FlushResult::Pending,
                Poll::Ready(Err(e)) => return FlushResult::Error(e),
                Poll::Ready(Ok(0)) => {
                    return FlushResult::Error(Error::new(
                        ErrorKind::WriteZero,
                        "upstream returned 0 bytes written",
                    ));
                }
                Poll::Ready(Ok(n)) => *position += n,
            }
        }

        FlushResult::Complete(0)
    }
}

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

        // Take ownership of state to avoid borrow conflicts.
        let state = std::mem::replace(&mut this.state, TlsWriterState::Idle);

        match state {
            TlsWriterState::Poisoned { error } => {
                this.state = TlsWriterState::Poisoned { error: None };
                let err = error.unwrap_or_else(|| Error::other("stream previously poisoned"));
                return Poll::Ready(Err(err));
            }

            TlsWriterState::WritingRecord {
                mut position,
                payload_size,
            } => {
                // Finish writing previous record before accepting new bytes.
                match Self::poll_flush_record_inner(
                    &mut this.upstream,
                    cx,
                    &this.record_buffer,
                    &mut position,
                ) {
                    FlushResult::Pending => {
                        this.state = TlsWriterState::WritingRecord {
                            position,
                            payload_size,
                        };
                        return Poll::Pending;
                    }
                    FlushResult::Error(e) => {
                        this.poison(Error::new(e.kind(), e.to_string()));
                        return Poll::Ready(Err(e));
                    }
                    FlushResult::Complete(_) => {
                        this.recycle_record_buffer();
                        this.state = TlsWriterState::Idle;
                        // continue to accept new buf below
                    }
                }
            }

            TlsWriterState::Idle => {
                this.state = TlsWriterState::Idle;
            }
        }

        // Now in Idle state
        if buf.is_empty() {
            return Poll::Ready(Ok(0));
        }

        // Chunk to maximum TLS payload size
        let chunk_size = buf.len().min(MAX_TLS_PAYLOAD);
        let chunk = &buf[..chunk_size];

        // Build the complete record (header + payload)
        this.build_record(chunk);

        match Pin::new(&mut this.upstream).poll_write(cx, &this.record_buffer) {
            Poll::Ready(Ok(n)) if n == this.record_buffer.len() => {
                this.recycle_record_buffer();
                Poll::Ready(Ok(chunk_size))
            }

            Poll::Ready(Ok(n)) => {
                this.state = TlsWriterState::WritingRecord {
                    position: n,
                    payload_size: chunk_size,
                };

                // We have accepted chunk_size bytes from caller.
                Poll::Ready(Ok(chunk_size))
            }

            Poll::Ready(Err(e)) => {
                this.poison(Error::new(e.kind(), e.to_string()));
                Poll::Ready(Err(e))
            }

            Poll::Pending => {
                this.state = TlsWriterState::WritingRecord {
                    position: 0,
                    payload_size: chunk_size,
                };

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

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

        let state = std::mem::replace(&mut this.state, TlsWriterState::Idle);

        match state {
            TlsWriterState::Poisoned { error } => {
                this.state = TlsWriterState::Poisoned { error: None };
                let err = error.unwrap_or_else(|| Error::other("stream previously poisoned"));
                return Poll::Ready(Err(err));
            }

            TlsWriterState::WritingRecord {
                mut position,
                payload_size,
            } => match Self::poll_flush_record_inner(
                &mut this.upstream,
                cx,
                &this.record_buffer,
                &mut position,
            ) {
                FlushResult::Pending => {
                    this.state = TlsWriterState::WritingRecord {
                        position,
                        payload_size,
                    };
                    return Poll::Pending;
                }
                FlushResult::Error(e) => {
                    this.poison(Error::new(e.kind(), e.to_string()));
                    return Poll::Ready(Err(e));
                }
                FlushResult::Complete(_) => {
                    this.recycle_record_buffer();
                    this.state = TlsWriterState::Idle;
                }
            },

            TlsWriterState::Idle => {
                this.state = TlsWriterState::Idle;
            }
        }

        Pin::new(&mut this.upstream).poll_flush(cx)
    }

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

        let state = std::mem::replace(&mut this.state, TlsWriterState::Idle);

        match state {
            TlsWriterState::WritingRecord {
                mut position,
                payload_size: _,
            } => {
                // Best-effort flush (do not block shutdown forever).
                let _ = Self::poll_flush_record_inner(
                    &mut this.upstream,
                    cx,
                    &this.record_buffer,
                    &mut position,
                );
                this.recycle_record_buffer();
                this.state = TlsWriterState::Idle;
            }
            _ => {
                this.state = TlsWriterState::Idle;
            }
        }

        Pin::new(&mut this.upstream).poll_shutdown(cx)
    }
}

impl<W: AsyncWrite + Unpin> FakeTlsWriter<W> {
    /// Write all data wrapped in TLS records.
    ///
    /// Convenience method that chunks into <= MAX_TLS_CIPHERTEXT_SIZE records.
    pub async fn write_all_tls(&mut self, data: &[u8]) -> Result<()> {
        let mut written = 0;
        while written < data.len() {
            let chunk_size = (data.len() - written).min(MAX_TLS_PAYLOAD);
            let chunk = &data[written..written + chunk_size];

            AsyncWriteExt::write_all(self, chunk).await?;
            written += chunk_size;
        }

        self.flush().await
    }
}

#[cfg(test)]
#[path = "tls_stream_size_adversarial_tests.rs"]
mod size_adversarial_tests;

// ============= Tests =============

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

    // ============= Test Helpers =============

    /// Build a valid TLS Application Data record
    fn build_tls_record(data: &[u8]) -> Vec<u8> {
        let mut record = vec![
            TLS_RECORD_APPLICATION,
            TLS_VERSION[0],
            TLS_VERSION[1],
            (data.len() >> 8) as u8,
            data.len() as u8,
        ];
        record.extend_from_slice(data);
        record
    }

    /// Build a Change Cipher Spec record
    fn build_ccs_record() -> Vec<u8> {
        vec![
            TLS_RECORD_CHANGE_CIPHER,
            TLS_VERSION[0],
            TLS_VERSION[1],
            0x00,
            0x01, // length = 1
            0x01, // CCS byte
        ]
    }

    /// Mock reader that returns data in chunks
    struct ChunkedReader {
        data: VecDeque<u8>,
        chunk_size: usize,
    }

    impl ChunkedReader {
        fn new(data: &[u8], chunk_size: usize) -> Self {
            Self {
                data: data.iter().copied().collect(),
                chunk_size,
            }
        }
    }

    impl AsyncRead for ChunkedReader {
        fn poll_read(
            mut self: Pin<&mut Self>,
            _cx: &mut Context<'_>,
            buf: &mut ReadBuf<'_>,
        ) -> Poll<Result<()>> {
            if self.data.is_empty() {
                return Poll::Ready(Ok(()));
            }

            let to_read = self.chunk_size.min(self.data.len()).min(buf.remaining());
            for _ in 0..to_read {
                if let Some(byte) = self.data.pop_front() {
                    buf.put_slice(&[byte]);
                }
            }

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

    // ============= FakeTlsReader Tests =============

    #[tokio::test]
    async fn test_tls_reader_single_record() {
        let payload = b"Hello, TLS!";
        let record = build_tls_record(payload);

        let reader = ChunkedReader::new(&record, 100);
        let mut tls_reader = FakeTlsReader::new(reader);

        let buf = tls_reader.read_exact(payload.len()).await.unwrap();
        assert_eq!(&buf[..], payload);
    }

    #[tokio::test]
    async fn test_tls_reader_multiple_records() {
        let payload1 = b"First record";
        let payload2 = b"Second record";

        let mut data = build_tls_record(payload1);
        data.extend_from_slice(&build_tls_record(payload2));

        let reader = ChunkedReader::new(&data, 100);
        let mut tls_reader = FakeTlsReader::new(reader);

        let buf1 = tls_reader.read_exact(payload1.len()).await.unwrap();
        assert_eq!(&buf1[..], payload1);

        let buf2 = tls_reader.read_exact(payload2.len()).await.unwrap();
        assert_eq!(&buf2[..], payload2);
    }

    #[tokio::test]
    async fn test_tls_reader_partial_header() {
        // Read header byte by byte
        let payload = b"Test";
        let record = build_tls_record(payload);

        let reader = ChunkedReader::new(&record, 1); // 1 byte at a time!
        let mut tls_reader = FakeTlsReader::new(reader);

        let buf = tls_reader.read_exact(payload.len()).await.unwrap();

        assert_eq!(&buf[..], payload);
    }

    #[tokio::test]
    async fn test_tls_reader_partial_body() {
        let payload = b"This is a longer payload that will be read in parts";
        let record = build_tls_record(payload);

        let reader = ChunkedReader::new(&record, 7); // Awkward chunk size
        let mut tls_reader = FakeTlsReader::new(reader);

        let buf = tls_reader.read_exact(payload.len()).await.unwrap();

        assert_eq!(&buf[..], payload);
    }

    #[tokio::test]
    async fn test_tls_reader_skip_ccs() {
        // CCS record followed by application data
        let mut data = build_ccs_record();
        let payload = b"After CCS";
        data.extend_from_slice(&build_tls_record(payload));

        let reader = ChunkedReader::new(&data, 100);
        let mut tls_reader = FakeTlsReader::new(reader);

        let buf = tls_reader.read_exact(payload.len()).await.unwrap();

        assert_eq!(&buf[..], payload);
    }

    #[tokio::test]
    async fn test_tls_reader_multiple_ccs() {
        // Multiple CCS records
        let mut data = build_ccs_record();
        data.extend_from_slice(&build_ccs_record());
        let payload = b"After multiple CCS";
        data.extend_from_slice(&build_tls_record(payload));

        let reader = ChunkedReader::new(&data, 3); // Small chunks
        let mut tls_reader = FakeTlsReader::new(reader);

        let buf = tls_reader.read_exact(payload.len()).await.unwrap();

        assert_eq!(&buf[..], payload);
    }

    #[tokio::test]
    async fn test_tls_reader_eof() {
        let reader = ChunkedReader::new(&[], 100);
        let mut tls_reader = FakeTlsReader::new(reader);

        let mut buf = vec![0u8; 10];
        let read = tls_reader.read(&mut buf).await.unwrap();

        assert_eq!(read, 0);
    }

    #[tokio::test]
    async fn test_tls_reader_state_names() {
        let reader = ChunkedReader::new(&[], 100);
        let tls_reader = FakeTlsReader::new(reader);

        assert_eq!(tls_reader.state_name(), "Idle");
        assert!(!tls_reader.is_poisoned());
    }

    // ============= FakeTlsWriter Tests =============

    #[tokio::test]
    async fn test_tls_writer_single_write() {
        let (client, mut server) = duplex(4096);
        let mut writer = FakeTlsWriter::new(client);

        let payload = b"Hello, TLS!";
        writer.write_all(payload).await.unwrap();
        writer.flush().await.unwrap();

        // Read the TLS record from server
        let mut header = [0u8; 5];
        server.read_exact(&mut header).await.unwrap();

        assert_eq!(header[0], TLS_RECORD_APPLICATION);
        assert_eq!(&header[1..3], &TLS_VERSION);

        let length = u16::from_be_bytes([header[3], header[4]]) as usize;
        assert_eq!(length, payload.len());

        let mut body = vec![0u8; length];
        server.read_exact(&mut body).await.unwrap();
        assert_eq!(&body, payload);
    }

    #[tokio::test]
    async fn test_tls_writer_large_data_chunking() {
        let (client, mut server) = duplex(65536);
        let mut writer = FakeTlsWriter::new(client);

        // Write data larger than MAX_TLS_PAYLOAD
        let payload: Vec<u8> = (0..20000).map(|i| (i % 256) as u8).collect();
        writer.write_all(&payload).await.unwrap();
        writer.flush().await.unwrap();

        // Read back - should be multiple records
        let mut received = Vec::new();
        let mut records_count = 0;

        while received.len() < payload.len() {
            let mut header = [0u8; 5];
            if server.read_exact(&mut header).await.is_err() {
                break;
            }

            assert_eq!(header[0], TLS_RECORD_APPLICATION);
            records_count += 1;

            let length = u16::from_be_bytes([header[3], header[4]]) as usize;
            assert!(length <= MAX_TLS_PAYLOAD);

            let mut body = vec![0u8; length];
            server.read_exact(&mut body).await.unwrap();
            received.extend_from_slice(&body);
        }

        assert_eq!(received, payload);
        assert!(records_count >= 2); // Should have multiple records
    }

    #[tokio::test]
    async fn test_tls_stream_roundtrip() {
        let (client, server) = duplex(4096);

        let mut writer = FakeTlsWriter::new(client);
        let mut reader = FakeTlsReader::new(server);

        let original = b"Hello, fake TLS!";
        writer.write_all_tls(original).await.unwrap();
        writer.flush().await.unwrap();

        let received = reader.read_exact(original.len()).await.unwrap();
        assert_eq!(&received[..], original);
    }

    #[tokio::test]
    async fn test_tls_stream_roundtrip_large() {
        let (client, server) = duplex(4096);

        let mut writer = FakeTlsWriter::new(client);
        let mut reader = FakeTlsReader::new(server);

        let original: Vec<u8> = (0..50000).map(|i| (i % 256) as u8).collect();

        // Write in background
        let write_data = original.clone();
        let write_handle = tokio::spawn(async move {
            writer.write_all_tls(&write_data).await.unwrap();
            writer.shutdown().await.unwrap();
        });

        // Read
        let mut received = Vec::new();
        let mut buf = vec![0u8; 1024];
        loop {
            let n = reader.read(&mut buf).await.unwrap();
            if n == 0 {
                break;
            }
            received.extend_from_slice(&buf[..n]);
        }

        write_handle.await.unwrap();
        assert_eq!(received, original);
    }

    #[tokio::test]
    async fn test_tls_writer_state_names() {
        let (client, _server) = duplex(4096);
        let writer = FakeTlsWriter::new(client);

        assert_eq!(writer.state_name(), "Idle");
        assert!(!writer.is_poisoned());
        assert!(!writer.has_pending());
    }

    // ============= Error Handling Tests =============

    #[tokio::test]
    async fn test_tls_reader_invalid_version() {
        let invalid_record = vec![
            TLS_RECORD_APPLICATION,
            0x04,
            0x00, // Invalid version
            0x00,
            0x05, // length = 5
            0x01,
            0x02,
            0x03,
            0x04,
            0x05,
        ];

        let reader = ChunkedReader::new(&invalid_record, 100);
        let mut tls_reader = FakeTlsReader::new(reader);

        let mut buf = vec![0u8; 5];
        let result = tls_reader.read(&mut buf).await;

        assert!(result.is_err());
        assert!(tls_reader.is_poisoned());
    }

    #[tokio::test]
    async fn test_tls_reader_unexpected_eof_header() {
        // Partial header
        let partial = vec![TLS_RECORD_APPLICATION, 0x03];

        let reader = ChunkedReader::new(&partial, 100);
        let mut tls_reader = FakeTlsReader::new(reader);

        let mut buf = vec![0u8; 10];
        let result = tls_reader.read(&mut buf).await;

        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_tls_reader_unexpected_eof_body() {
        // Valid header but truncated body
        let mut record = vec![
            TLS_RECORD_APPLICATION,
            TLS_VERSION[0],
            TLS_VERSION[1],
            0x00,
            0x10, // length = 16
        ];
        record.extend_from_slice(&[0u8; 8]); // Only 8 bytes of body

        let reader = ChunkedReader::new(&record, 100);
        let mut tls_reader = FakeTlsReader::new(reader);

        let mut buf = vec![0u8; 16];
        let result = tls_reader.read(&mut buf).await;

        assert!(result.is_err());
    }

    // ============= Header Parsing Tests =============

    #[test]
    fn test_tls_record_header_parse() {
        let header = [0x17, 0x03, 0x03, 0x01, 0x00];
        let parsed = TlsRecordHeader::parse(&header).unwrap();

        assert_eq!(parsed.record_type, TLS_RECORD_APPLICATION);
        assert_eq!(parsed.version, TLS_VERSION);
        assert_eq!(parsed.length, 256);
    }

    #[test]
    fn test_tls_record_header_validate() {
        let valid = TlsRecordHeader {
            record_type: TLS_RECORD_APPLICATION,
            version: TLS_VERSION,
            length: 100,
        };
        assert!(valid.validate().is_ok());

        let invalid_version = TlsRecordHeader {
            record_type: TLS_RECORD_APPLICATION,
            version: [0x04, 0x00],
            length: 100,
        };
        assert!(invalid_version.validate().is_err());

        let too_large = TlsRecordHeader {
            record_type: TLS_RECORD_APPLICATION,
            version: TLS_VERSION,
            length: 20000,
        };
        assert!(too_large.validate().is_err());
    }

    #[test]
    fn test_tls_record_header_to_bytes() {
        let header = TlsRecordHeader {
            record_type: TLS_RECORD_APPLICATION,
            version: TLS_VERSION,
            length: 0x1234,
        };

        let bytes = header.to_bytes();
        assert_eq!(bytes, [0x17, 0x03, 0x03, 0x12, 0x34]);
    }
}

#[cfg(test)]
#[path = "tls_stream_pending_plaintext_security_tests.rs"]
mod pending_plaintext_security_tests;
