// Quinn-compatible API wrapper around quiche
//
// This module provides a drop-in replacement for Quinn using quiche underneath.
// Uses channel-based architecture for async/await support while maintaining
// single-threaded access to quiche::Connection for maximum performance.

pub mod crypto {
    pub mod rustls {
        use std::sync::Arc;

        pub struct QuicServerConfig {
            pub(super) rustls_config: Arc<tokio_rustls::rustls::ServerConfig>,
            // Store cert/key for quiche - will be populated by rustls_util
            pub(super) cert_pem: Option<Vec<u8>>,
            pub(super) key_pem: Option<Vec<u8>>,
        }

        impl QuicServerConfig {
            pub fn with_cert_key(
                rustls_config: Arc<tokio_rustls::rustls::ServerConfig>,
                cert_pem: Vec<u8>,
                key_pem: Vec<u8>,
            ) -> Self {
                QuicServerConfig {
                    rustls_config,
                    cert_pem: Some(cert_pem),
                    key_pem: Some(key_pem),
                }
            }
        }

        impl TryFrom<Arc<tokio_rustls::rustls::ServerConfig>> for QuicServerConfig {
            type Error = String;

            fn try_from(config: Arc<tokio_rustls::rustls::ServerConfig>) -> Result<Self, Self::Error> {
                // Fallback for compatibility - won't work without cert/key
                Ok(QuicServerConfig {
                    rustls_config: config,
                    cert_pem: None,
                    key_pem: None,
                })
            }
        }
    }
}

use bytes::{Bytes, BytesMut};
use std::collections::{HashMap, VecDeque};
use std::io;
use std::net::SocketAddr;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Arc;
use std::task::{Context, Poll};
use std::time::{Duration, Instant};
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
use tokio::net::UdpSocket;
use tokio::sync::{mpsc, oneshot, Mutex};
use tokio::time;

const MAX_DATAGRAM_SIZE: usize = 1350;
const RECV_BUFFER_SIZE: usize = 65535;

// Re-export types for compatibility
pub use crypto::rustls::QuicServerConfig;

/// Configuration for the endpoint
pub struct EndpointConfig {
    _private: (),
}

impl EndpointConfig {
    pub fn default() -> Self {
        Self { _private: () }
    }
}

/// Server configuration
pub struct ServerConfig {
    crypto: Arc<QuicServerConfig>,
    pub transport: Arc<TransportConfig>,
}

impl ServerConfig {
    pub fn with_crypto(crypto: Arc<QuicServerConfig>) -> Self {
        Self {
            crypto,
            transport: Arc::new(TransportConfig::default()),
        }
    }
}

/// Transport configuration
pub struct TransportConfig {
    max_concurrent_bidi_streams: AtomicUsize,
    max_concurrent_uni_streams: AtomicUsize,
    max_idle_timeout: Mutex<Option<Duration>>,
    keep_alive_interval: Mutex<Option<Duration>>,
    send_window: AtomicUsize,
    receive_window: AtomicUsize,
    stream_receive_window: AtomicUsize,
}

impl Default for TransportConfig {
    fn default() -> Self {
        Self {
            max_concurrent_bidi_streams: AtomicUsize::new(100),
            max_concurrent_uni_streams: AtomicUsize::new(100),
            max_idle_timeout: Mutex::new(Some(Duration::from_secs(30))),
            keep_alive_interval: Mutex::new(None),
            send_window: AtomicUsize::new(8 * 1024 * 1024),
            receive_window: AtomicUsize::new(8 * 1024 * 1024),
            stream_receive_window: AtomicUsize::new(1 * 1024 * 1024),
        }
    }
}

impl TransportConfig {
    pub fn max_concurrent_bidi_streams(&mut self, value: u64) -> &mut Self {
        self.max_concurrent_bidi_streams.store(value as usize, Ordering::Relaxed);
        self
    }

    pub fn max_concurrent_uni_streams(&mut self, value: u64) -> &mut Self {
        self.max_concurrent_uni_streams.store(value as usize, Ordering::Relaxed);
        self
    }

    pub fn max_idle_timeout(&mut self, value: Option<Duration>) -> &mut Self {
        *self.max_idle_timeout.blocking_lock() = value;
        self
    }

    pub fn keep_alive_interval(&mut self, value: Option<Duration>) -> &mut Self {
        *self.keep_alive_interval.blocking_lock() = value;
        self
    }

    pub fn send_window(&mut self, value: u64) -> &mut Self {
        self.send_window.store(value as usize, Ordering::Relaxed);
        self
    }

    pub fn receive_window(&mut self, value: u64) -> &mut Self {
        self.receive_window.store(value as usize, Ordering::Relaxed);
        self
    }

    pub fn stream_receive_window(&mut self, value: u64) -> &mut Self {
        self.stream_receive_window.store(value as usize, Ordering::Relaxed);
        self
    }
}

/// Tokio runtime marker
pub struct TokioRuntime;

/// QUIC endpoint that can accept connections
pub struct Endpoint {
    _local_addr: SocketAddr,
    incoming_rx: Mutex<mpsc::UnboundedReceiver<Incoming>>,
}

impl Endpoint {
    pub fn new(
        _endpoint_config: EndpointConfig,
        server_config: Option<ServerConfig>,
        socket: std::net::UdpSocket,
        _runtime: Arc<TokioRuntime>,
    ) -> io::Result<Self> {
        let server_config = server_config.ok_or_else(|| {
            io::Error::other("server config required for server endpoint")
        })?;

        socket.set_nonblocking(true)?;
        let local_addr = socket.local_addr()?;
        let socket = Arc::new(UdpSocket::from_std(socket)?);

        let (incoming_tx, incoming_rx) = mpsc::unbounded_channel();

        // Spawn endpoint task
        tokio::spawn(run_endpoint(socket, server_config, incoming_tx));

        Ok(Endpoint {
            _local_addr: local_addr,
            incoming_rx: Mutex::new(incoming_rx),
        })
    }

    pub async fn accept(&self) -> Option<Incoming> {
        self.incoming_rx.lock().await.recv().await
    }
}

/// Incoming connection
pub struct Incoming {
    conn_rx: oneshot::Receiver<Result<Connection, io::Error>>,
}

impl Incoming {
    pub async fn await(self) -> io::Result<Connection> {
        self.conn_rx.await
            .map_err(|_| io::Error::other("connection setup failed"))?
    }
}

enum ConnectionCmd {
    AcceptBi {
        resp: oneshot::Sender<Result<(SendStream, RecvStream), ConnectionError>>,
    },
    AcceptUni {
        resp: oneshot::Sender<Result<RecvStream, ConnectionError>>,
    },
    ReadDatagram {
        resp: oneshot::Sender<Result<Bytes, ConnectionError>>,
    },
    SendDatagram {
        data: Bytes,
        resp: oneshot::Sender<Result<(), io::Error>>,
    },
    StreamSend {
        stream_id: u64,
        data: Bytes,
        fin: bool,
        resp: oneshot::Sender<Result<usize, io::Error>>,
    },
    StreamRecv {
        stream_id: u64,
        buf_size: usize,
        resp: oneshot::Sender<Result<Option<(Bytes, bool)>, io::Error>>,
    },
    StreamShutdown {
        stream_id: u64,
        direction: quiche::Shutdown,
        error_code: u64,
        resp: oneshot::Sender<Result<(), io::Error>>,
    },
    ExportKeyingMaterial {
        output_len: usize,
        label: Vec<u8>,
        context: Vec<u8>,
        resp: oneshot::Sender<Result<Vec<u8>, io::Error>>,
    },
    QueryMaxDatagramSize {
        resp: oneshot::Sender<Option<usize>>,
    },
}

/// Established QUIC connection
#[derive(Clone)]
pub struct Connection {
    cmd_tx: mpsc::UnboundedSender<ConnectionCmd>,
    state: Arc<ConnectionState>,
}

struct ConnectionState {
    _peer_addr: SocketAddr,
    is_closed: AtomicBool,
}

impl Connection {
    pub async fn accept_bi(&self) -> Result<(SendStream, RecvStream), ConnectionError> {
        let (tx, rx) = oneshot::channel();
        self.cmd_tx.send(ConnectionCmd::AcceptBi { resp: tx })
            .map_err(|_| ConnectionError::LocallyClosed)?;
        rx.await.map_err(|_| ConnectionError::LocallyClosed)?
    }

    pub async fn accept_uni(&self) -> Result<RecvStream, ConnectionError> {
        let (tx, rx) = oneshot::channel();
        self.cmd_tx.send(ConnectionCmd::AcceptUni { resp: tx })
            .map_err(|_| ConnectionError::LocallyClosed)?;
        rx.await.map_err(|_| ConnectionError::LocallyClosed)?
    }

    pub async fn read_datagram(&self) -> Result<Bytes, ConnectionError> {
        let (tx, rx) = oneshot::channel();
        self.cmd_tx.send(ConnectionCmd::ReadDatagram { resp: tx })
            .map_err(|_| ConnectionError::LocallyClosed)?;
        rx.await.map_err(|_| ConnectionError::LocallyClosed)?
    }

    pub fn send_datagram(&self, data: Bytes) -> Result<(), io::Error> {
        let (tx, rx) = oneshot::channel();
        self.cmd_tx.send(ConnectionCmd::SendDatagram { data, resp: tx })
            .map_err(|_| io::Error::other("connection closed"))?;

        // Block on response (datagram send is immediate in quiche)
        rx.blocking_recv()
            .map_err(|_| io::Error::other("connection closed"))?
    }

    pub fn max_datagram_size(&self) -> Option<usize> {
        let (tx, rx) = oneshot::channel();
        self.cmd_tx.send(ConnectionCmd::QueryMaxDatagramSize { resp: tx }).ok()?;
        rx.blocking_recv().ok()?
    }

    pub fn export_keying_material(
        &self,
        output: &mut [u8],
        label: &[u8],
        context: &[u8],
    ) -> Result<(), io::Error> {
        let (tx, rx) = oneshot::channel();
        self.cmd_tx.send(ConnectionCmd::ExportKeyingMaterial {
            output_len: output.len(),
            label: label.to_vec(),
            context: context.to_vec(),
            resp: tx,
        }).map_err(|_| io::Error::other("connection closed"))?;

        let result = rx.blocking_recv()
            .map_err(|_| io::Error::other("connection closed"))??;

        output.copy_from_slice(&result);
        Ok(())
    }
}

/// Send half of a stream
pub struct SendStream {
    cmd_tx: mpsc::UnboundedSender<ConnectionCmd>,
    stream_id: u64,
}

impl SendStream {
    pub async fn write(&mut self, buf: &[u8]) -> Result<usize, io::Error> {
        let (tx, rx) = oneshot::channel();
        self.cmd_tx.send(ConnectionCmd::StreamSend {
            stream_id: self.stream_id,
            data: Bytes::copy_from_slice(buf),
            fin: false,
            resp: tx,
        }).map_err(|_| io::Error::other("connection closed"))?;

        rx.await.map_err(|_| io::Error::other("connection closed"))?
    }

    pub async fn finish(&mut self) -> Result<(), io::Error> {
        let (tx, rx) = oneshot::channel();
        self.cmd_tx.send(ConnectionCmd::StreamSend {
            stream_id: self.stream_id,
            data: Bytes::new(),
            fin: true,
            resp: tx,
        }).map_err(|_| io::Error::other("connection closed"))?;

        rx.await.map_err(|_| io::Error::other("connection closed"))?;
        Ok(())
    }

    pub async fn shutdown(&mut self) -> Result<(), io::Error> {
        let (tx, rx) = oneshot::channel();
        self.cmd_tx.send(ConnectionCmd::StreamShutdown {
            stream_id: self.stream_id,
            direction: quiche::Shutdown::Write,
            error_code: 0,
            resp: tx,
        }).map_err(|_| io::Error::other("connection closed"))?;

        rx.await.map_err(|_| io::Error::other("connection closed"))?
    }
}

/// Receive half of a stream
pub struct RecvStream {
    cmd_tx: mpsc::UnboundedSender<ConnectionCmd>,
    stream_id: u64,
    buffer: BytesMut,
    finished: bool,
    pending_read: Option<oneshot::Receiver<Result<Option<(Bytes, bool)>, io::Error>>>,
}

impl RecvStream {
    pub async fn read(&mut self, buf: &mut [u8]) -> Result<Option<usize>, io::Error> {
        if self.finished && self.buffer.is_empty() {
            return Ok(None);
        }

        // Return buffered data first
        if !self.buffer.is_empty() {
            let to_copy = std::cmp::min(buf.len(), self.buffer.len());
            buf[..to_copy].copy_from_slice(&self.buffer[..to_copy]);
            self.buffer.advance(to_copy);
            return Ok(Some(to_copy));
        }

        // Request more data
        let (tx, rx) = oneshot::channel();
        self.cmd_tx.send(ConnectionCmd::StreamRecv {
            stream_id: self.stream_id,
            buf_size: buf.len().max(8192),
            resp: tx,
        }).map_err(|_| io::Error::other("connection closed"))?;

        match rx.await.map_err(|_| io::Error::other("connection closed"))?? {
            Some((data, fin)) => {
                self.finished = fin;
                let to_copy = std::cmp::min(buf.len(), data.len());
                buf[..to_copy].copy_from_slice(&data[..to_copy]);

                if to_copy < data.len() {
                    self.buffer.extend_from_slice(&data[to_copy..]);
                }

                Ok(Some(to_copy))
            }
            None => {
                // Would block - no data available yet
                Ok(Some(0))
            }
        }
    }

    pub async fn stop(&mut self, error_code: u64) -> Result<(), io::Error> {
        let (tx, rx) = oneshot::channel();
        self.cmd_tx.send(ConnectionCmd::StreamShutdown {
            stream_id: self.stream_id,
            direction: quiche::Shutdown::Read,
            error_code,
            resp: tx,
        }).map_err(|_| io::Error::other("connection closed"))?;

        rx.await.map_err(|_| io::Error::other("connection closed"))?
    }
}

// Implement AsyncRead for RecvStream
impl AsyncRead for RecvStream {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        // Return buffered data if available
        if !self.buffer.is_empty() {
            let to_copy = std::cmp::min(buf.remaining(), self.buffer.len());
            buf.put_slice(&self.buffer[..to_copy]);
            self.buffer.advance(to_copy);
            return Poll::Ready(Ok(()));
        }

        // Return EOF if finished
        if self.finished {
            return Poll::Ready(Ok(()));
        }

        // Check if we have a pending read
        if let Some(mut rx) = self.pending_read.take() {
            match Pin::new(&mut rx).poll(cx) {
                Poll::Ready(Ok(Ok(Some((data, fin))))) => {
                    self.finished = fin;
                    let to_copy = std::cmp::min(buf.remaining(), data.len());
                    buf.put_slice(&data[..to_copy]);
                    if to_copy < data.len() {
                        self.buffer.extend_from_slice(&data[to_copy..]);
                    }
                    Poll::Ready(Ok(()))
                }
                Poll::Ready(Ok(Ok(None))) => {
                    // Would block - retry
                    Poll::Pending
                }
                Poll::Ready(Ok(Err(e))) => Poll::Ready(Err(e)),
                Poll::Ready(Err(_)) => Poll::Ready(Err(io::Error::other("connection closed"))),
                Poll::Pending => {
                    self.pending_read = Some(rx);
                    Poll::Pending
                }
            }
        } else {
            // Start a new read
            let (tx, rx) = oneshot::channel();
            if self.cmd_tx.send(ConnectionCmd::StreamRecv {
                stream_id: self.stream_id,
                buf_size: buf.remaining().max(8192),
                resp: tx,
            }).is_err() {
                return Poll::Ready(Err(io::Error::other("connection closed")));
            }
            self.pending_read = Some(rx);
            cx.waker().wake_by_ref();
            Poll::Pending
        }
    }
}

// Implement AsyncWrite for SendStream
impl AsyncWrite for SendStream {
    fn poll_write(
        self: Pin<&mut Self>,
        _cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        let this = self.get_mut();
        let (tx, rx) = oneshot::channel();

        if this.cmd_tx.send(ConnectionCmd::StreamSend {
            stream_id: this.stream_id,
            data: Bytes::copy_from_slice(buf),
            fin: false,
            resp: tx,
        }).is_err() {
            return Poll::Ready(Err(io::Error::other("connection closed")));
        }

        // Block on response - quiche stream_send is synchronous
        match rx.blocking_recv() {
            Ok(Ok(written)) => Poll::Ready(Ok(written)),
            Ok(Err(e)) => Poll::Ready(Err(e)),
            Err(_) => Poll::Ready(Err(io::Error::other("connection closed"))),
        }
    }

    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        // No-op for QUIC
        Poll::Ready(Ok(()))
    }

    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        let this = self.get_mut();
        let (tx, rx) = oneshot::channel();

        if this.cmd_tx.send(ConnectionCmd::StreamShutdown {
            stream_id: this.stream_id,
            direction: quiche::Shutdown::Write,
            error_code: 0,
            resp: tx,
        }).is_err() {
            return Poll::Ready(Err(io::Error::other("connection closed")));
        }

        match rx.blocking_recv() {
            Ok(Ok(())) => Poll::Ready(Ok(())),
            Ok(Err(e)) => Poll::Ready(Err(e)),
            Err(_) => Poll::Ready(Err(io::Error::other("connection closed"))),
        }
    }
}

#[derive(Debug)]
pub enum ConnectionError {
    ApplicationClosed(ApplicationClose),
    ConnectionClosed(ConnectionClose),
    TimedOut,
    LocallyClosed,
}

#[derive(Debug)]
pub struct ApplicationClose {
    pub error_code: u64,
    pub reason: Bytes,
}

#[derive(Debug)]
pub struct ConnectionClose {
    pub error_code: u64,
    pub reason: Bytes,
}

impl std::fmt::Display for ConnectionError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ConnectionError::ApplicationClosed(c) => write!(f, "application closed: {}", c.error_code),
            ConnectionError::ConnectionClosed(c) => write!(f, "connection closed: {}", c.error_code),
            ConnectionError::TimedOut => write!(f, "timed out"),
            ConnectionError::LocallyClosed => write!(f, "locally closed"),
        }
    }
}

impl std::error::Error for ConnectionError {}

// Endpoint task - handles all incoming connections
async fn run_endpoint(
    socket: Arc<UdpSocket>,
    config: ServerConfig,
    incoming_tx: mpsc::UnboundedSender<Incoming>,
) {
    let mut quiche_config = create_quiche_config(&config);
    let local_addr = socket.local_addr().unwrap();

    let scid_seed = ring::hmac::Key::generate(
        ring::hmac::HMAC_SHA256,
        &ring::rand::SystemRandom::new(),
    ).unwrap();

    let mut recv_buf = vec![0; RECV_BUFFER_SIZE];
    let mut connections: HashMap<Vec<u8>, mpsc::UnboundedSender<Bytes>> = HashMap::new();

    loop {
        // Receive packet from socket
        let (len, peer_addr) = match socket.recv_from(&mut recv_buf).await {
            Ok(v) => v,
            Err(e) => {
                log::error!("Socket recv error: {}", e);
                break;
            }
        };

        let pkt_buf = Bytes::copy_from_slice(&recv_buf[..len]);

        // Parse header
        let hdr = match quiche::Header::from_slice(&pkt_buf, quiche::MAX_CONN_ID_LEN) {
            Ok(h) => h,
            Err(e) => {
                log::warn!("Failed to parse header: {:?}", e);
                continue;
            }
        };

        // Check if existing connection
        if let Some(conn_tx) = connections.get(hdr.dcid.as_ref()) {
            // Send to existing connection's task
            if conn_tx.send(pkt_buf).is_err() {
                connections.remove(hdr.dcid.as_ref());
            }
            continue;
        }

        // New connection
        if hdr.ty != quiche::Type::Initial {
            log::warn!("Non-Initial packet for unknown connection");
            continue;
        }

        // Create connection ID
        let conn_id = ring::hmac::sign(&scid_seed, hdr.dcid.as_ref());
        let scid = quiche::ConnectionId::from_ref(&conn_id.as_ref()[..quiche::MAX_CONN_ID_LEN]);

        // Create quiche connection
        let conn = match quiche::accept(
            &scid,
            None,
            local_addr,
            peer_addr,
            &mut quiche_config,
        ) {
            Ok(c) => c,
            Err(e) => {
                log::error!("Failed to accept connection: {:?}", e);
                continue;
            }
        };

        // Create channel for this connection's packets
        let (conn_pkt_tx, conn_pkt_rx) = mpsc::unbounded_channel();
        connections.insert(scid.to_vec(), conn_pkt_tx.clone());

        // Send initial packet to connection
        let _ = conn_pkt_tx.send(pkt_buf);

        // Spawn connection task
        let (conn_ready_tx, conn_ready_rx) = oneshot::channel();
        let socket_clone = socket.clone();

        tokio::spawn(async move {
            run_connection(conn, socket_clone, peer_addr, conn_pkt_rx, conn_ready_tx).await;
        });

        // Send to incoming queue
        let _ = incoming_tx.send(Incoming {
            conn_rx: conn_ready_rx,
        });
    }
}

fn create_quiche_config(config: &ServerConfig) -> quiche::Config {
    let transport = &config.transport;
    let mut qconfig = quiche::Config::new(quiche::PROTOCOL_VERSION).unwrap();

    // Load TLS cert and key
    if let (Some(cert_pem), Some(key_pem)) = (&config.crypto.cert_pem, &config.crypto.key_pem) {
        use std::io::Write;

        // Write to temp files (quiche needs file paths)
        let temp_dir = std::env::temp_dir();
        let cert_path = temp_dir.join(format!("quiche_cert_{}.pem", std::process::id()));
        let key_path = temp_dir.join(format!("quiche_key_{}.pem", std::process::id()));

        if let Ok(mut f) = std::fs::File::create(&cert_path) {
            let _ = f.write_all(cert_pem);
        }
        if let Ok(mut f) = std::fs::File::create(&key_path) {
            let _ = f.write_all(key_pem);
        }

        // Load into quiche
        let _ = qconfig.load_cert_chain_from_pem_file(cert_path.to_str().unwrap());
        let _ = qconfig.load_priv_key_from_pem_file(key_path.to_str().unwrap());

        // Clean up temp files
        let _ = std::fs::remove_file(&cert_path);
        let _ = std::fs::remove_file(&key_path);
    }

    // Set ALPN protocols from rustls config
    let alpn_protocols: Vec<Vec<u8>> = config.crypto.rustls_config
        .alpn_protocols
        .iter()
        .map(|p| p.to_vec())
        .collect();
    if !alpn_protocols.is_empty() {
        let alpn_refs: Vec<&[u8]> = alpn_protocols.iter().map(|p| p.as_slice()).collect();
        let _ = qconfig.set_application_protos(&alpn_refs);
    }

    // Configure transport parameters
    let max_idle = transport.max_idle_timeout.blocking_lock();
    if let Some(timeout) = *max_idle {
        qconfig.set_max_idle_timeout(timeout.as_millis() as u64);
    }
    drop(max_idle);

    qconfig.set_max_recv_udp_payload_size(RECV_BUFFER_SIZE);
    qconfig.set_initial_max_data(transport.receive_window.load(Ordering::Relaxed) as u64);
    qconfig.set_initial_max_stream_data_bidi_local(transport.stream_receive_window.load(Ordering::Relaxed) as u64);
    qconfig.set_initial_max_stream_data_bidi_remote(transport.stream_receive_window.load(Ordering::Relaxed) as u64);
    qconfig.set_initial_max_stream_data_uni(transport.stream_receive_window.load(Ordering::Relaxed) as u64);
    qconfig.set_initial_max_streams_bidi(transport.max_concurrent_bidi_streams.load(Ordering::Relaxed) as u64);
    qconfig.set_initial_max_streams_uni(transport.max_concurrent_uni_streams.load(Ordering::Relaxed) as u64);
    qconfig.enable_dgram(true, 1000, 1000);

    qconfig
}

// Connection task - handles a single connection
async fn run_connection(
    mut conn: quiche::Connection,
    socket: Arc<UdpSocket>,
    peer_addr: SocketAddr,
    mut pkt_rx: mpsc::UnboundedReceiver<Bytes>,
    ready_tx: oneshot::Sender<Result<Connection, io::Error>>,
) {
    let local_addr = socket.local_addr().unwrap();
    let (cmd_tx, mut cmd_rx) = mpsc::unbounded_channel();

    let mut send_buf = vec![0; MAX_DATAGRAM_SIZE];

    // Pending operations
    let mut pending_bi: VecDeque<oneshot::Sender<Result<(SendStream, RecvStream), ConnectionError>>> = VecDeque::new();
    let mut pending_uni: VecDeque<oneshot::Sender<Result<RecvStream, ConnectionError>>> = VecDeque::new();
    let mut pending_dgram: VecDeque<oneshot::Sender<Result<Bytes, ConnectionError>>> = VecDeque::new();

    let mut ready_sent = false;
    let mut interval = time::interval(Duration::from_millis(25));

    loop {
        tokio::select! {
            // Handle incoming packets
            Some(pkt) = pkt_rx.recv() => {
                let recv_info = quiche::RecvInfo {
                    from: peer_addr,
                    to: local_addr,
                };

                if let Err(e) = conn.recv(&pkt, recv_info) {
                    log::warn!("Connection recv error: {:?}", e);
                }

                // Send connection handle once established
                if !ready_sent && conn.is_established() {
                    let connection = Connection {
                        cmd_tx: cmd_tx.clone(),
                        state: Arc::new(ConnectionState {
                            _peer_addr: peer_addr,
                            is_closed: AtomicBool::new(false),
                        }),
                    };
                    let _ = ready_tx.send(Ok(connection));
                    ready_sent = true;
                }

                // Flush pending operations
                flush_pending(&mut conn, &cmd_tx, &mut pending_bi, &mut pending_uni, &mut pending_dgram);

                // Send outgoing packets
                send_packets(&mut conn, &socket, &mut send_buf).await;
            }

            // Handle commands
            Some(cmd) = cmd_rx.recv() => {
                handle_command(&mut conn, cmd, &cmd_tx, &mut pending_bi, &mut pending_uni, &mut pending_dgram);
                send_packets(&mut conn, &socket, &mut send_buf).await;
            }

            // Timeout handling
            _ = interval.tick() => {
                conn.on_timeout();
                send_packets(&mut conn, &socket, &mut send_buf).await;
            }
        }

        if conn.is_closed() {
            break;
        }
    }
}

fn flush_pending(
    conn: &mut quiche::Connection,
    cmd_tx: &mpsc::UnboundedSender<ConnectionCmd>,
    pending_bi: &mut VecDeque<oneshot::Sender<Result<(SendStream, RecvStream), ConnectionError>>>,
    pending_uni: &mut VecDeque<oneshot::Sender<Result<RecvStream, ConnectionError>>>,
    pending_dgram: &mut VecDeque<oneshot::Sender<Result<Bytes, ConnectionError>>>,
) {
    // Check for new streams
    for stream_id in conn.readable() {
        let is_bidi = stream_id % 4 < 2;

        if is_bidi {
            if let Some(sender) = pending_bi.pop_front() {
                let send = SendStream {
                    cmd_tx: cmd_tx.clone(),
                    stream_id,
                };
                let recv = RecvStream {
                    cmd_tx: cmd_tx.clone(),
                    stream_id,
                    buffer: BytesMut::new(),
                    finished: false,
                    pending_read: None,
                };
                let _ = sender.send(Ok((send, recv)));
            }
        } else {
            if let Some(sender) = pending_uni.pop_front() {
                let recv = RecvStream {
                    cmd_tx: cmd_tx.clone(),
                    stream_id,
                    buffer: BytesMut::new(),
                    finished: false,
                    pending_read: None,
                };
                let _ = sender.send(Ok(recv));
            }
        }
    }

    // Check for datagrams
    while let Some(sender) = pending_dgram.front() {
        let mut buf = vec![0; 65535];
        match conn.dgram_recv(&mut buf) {
            Ok(len) => {
                buf.truncate(len);
                let sender = pending_dgram.pop_front().unwrap();
                let _ = sender.send(Ok(Bytes::from(buf)));
            }
            Err(quiche::Error::Done) => break,
            Err(_) => {
                pending_dgram.pop_front();
            }
        }
    }
}

fn handle_command(
    conn: &mut quiche::Connection,
    cmd: ConnectionCmd,
    cmd_tx: &mpsc::UnboundedSender<ConnectionCmd>,
    pending_bi: &mut VecDeque<oneshot::Sender<Result<(SendStream, RecvStream), ConnectionError>>>,
    pending_uni: &mut VecDeque<oneshot::Sender<Result<RecvStream, ConnectionError>>>,
    pending_dgram: &mut VecDeque<oneshot::Sender<Result<Bytes, ConnectionError>>>,
) {
    match cmd {
        ConnectionCmd::AcceptBi { resp } => {
            pending_bi.push_back(resp);
        }
        ConnectionCmd::AcceptUni { resp } => {
            pending_uni.push_back(resp);
        }
        ConnectionCmd::ReadDatagram { resp } => {
            pending_dgram.push_back(resp);
        }
        ConnectionCmd::SendDatagram { data, resp } => {
            let result = conn.dgram_send(&data)
                .map_err(|e| io::Error::other(format!("{:?}", e)));
            let _ = resp.send(result);
        }
        ConnectionCmd::StreamSend { stream_id, data, fin, resp } => {
            let result = conn.stream_send(stream_id, &data, fin)
                .map_err(|e| io::Error::other(format!("{:?}", e)));
            let _ = resp.send(result);
        }
        ConnectionCmd::StreamRecv { stream_id, buf_size, resp } => {
            let mut buf = vec![0; buf_size];
            let result = match conn.stream_recv(stream_id, &mut buf) {
                Ok((len, fin)) => {
                    buf.truncate(len);
                    Ok(Some((Bytes::from(buf), fin)))
                }
                Err(quiche::Error::Done) => Ok(None),
                Err(e) => Err(io::Error::other(format!("{:?}", e))),
            };
            let _ = resp.send(result);
        }
        ConnectionCmd::StreamShutdown { stream_id, direction, error_code, resp } => {
            let result = conn.stream_shutdown(stream_id, direction, error_code)
                .map_err(|e| io::Error::other(format!("{:?}", e)));
            let _ = resp.send(result);
        }
        ConnectionCmd::ExportKeyingMaterial { output_len, label, context, resp } => {
            let mut output = vec![0; output_len];
            let result = conn.export_keying_material(&mut output, &label, &context)
                .map(|_| output)
                .map_err(|e| io::Error::other(format!("{:?}", e)));
            let _ = resp.send(result);
        }
        ConnectionCmd::QueryMaxDatagramSize { resp } => {
            let size = conn.dgram_max_writable_len();
            let _ = resp.send(if size > 0 { Some(size) } else { None });
        }
    }
}

async fn send_packets(
    conn: &mut quiche::Connection,
    socket: &UdpSocket,
    send_buf: &mut [u8],
) {
    loop {
        let (size, send_info) = match conn.send(send_buf) {
            Ok(v) => v,
            Err(quiche::Error::Done) => return,
            Err(e) => {
                log::error!("Connection send error: {:?}", e);
                return;
            }
        };

        if let Err(e) = socket.send_to(&send_buf[..size], send_info.to).await {
            log::error!("Socket send error: {}", e);
            return;
        }
    }
}
