//! A scripted WebSocket server for tests: it accepts one handshake per connection and then lets
//! the test send raw frames and read the client's. Built on `tls_server`. Do not edit.

#![allow(dead_code)] // each test file uses a different part of this module

use std::io::{Read, Write};
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;

use gatewayd::ws::handshake::accept_for;
use rustls::ServerConfig;

use crate::tls_server::{Conn, serve};

pub struct Peer {
    pub conn: Box<dyn Conn>,
    /// The request head the client sent, for tests that check it.
    pub request: String,
}

/// A frame from the client: opcode, whether it was masked, and the unmasked payload.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ClientFrame {
    pub opcode: u8,
    pub masked: bool,
    pub mask: [u8; 4],
    pub payload: Vec<u8>,
}

impl Peer {
    pub fn send(&mut self, bytes: &[u8]) {
        let _ = self.conn.write_all(bytes);
        let _ = self.conn.flush();
    }

    /// An unmasked server frame, FIN set.
    pub fn frame(&mut self, opcode: u8, payload: &[u8]) {
        let mut out = vec![0x80 | opcode];
        if payload.len() < 126 {
            out.push(payload.len() as u8);
        } else {
            out.push(126);
            out.extend_from_slice(&(payload.len() as u16).to_be_bytes());
        }
        out.extend_from_slice(payload);
        self.send(&out);
    }

    pub fn text(&mut self, text: &str) {
        self.frame(0x1, text.as_bytes());
    }

    fn read_exact(&mut self, n: usize) -> Option<Vec<u8>> {
        let mut buf = vec![0u8; n];
        self.conn.read_exact(&mut buf).ok()?;
        Some(buf)
    }

    /// The next frame from the client, or `None` when it has gone.
    pub fn read_frame(&mut self) -> Option<ClientFrame> {
        let head = self.read_exact(2)?;
        let opcode = head[0] & 0x0F;
        let masked = head[1] & 0x80 != 0;
        let len = match head[1] & 0x7F {
            126 => u16::from_be_bytes(self.read_exact(2)?.try_into().ok()?) as usize,
            127 => u64::from_be_bytes(self.read_exact(8)?.try_into().ok()?) as usize,
            n => n as usize,
        };
        let mask: [u8; 4] = if masked {
            self.read_exact(4)?.try_into().ok()?
        } else {
            [0; 4]
        };
        let raw = self.read_exact(len)?;
        let payload = raw
            .iter()
            .zip(mask.iter().cycle())
            .map(|(b, m)| b ^ m)
            .collect();
        Some(ClientFrame {
            opcode,
            masked,
            mask,
            payload,
        })
    }

    pub fn pause(&self, d: Duration) {
        std::thread::sleep(d);
    }
}

/// Serve WebSocket connections: complete the handshake (or answer `refuse_with` instead), then run
/// `script` on the connection.
pub fn serve_ws<F>(
    tls: Option<Arc<ServerConfig>>,
    refuse_with: Option<&'static str>,
    script: F,
) -> SocketAddr
where
    F: Fn(Peer) + Send + Sync + 'static,
{
    serve(tls, move |mut conn| {
        let mut head = Vec::new();
        let mut byte = [0u8; 1];
        while !head.ends_with(b"\r\n\r\n") {
            if conn.read(&mut byte).map(|n| n == 0).unwrap_or(true) {
                return;
            }
            head.push(byte[0]);
        }
        let request = String::from_utf8_lossy(&head).into_owned();
        if let Some(reply) = refuse_with {
            let _ = conn.write_all(reply.as_bytes());
            return;
        }
        let key = request
            .lines()
            .find_map(|l| l.strip_prefix("Sec-WebSocket-Key: "))
            .unwrap_or_default()
            .trim()
            .to_string();
        let reply = format!(
            "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: {}\r\n\r\n",
            accept_for(&key)
        );
        let _ = conn.write_all(reply.as_bytes());
        let _ = conn.flush();
        script(Peer { conn, request });
    })
}