//! 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 16384 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
//! - Proper handling of all TLS record types
//!
//! Important nuance (Telegram FakeTLS):
//! - The TLS spec limits "plaintext fragments" to 2^14 (16384) bytes.
//! - However, the on-the-wire record length can exceed 16384 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 16384 + 24 bytes. We accept that as MAX_TLS_CHUNK_SIZE.
//!
//! If you reject those (e.g. validate length <= 16384), 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, AsyncWrite, AsyncReadExt, AsyncWriteExt, ReadBuf};

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

// ============= 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 ciphertexts often add ~16-24 bytes AEAD overhead, so to mimic
/// on-the-wire record sizes we allow up to 16384 + 24 bytes of plaintext.
const MAX_TLS_PAYLOAD: usize = 16384 + 24;

/// 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;

// ============= 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_CHUNK_SIZE (16384 + 24).
    /// - For other record types we keep stricter bounds to avoid memory abuse.
    fn validate(&self) -> Result<()> {
        // Version: accept TLS 1.0 header (ClientHello quirk) and TLS_VERSION (0x0303).
        if self.version != [0x03, 0x01] && self.version != TLS_VERSION {
            return Err(Error::new(
                ErrorKind::InvalidData,
                format!("Invalid TLS version: {:02x?}", self.version),
            ));
        }

        let len = self.length as usize;

        // Length checks depend on record type.
        // Telegram FakeTLS: ApplicationData length may be 16384 + 24.
        match self.record_type {
            TLS_RECORD_APPLICATION => {
                if len > MAX_TLS_CHUNK_SIZE {
                    return Err(Error::new(
                        ErrorKind::InvalidData,
                        format!("TLS record too large: {} bytes (max {})", len, MAX_TLS_CHUNK_SIZE),
                    ));
                }
            }

            // ChangeCipherSpec/Alert/Handshake should never be that large for our usage
            // (post-handshake we don't expect Handshake at all).
            // Keep strict to reduce attack surface.
            _ => {
                if len > MAX_TLS_PAYLOAD {
                    return Err(Error::new(
                        ErrorKind::InvalidData,
                        format!("TLS control record too large: {} bytes (max {})", len, MAX_TLS_PAYLOAD),
                    ));
                }
            }
        }

        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,
}

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

    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 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::new(ErrorKind::Other, "stream previously poisoned")
            }),
            _ => io::Error::new(ErrorKind::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::new(ErrorKind::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);

                    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,
) -> BodyPollResult {
    // NOTE: This implementation uses a temporary Vec to avoid tricky borrow/lifetime
    // issues with BytesMut spare capacity and ReadBuf across polls.
    // It's safe and correct; optimization is possible if needed.
    while buffer.len() < target_len {
        let remaining = target_len - buffer.len();

        let mut temp = vec![0u8; remaining.min(8192)];
        let mut read_buf = ReadBuf::new(&mut temp);

        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(())) => {
                let n = 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(&temp[..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 {
        record: WriteBuffer,
        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 <= 16384 payload bytes (MAX_TLS_PAYLOAD).
/// 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_CHUNK_SIZE.
pub struct FakeTlsWriter<W> {
    upstream: W,
    state: TlsWriterState,
}

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

    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 { record, .. } if !record.is_empty())
    }

    fn poison(&mut self, error: io::Error) {
        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::new(ErrorKind::Other, "stream previously poisoned")
            }),
            _ => io::Error::new(ErrorKind::Other, "stream not poisoned"),
        }
    }

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

        let mut record = BytesMut::with_capacity(TLS_HEADER_SIZE + data.len());
        record.extend_from_slice(&header.to_bytes());
        record.extend_from_slice(data);
        record
    }
}

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: &mut WriteBuffer,
    ) -> FlushResult {
        while !record.is_empty() {
            let data = record.pending();
            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)) => record.advance(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::new(ErrorKind::Other, "stream previously poisoned")
                });
                return Poll::Ready(Err(err));
            }

            TlsWriterState::WritingRecord { mut record, payload_size } => {
                // Finish writing previous record before accepting new bytes.
                match Self::poll_flush_record_inner(&mut this.upstream, cx, &mut record) {
                    FlushResult::Pending => {
                        this.state = TlsWriterState::WritingRecord { record, 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.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)
        let record_data = Self::build_record(chunk);

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

            Poll::Ready(Ok(n)) => {
                // Partial write of the record: store remainder.
                let mut write_buffer = WriteBuffer::with_max_size(MAX_PENDING_WRITE);
                // record_data length is <= 16389, fits MAX_PENDING_WRITE
                let _ = write_buffer.extend(&record_data[n..]);

                this.state = TlsWriterState::WritingRecord {
                    record: write_buffer,
                    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 => {
                // Buffer entire record and report success for this chunk.
                let mut write_buffer = WriteBuffer::with_max_size(MAX_PENDING_WRITE);
                let _ = write_buffer.extend(&record_data);

                this.state = TlsWriterState::WritingRecord {
                    record: write_buffer,
                    payload_size: chunk_size,
                };

                // Wake to retry flushing soon.
                cx.waker().wake_by_ref();

                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::new(ErrorKind::Other, "stream previously poisoned")
                });
                return Poll::Ready(Err(err));
            }

            TlsWriterState::WritingRecord { mut record, payload_size } => {
                match Self::poll_flush_record_inner(&mut this.upstream, cx, &mut record) {
                    FlushResult::Pending => {
                        this.state = TlsWriterState::WritingRecord { record, 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.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 record, payload_size: _ } => {
                // Best-effort flush (do not block shutdown forever).
                let _ = Self::poll_flush_record_inner(&mut this.upstream, cx, &mut record);
                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 <= 16384 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
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::VecDeque;
    use tokio::io::{duplex, AsyncReadExt, AsyncWriteExt};
    
    // ============= 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]);
    }
}