use std::borrow::Cow;
use std::collections::HashMap;
use std::io;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;

use avril_caps::Capabilities;
use avril_core::serde_v8::{JsBuffer, ToJsBuffer};
use avril_core::{Extension, Op};
use base64::Engine;
use base64::engine::general_purpose::STANDARD as BASE64;
use bytes::Bytes;
use fastwebsockets::{Frame, OpCode, WebSocketError, WebSocketWrite};
use futures::channel::oneshot;
use http_body_util::Empty;
use hyper::body::Incoming;
use hyper::header::{self, HeaderValue};
use hyper::{Method, Request, StatusCode};
use hyper_util::rt::TokioIo;
use rustls::RootCertStore;
use sha1::{Digest, Sha1};
use tokio::io::{AsyncRead, AsyncWrite, ReadHalf, WriteHalf};
use tokio::net::TcpStream;
use tokio::sync::{Mutex as AsyncMutex, mpsc};
use tokio_rustls::{TlsConnector, rustls::pki_types::ServerName};
use tokio_util::sync::CancellationToken;
use tokio_util::task::AbortOnDropHandle;
use url::{Host, Url};

mod reader;

type Error = Box<dyn std::error::Error + Send + Sync>;
/// Undelivered events; the guest pump drains this promptly, stalling the read
/// loop (and with it the peer) when it falls behind.
const EVENT_QUEUE: usize = 64;
/// Frames read but not yet handled; bounds what a stalled guest can buffer.
const READ_QUEUE: usize = 16;
/// RFC 6455 §7.1.5 recommends waiting a reasonable time for the peer's close frame.
const CLOSE_TIMEOUT: Duration = Duration::from_secs(30);
const TCP_CLOSE_TIMEOUT: Duration = Duration::from_secs(2);
const MAX_ID: u64 = 9_007_199_254_740_991;

trait Duplex: AsyncRead + AsyncWrite {}
impl<T: AsyncRead + AsyncWrite> Duplex for T {}
type Socket = Box<dyn Duplex + Send + Unpin>;
type Negotiated = TokioIo<hyper::upgrade::Upgraded>;
/// Reading runs in its own task: decoding consumes frame bytes across
/// await points, so a dropped read desynchronizes the stream. It is reached
/// through a channel instead, whose receive is cancel-safe.
type Reader = reader::Reader<ReadHalf<Negotiated>>;
type Writer = WebSocketWrite<WriteHalf<Negotiated>>;

/// A received frame, owned so it can cross a task boundary.
struct Received {
    opcode: OpCode,
    payload: Vec<u8>,
}

enum FrameKind {
    Text,
    Binary,
}

enum Command {
    Send(FrameKind, Vec<u8>, oneshot::Sender<bool>),
    Close(Option<Vec<u8>>),
}

struct Resource {
    /// Aborts a connection that was not yet established, per close() during CONNECTING.
    canceled: CancellationToken,
    commands: mpsc::UnboundedSender<Command>,
    events: AsyncMutex<mpsc::Receiver<Event>>,
}

/// One event per pump turn, shaped for `serde_v8` as a plain object.
#[derive(serde::Serialize)]
#[serde(tag = "kind", rename_all = "lowercase")]
enum Event {
    Open {
        protocol: String,
    },
    Text {
        data: String,
    },
    Binary {
        data: ToJsBuffer,
    },
    Error {
        message: String,
    },
    Closing,
    Close {
        clean: bool,
        code: u16,
        reason: String,
    },
}

struct CreatedSocket {
    id: u64,
    resource: Arc<Resource>,
    commands: mpsc::UnboundedReceiver<Command>,
    events: mpsc::Sender<Event>,
}

pub struct Sockets {
    runtime: tokio::runtime::Handle,
    tls: TlsConnector,
    resources: Arc<Mutex<HashMap<u64, Arc<Resource>>>>,
    next: AtomicU64,
    caps: Arc<Capabilities>,
}

impl Sockets {
    fn create(&self) -> io::Result<CreatedSocket> {
        let id = self
            .next
            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |id| {
                (id <= MAX_ID).then_some(id + 1)
            })
            .map_err(|_| io::Error::other("WebSocket IDs exhausted"))?;
        let (commands, receiving) = mpsc::unbounded_channel();
        let (events, event_queue) = mpsc::channel(EVENT_QUEUE);
        let resource = Arc::new(Resource {
            canceled: CancellationToken::new(),
            commands,
            events: AsyncMutex::new(event_queue),
        });
        self.resources.lock().unwrap().insert(id, resource.clone());
        Ok(CreatedSocket {
            id,
            resource,
            commands: receiving,
            events,
        })
    }

    pub fn accept(&self, upgrade: avril_http::Upgrade, protocol: String) -> io::Result<u64> {
        let CreatedSocket { id, resource, commands, events } = self.create()?;
        self.runtime.spawn(async move {
            let server_abort = upgrade.canceled.clone();
            let accepted = tokio::select! {
                _ = resource.canceled.cancelled() => Err("WebSocket upgrade canceled".into()),
                result = upgrade.accept() => result,
            };
            let (stream, _lease) = match accepted {
                Ok(accepted) => accepted,
                Err(error) => {
                    fail(&events, error.to_string()).await;
                    return;
                }
            };
            let (reader, writer) = tokio::io::split(TokioIo::new(stream));
            let (_, writer) = fastwebsockets::after_handshake_split(
                tokio::io::empty(),
                writer,
                fastwebsockets::Role::Server,
            );
            connected(
                (reader::Reader::new(reader, true), writer, protocol),
                resource.canceled.clone(),
                commands,
                events,
                Some(server_abort),
            )
            .await;
        });
        Ok(id)
    }

    fn resource(&self, id: u64) -> Result<Arc<Resource>, Error> {
        self.resources
            .lock()
            .unwrap()
            .get(&id)
            .cloned()
            .ok_or_else(|| "WebSocket is closed".into())
    }
}

fn accept_key(key: &str) -> String {
    let mut hasher = Sha1::new();
    hasher.update(key.as_bytes());
    hasher.update(b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11");
    BASE64.encode(hasher.finalize())
}

async fn attempt(
    tls: &TlsConnector,
    url: &Url,
    protocols: &[String],
    user_agent: &str,
) -> Result<(Reader, Writer, String), Error> {
    let user_agent = user_agent
        .chars()
        .map(|c| u8::try_from(c as u32))
        .collect::<Result<Vec<_>, _>>()?;
    let user_agent = HeaderValue::from_bytes(&user_agent)?;
    // hosts serialize with brackets for IPv6, which neither TcpStream nor SNI accept.
    let (host, server_name) = match url.host() {
        Some(Host::Domain(domain)) => (
            Cow::Borrowed(domain),
            ServerName::try_from(domain.to_owned())?,
        ),
        Some(Host::Ipv4(ip)) => (Cow::Owned(ip.to_string()), ServerName::IpAddress(ip.into())),
        Some(Host::Ipv6(ip)) => (Cow::Owned(ip.to_string()), ServerName::IpAddress(ip.into())),
        None => return Err("Missing WebSocket host".into()),
    };
    let port = url
        .port_or_known_default()
        .ok_or("Missing WebSocket port")?;
    let socket: Socket = if url.scheme() == "wss" {
        let stream = TcpStream::connect((host.as_ref(), port)).await?;
        Box::new(tls.connect(server_name, stream).await?)
    } else {
        Box::new(TcpStream::connect((host.as_ref(), port)).await?)
    };
    // The request target is origin-form, as the fetch-integrated spec requires.
    let target = match url.query() {
        Some(query) => format!("{}?{query}", url.path()),
        None => url.path().to_owned(),
    };
    let host_header = match url.port() {
        Some(port) => format!("{}:{port}", url.host().ok_or("Missing WebSocket host")?),
        None => url.host().ok_or("Missing WebSocket host")?.to_string(),
    };
    let key = fastwebsockets::handshake::generate_key();
    let mut request = Request::builder()
        .method(Method::GET)
        .uri(target)
        .header(header::HOST, host_header)
        .header(header::UPGRADE, "websocket")
        .header(header::CONNECTION, "Upgrade")
        .header("sec-websocket-key", key.as_str())
        .header("sec-websocket-version", "13")
        .header(header::USER_AGENT, user_agent);
    if !protocols.is_empty() {
        request = request.header("sec-websocket-protocol", protocols.join(", "));
    }
    let request = request.body(Empty::<Bytes>::new())?;
    let (mut sender, connection) =
        hyper::client::conn::http1::handshake(TokioIo::new(socket)).await?;
    // An abandoned handshake must not leave Hyper holding the transport.
    let _connection = AbortOnDropHandle::new(tokio::spawn(connection.with_upgrades()));
    let mut response = sender.send_request(request).await?;
    let protocol = validate(&response, &key, protocols)?;
    let upgraded = hyper::upgrade::on(&mut response).await?;
    let (reader, writer) = tokio::io::split(TokioIo::new(upgraded));
    let (_, writer) = fastwebsockets::after_handshake_split(
        tokio::io::empty(),
        writer,
        fastwebsockets::Role::Client,
    );
    Ok((reader::Reader::new(reader, false), writer, protocol))
}

fn validate(
    response: &hyper::Response<Incoming>,
    key: &str,
    protocols: &[String],
) -> Result<String, Error> {
    let headers = response.headers();
    let contains_token = |name: &str, token: &str| {
        headers.get_all(name).iter().any(|value| {
            value.to_str().is_ok_and(|value| {
                value
                    .split(',')
                    .any(|value| value.trim().eq_ignore_ascii_case(token))
            })
        })
    };
    if response.status() != StatusCode::SWITCHING_PROTOCOLS
        || !contains_token("upgrade", "websocket")
        || !contains_token("connection", "upgrade")
    {
        return Err("The WebSocket handshake response is invalid".into());
    }
    let accepted: Vec<&HeaderValue> = headers.get_all("sec-websocket-accept").iter().collect();
    if accepted.len() != 1 || accepted[0].to_str().ok() != Some(accept_key(key).as_str()) {
        return Err("The WebSocket handshake response is invalid".into());
    }
    // No extensions are offered, so any response extension is unsupported.
    if headers.contains_key("sec-websocket-extensions") {
        return Err("The WebSocket server selected an unsupported extension".into());
    }
    let chosen: Vec<&HeaderValue> = headers.get_all("sec-websocket-protocol").iter().collect();
    if protocols.is_empty() {
        if !chosen.is_empty() {
            return Err("The WebSocket server selected an unexpected subprotocol".into());
        }
        return Ok(String::new());
    }
    if chosen.len() != 1 {
        return Err("The WebSocket server did not select one subprotocol".into());
    }
    let value = chosen[0]
        .to_str()
        .map_err(|_| "The WebSocket server selected an invalid subprotocol")?;
    if value.is_empty() || !protocols.iter().any(|protocol| protocol == value) {
        return Err("The WebSocket server selected an invalid subprotocol".into());
    }
    Ok(value.to_owned())
}

/// Validate a close frame payload, returning its code and reason or the failure code.
fn close_payload(payload: &[u8]) -> Result<Option<(u16, String)>, u16> {
    match payload.len() {
        0 => Ok(None),
        1 => Err(1002),
        _ => {
            let code = u16::from_be_bytes([payload[0], payload[1]]);
            let reason = match std::str::from_utf8(&payload[2..]) {
                Ok(reason) => reason.to_owned(),
                Err(_) => return Err(1007),
            };
            if matches!(code, 1000..=1003 | 1007..=1014 | 3000..=4999) {
                Ok(Some((code, reason)))
            } else {
                Err(1002)
            }
        }
    }
}

/// The close frame code for a failed connection, per RFC 6455 §7.1.7.
fn failure_code(error: &WebSocketError) -> u16 {
    match error {
        WebSocketError::InvalidUTF8 => 1007,
        WebSocketError::FrameTooLarge => 1009,
        WebSocketError::IoError(_) | WebSocketError::UnexpectedEOF => 1006,
        _ => 1002,
    }
}

/// Forward frames to the event loop, which owns the write half.
async fn read_frames(
    mut reader: reader::Reader<impl AsyncRead + Unpin>,
    received: mpsc::Sender<Result<Received, WebSocketError>>,
    transport_closed: oneshot::Sender<()>,
    shutdown: CancellationToken,
) {
    loop {
        let frame = tokio::select! {
            _ = shutdown.cancelled() => return,
            frame = reader.read() => frame,
        };
        match frame {
            Ok(frame) => {
                let close = frame.opcode == OpCode::Close;
                if received.send(Ok(frame)).await.is_err() {
                    return;
                }
                if close {
                    // After Close, only transport termination matters; do not assemble more messages.
                    let mut discard = tokio::io::sink();
                    tokio::select! {
                        _ = shutdown.cancelled() => return,
                        _ = tokio::io::copy(&mut reader.stream, &mut discard) => {
                            let _ = transport_closed.send(());
                        }
                    }
                    return;
                }
            }
            Err(error) => {
                let _ = received.send(Err(error)).await;
                return;
            }
        }
    }
}

async fn write(
    writer: &mut WebSocketWrite<impl AsyncWrite + Unpin>,
    frame: Frame<'_>,
) -> Result<(), WebSocketError> {
    writer.write_frame(frame).await?;
    // TLS buffers writes until flushed.
    writer.flush().await
}

async fn push(events: &mpsc::Sender<Event>, event: Event) {
    // A dropped receiver means the connection is already gone.
    let _ = events.send(event).await;
}

async fn closed(events: &mpsc::Sender<Event>, clean: bool, code: u16, reason: String) {
    push(events, Event::Close { clean, code, reason }).await;
}

/// Fail the connection without an established stream.
async fn fail(events: &mpsc::Sender<Event>, message: String) {
    push(events, Event::Error { message }).await;
    closed(events, false, 1006, String::new()).await;
}

/// Fail an established connection: send the close frame for the failure, then error and close.
async fn abort(
    writer: &mut WebSocketWrite<impl AsyncWrite + Unpin>,
    events: &mpsc::Sender<Event>,
    message: String,
    code: u16,
    wrote_close: bool,
) {
    // An abruptly closed transport cannot carry a close frame (code 1006).
    if !wrote_close && code != 1006 {
        let _ = tokio::time::timeout(CLOSE_TIMEOUT, write(writer, Frame::close(code, b""))).await;
    }
    push(events, Event::Error { message }).await;
    closed(events, false, 1006, String::new()).await;
}

async fn connect(
    tls: TlsConnector,
    canceled: CancellationToken,
    commands: mpsc::UnboundedReceiver<Command>,
    events: mpsc::Sender<Event>,
    url: Url,
    protocols: Vec<String>,
    user_agent: String,
) {
    let established = tokio::select! {
        _ = canceled.cancelled() => {
            Err("The WebSocket connection was closed while connecting".into())
        }
        result = attempt(&tls, &url, &protocols, &user_agent) => match result {
            Ok(established) => Ok(established),
            Err(error) => Err(format!("failed to connect to WebSocket: {error}")),
        },
    };
    let established = match established {
        Ok(established) => established,
        Err(message) => {
            fail(&events, message).await;
            return;
        }
    };
    if canceled.is_cancelled() {
        fail(
            &events,
            "The WebSocket connection was closed while connecting".into(),
        )
        .await;
        return;
    }
    connected(established, canceled, commands, events, None).await;
}

async fn connected(
    (reader, mut writer, protocol): (Reader, Writer, String),
    canceled: CancellationToken,
    commands: mpsc::UnboundedReceiver<Command>,
    events: mpsc::Sender<Event>,
    server_abort: Option<CancellationToken>,
) {
    push(&events, Event::Open { protocol }).await;
    let (received, frames) = mpsc::channel(READ_QUEUE);
    let shutdown = CancellationToken::new();
    let (ended, transport_closed) = oneshot::channel();
    let reading = tokio::spawn(read_frames(reader, received, ended, shutdown.clone()));
    tokio::select! {
        // close() during CONNECTING can land after `open` was queued. Keep
        // cancellation outside the loop so it also interrupts blocked writes.
        _ = canceled.cancelled() => {
            fail(&events, "The WebSocket connection was closed while connecting".into()).await;
        }
        _ = async {
            if let Some(token) = &server_abort { token.cancelled().await; }
            else { std::future::pending::<()>().await; }
        } => { closed(&events, false, 1006, String::new()).await; }
        _ = events.closed() => {}
        _ = stream_loop(&mut writer, commands, frames, transport_closed, &events, server_abort.is_some()) => {}
    }
    // Drops the read half, closing the transport with the connection.
    shutdown.cancel();
    let _ = reading.await;
}

async fn stream_loop(
    writer: &mut WebSocketWrite<impl AsyncWrite + Unpin>,
    mut commands: mpsc::UnboundedReceiver<Command>,
    mut frames: mpsc::Receiver<Result<Received, WebSocketError>>,
    transport_closed: oneshot::Receiver<()>,
    events: &mpsc::Sender<Event>,
    server: bool,
) {
    let mut wrote_close = false;
    let mut timer = Box::pin(tokio::time::sleep(CLOSE_TIMEOUT));
    loop {
        tokio::select! {
            _ = timer.as_mut(), if wrote_close => {
                // The peer never completed the closing handshake.
                closed(events, false, 1006, String::new()).await;
                return;
            }
            command = commands.recv() => match command {
                Some(Command::Send(kind, payload, ack)) => {
                    let result = if wrote_close {
                        Ok(())
                    } else {
                        let frame = match kind {
                            FrameKind::Text => Frame::text(payload.into()),
                            FrameKind::Binary => Frame::binary(payload.into()),
                        };
                        write(writer, frame).await
                    };
                    let _ = ack.send(result.is_ok() && !wrote_close);
                    if let Err(error) = result {
                        abort(writer, events, error.to_string(), 1006, wrote_close).await;
                        return;
                    }
                }
                Some(Command::Close(payload)) => {
                    if wrote_close {
                        continue;
                    }
                    let frame = match payload {
                        Some(payload) => Frame::close_raw(payload.into()),
                        None => Frame::close_raw(Vec::new().into()),
                    };
                    let deadline = tokio::time::Instant::now() + CLOSE_TIMEOUT;
                    match tokio::time::timeout_at(deadline, write(writer, frame)).await {
                        Ok(Ok(())) => {
                            wrote_close = true;
                            timer.as_mut().reset(deadline);
                        }
                        Ok(Err(error)) => {
                            abort(writer, events, error.to_string(), 1006, wrote_close).await;
                            return;
                        }
                        Err(_) => {
                            closed(events, false, 1006, String::new()).await;
                            return;
                        }
                    }
                }
                None => {
                    closed(events, false, 1006, String::new()).await;
                    return;
                }
            },
            frame = frames.recv() => match frame {
                Some(Err(error)) => {
                    if wrote_close {
                        closed(events, false, 1006, String::new()).await;
                    } else {
                        abort(writer, events, error.to_string(), failure_code(&error), false).await;
                    }
                    return;
                }
                None => {
                    closed(events, false, 1006, String::new()).await;
                    return;
                }
                Some(Ok(frame)) => match frame.opcode {
                    OpCode::Ping if !wrote_close => {
                        if let Err(error) = write(writer, Frame::pong(frame.payload.into())).await {
                            abort(writer, events, error.to_string(), 1006, wrote_close).await;
                            return;
                        }
                    }
                    OpCode::Close => match close_payload(&frame.payload) {
                        Ok(payload) => {
                            push(events, Event::Closing).await;
                            let (code, reason) = match payload {
                                Some((code, reason)) => (code, reason),
                                None => (1005, String::new()),
                            };
                            if !wrote_close {
                                let echo = Frame::close_raw(frame.payload.into());
                                let result = tokio::time::timeout(CLOSE_TIMEOUT, write(writer, echo)).await;
                                let message = match result {
                                    Ok(Ok(())) => None,
                                    Ok(Err(error)) => Some(error.to_string()),
                                    Err(_) => Some("The WebSocket close reply timed out".into()),
                                };
                                if let Some(message) = message {
                                    push(events, Event::Error { message }).await;
                                    closed(events, false, code, reason).await;
                                    return;
                                }
                            }
                            // The handshake is complete. Closing our own transport after
                            // the grace period is still a clean closure (RFC 6455 §7.1.1).
                            if !server { let _ = tokio::time::timeout(TCP_CLOSE_TIMEOUT, transport_closed).await; }
                            closed(events, true, code, reason).await;
                            return;
                        }
                        Err(code) => {
                            abort(
                                writer,
                                events,
                                "The WebSocket connection received an invalid close frame".into(),
                                code,
                                wrote_close,
                            )
                            .await;
                            return;
                        }
                    },
                    OpCode::Text => match String::from_utf8(frame.payload) {
                        Ok(data) => {
                            push(events, Event::Text { data }).await;
                        }
                        Err(_) => {
                            abort(
                                writer,
                                events,
                                "The WebSocket connection received invalid UTF-8 in a text message".into(),
                                1007,
                                wrote_close,
                            )
                            .await;
                            return;
                        }
                    },
                    OpCode::Binary => {
                        let data = ToJsBuffer::from(frame.payload);
                        push(events, Event::Binary { data }).await;
                    }
                    _ => {}
                },
            },
        }
    }
}

pub fn init(
    runtime: tokio::runtime::Handle,
    caps: Arc<Capabilities>,
) -> io::Result<(Extension, Arc<Sockets>)> {
    let mut roots = RootCertStore::empty();
    roots.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
    let config = rustls::ClientConfig::builder_with_provider(Arc::new(
        rustls::crypto::ring::default_provider(),
    ))
    .with_safe_default_protocol_versions()
    .map_err(io::Error::other)?
    .with_root_certificates(roots)
    .with_no_client_auth();
    let tls = TlsConnector::from(Arc::new(config));
    let state = Arc::new(Sockets {
        runtime,
        tls,
        resources: Arc::new(Mutex::new(HashMap::new())),
        next: AtomicU64::new(1),
        caps,
    });
    let create = state.clone();
    let send = state.clone();
    let close = state.clone();
    let fail = state.clone();
    let retire = state.clone();
    let wait = state.clone();
    Ok((
        Extension {
            name: "ext:websocket",
            dependencies: &["ext:blob", "ext:encoding", "ext:url", "ext:web"],
            ops: vec![
                Op::new(
                    "op_ws_create",
                    move |(url, protocols, user_agent): (String, Vec<String>, String)| {
                        let url = Url::parse(&url).map_err(io::Error::other)?;
                        if !matches!(url.scheme(), "ws" | "wss") {
                            return Err(io::Error::other("Invalid WebSocket scheme"));
                        }
                        let host = url
                            .host_str()
                            .ok_or_else(|| io::Error::other("Missing WebSocket host"))?;
                        let port = url
                            .port_or_known_default()
                            .ok_or_else(|| io::Error::other("Missing WebSocket port"))?;
                        create
                            .caps
                            .check_net(host, port)
                            .map_err(io::Error::other)?;
                        let CreatedSocket {
                            id,
                            resource,
                            commands: receiving,
                            events,
                        } = create.create()?;
                        create.runtime.spawn(connect(
                            create.tls.clone(),
                            resource.canceled.clone(),
                            receiving,
                            events,
                            url,
                            protocols,
                            user_agent,
                        ));
                        Ok::<_, io::Error>(id)
                    },
                ),
                // The command is enqueued at dispatch so queue order equals call order;
                // the future only awaits the write acknowledgement.
                Op::asynchronous_unreferenced(
                    "op_ws_send",
                    move |(id, text, bytes): (u64, bool, Option<JsBuffer>)| {
                        let (ack, receive) = oneshot::channel();
                        if let Ok(resource) = send.resource(id) {
                            let kind = if text {
                                FrameKind::Text
                            } else {
                                FrameKind::Binary
                            };
                            let payload = bytes.map(|bytes| bytes.to_vec()).unwrap_or_default();
                            let _ = resource.commands.send(Command::Send(kind, payload, ack));
                        } else {
                            drop(ack);
                        }
                        async move { Ok::<_, io::Error>(receive.await.unwrap_or(false)) }
                    },
                ),
                Op::new(
                    "op_ws_close",
                    move |(id, code, reason): (u64, Option<u16>, Option<String>)| {
                        let payload = match (code, reason) {
                            (Some(code), Some(reason)) => {
                                let mut payload = code.to_be_bytes().to_vec();
                                payload.extend_from_slice(reason.as_bytes());
                                Some(payload)
                            }
                            (Some(code), None) => Some(code.to_be_bytes().to_vec()),
                            // A reason without a code is not a valid close payload; use 1000.
                            (None, Some(reason)) => {
                                let mut payload = 1000u16.to_be_bytes().to_vec();
                                payload.extend_from_slice(reason.as_bytes());
                                Some(payload)
                            }
                            (None, None) => None,
                        };
                        if let Ok(resource) = close.resource(id) {
                            let _ = resource.commands.send(Command::Close(payload));
                        }
                        Ok::<_, io::Error>(())
                    },
                ),
                Op::new("op_ws_fail", move |id: u64| {
                    if let Ok(resource) = fail.resource(id) {
                        resource.canceled.cancel();
                    }
                    Ok::<_, io::Error>(())
                }),
                Op::new("op_ws_retire", move |id: u64| {
                    if let Some(resource) = retire.resources.lock().unwrap().remove(&id) {
                        resource.canceled.cancel();
                    }
                    Ok::<_, io::Error>(())
                }),
                Op::asynchronous("op_ws_wait", move |id: u64| {
                    let resource = wait.resource(id);
                    async move {
                        let Ok(resource) = resource else {
                            return Err(io::Error::other("WebSocket is closed"));
                        };
                        let mut events = resource.events.lock().await;
                        events
                            .recv()
                            .await
                            .ok_or_else(|| io::Error::other("WebSocket task ended"))
                    }
                }),
            ],
            modules: &[
                (
                    "ext:websocket/main.js",
                    include_str!(concat!(env!("OUT_DIR"), "/main.js")),
                ),
                (
                    "ext:websocket/websocket.js",
                    include_str!(concat!(env!("OUT_DIR"), "/websocket.js")),
                ),
            ],
            entry_point: Some("ext:websocket/main.js"),
        },
        state,
    ))
}

#[cfg(test)]
mod tests {
    use super::*;

    enum Input {
        Command(Command),
        Frame(Result<Received, WebSocketError>),
    }

    async fn events_with_blocked_transport(input: Input) -> Vec<Event> {
        let (socket, _peer) = tokio::io::duplex(1);
        let ws = fastwebsockets::WebSocket::after_handshake(socket, fastwebsockets::Role::Client);
        let (_, mut writer) = ws.split(tokio::io::split);
        let (commands, queued) = mpsc::unbounded_channel();
        let (frames, received) = mpsc::channel(1);
        let (events, mut emitted) = mpsc::channel(2);
        let (_ended, transport_closed) = oneshot::channel();
        match input {
            Input::Command(command) => commands.send(command).unwrap(),
            Input::Frame(frame) => frames.send(frame).await.unwrap(),
        }
        let started = tokio::time::Instant::now();
        let connection = tokio::spawn(async move {
            stream_loop(
                &mut writer,
                queued,
                received,
                transport_closed,
                &events,
                false,
            )
            .await;
        });
        let mut result = Vec::new();
        while let Some(event) = emitted.recv().await {
            result.push(event);
        }
        connection.await.unwrap();
        assert_eq!(started.elapsed(), Duration::from_secs(30));
        result
    }

    #[tokio::test(start_paused = true)]
    async fn client_close_write_cannot_stall_shutdown() {
        let events = events_with_blocked_transport(Input::Command(Command::Close(None))).await;
        assert!(matches!(
            events.as_slice(),
            [Event::Close { clean: false, code: 1006, reason }] if reason.is_empty()
        ));
    }

    #[tokio::test(start_paused = true)]
    async fn peer_close_reply_cannot_stall_shutdown() {
        let events = events_with_blocked_transport(Input::Frame(Ok(Received {
            opcode: OpCode::Close,
            payload: [3001u16.to_be_bytes().as_slice(), b"bye"].concat(),
        })))
        .await;
        assert!(matches!(
            events.as_slice(),
            [Event::Closing, Event::Error { .. }, Event::Close { clean: false, code: 3001, reason }]
                if reason == "bye"
        ));
    }

    #[tokio::test(start_paused = true)]
    async fn failure_close_write_cannot_stall_shutdown() {
        let events =
            events_with_blocked_transport(Input::Frame(Err(WebSocketError::InvalidValue))).await;
        assert!(matches!(
            events.as_slice(),
            [Event::Error { .. }, Event::Close { clean: false, code: 1006, reason }]
                if reason.is_empty()
        ));
    }

    struct Connection {
        peer: tokio::io::DuplexStream,
        events: mpsc::Receiver<Event>,
        task: tokio::task::JoinHandle<()>,
        _commands: mpsc::UnboundedSender<Command>,
    }

    async fn completed_handshake(client_initiated: bool) -> Connection {
        use tokio::io::{AsyncReadExt, AsyncWriteExt};

        let (socket, mut peer) = tokio::io::duplex(1024);
        let (reader, writer) = tokio::io::split(socket);
        let (_, mut writer) = fastwebsockets::after_handshake_split(
            tokio::io::empty(),
            writer,
            fastwebsockets::Role::Client,
        );
        let (commands, queued) = mpsc::unbounded_channel();
        let (received, frames) = mpsc::channel(READ_QUEUE);
        let (ended, transport_closed) = oneshot::channel();
        let (events, mut emitted) = mpsc::channel(EVENT_QUEUE);
        let shutdown = CancellationToken::new();
        let reading = tokio::spawn(read_frames(
            reader::Reader::new(reader, false),
            received,
            ended,
            shutdown.clone(),
        ));
        let task = tokio::spawn(async move {
            stream_loop(
                &mut writer,
                queued,
                frames,
                transport_closed,
                &events,
                false,
            )
            .await;
            shutdown.cancel();
            reading.await.unwrap();
        });
        let payload = [3001u16.to_be_bytes().as_slice(), b"bye"].concat();
        let server_close = [b"\x88\x05".as_slice(), &payload].concat();
        if client_initiated {
            commands
                .send(Command::Close(Some(payload.clone())))
                .unwrap();
        } else {
            peer.write_all(&server_close).await.unwrap();
        }
        let mut reply = [0; 11];
        peer.read_exact(&mut reply).await.unwrap();
        assert_eq!(&reply[..2], b"\x88\x85");
        for index in 0..5 {
            assert_eq!(reply[6 + index] ^ reply[2 + index % 4], payload[index]);
        }
        if client_initiated {
            peer.write_all(&server_close).await.unwrap();
        }
        assert!(matches!(emitted.recv().await, Some(Event::Closing)));
        Connection {
            peer,
            events: emitted,
            task,
            _commands: commands,
        }
    }

    async fn assert_waits_for_eof(client_initiated: bool) {
        use tokio::io::AsyncWriteExt;

        let mut connection = completed_handshake(client_initiated).await;
        tokio::time::advance(Duration::from_secs(1)).await;
        assert!(matches!(
            connection.events.try_recv(),
            Err(mpsc::error::TryRecvError::Empty)
        ));
        connection.peer.shutdown().await.unwrap();
        assert!(matches!(connection.events.recv().await,
            Some(Event::Close { clean: true, code: 3001, reason }) if reason == "bye"));
        connection.task.await.unwrap();
        assert!(connection.events.recv().await.is_none());
    }

    #[tokio::test(start_paused = true)]
    async fn client_close_waits_for_transport_eof() {
        assert_waits_for_eof(true).await;
    }

    #[tokio::test(start_paused = true)]
    async fn peer_close_waits_for_transport_eof() {
        assert_waits_for_eof(false).await;
    }

    #[tokio::test(start_paused = true)]
    async fn tcp_close_timeout_still_reports_clean_close() {
        use tokio::io::AsyncReadExt;

        let mut connection = completed_handshake(false).await;
        let started = tokio::time::Instant::now();
        assert!(matches!(connection.events.recv().await,
            Some(Event::Close { clean: true, code: 3001, reason }) if reason == "bye"));
        assert_eq!(started.elapsed(), Duration::from_secs(2));
        connection.task.await.unwrap();
        let mut byte = [0];
        assert_eq!(connection.peer.read(&mut byte).await.unwrap(), 0);
        assert!(connection.events.recv().await.is_none());
    }
}