use std::io::{self, Read, Write};
use std::net::{TcpListener, TcpStream};
use std::time::{Duration, Instant};

pub const CONTINUATION: u8 = 0;
pub const TEXT: u8 = 1;
pub const BINARY: u8 = 2;
pub const CLOSE: u8 = 8;
pub const PING: u8 = 9;
pub const PONG: u8 = 10;

pub fn accept(key: &str) -> String {
    use base64::Engine;
    use sha1::{Digest, Sha1};
    let mut hasher = Sha1::new();
    hasher.update(key.as_bytes());
    hasher.update(b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11");
    base64::engine::general_purpose::STANDARD.encode(hasher.finalize())
}

pub struct Frame {
    pub opcode: u8,
    pub payload: Vec<u8>,
}

impl Frame {
    pub fn close(code: u16, reason: &str) -> Frame {
        let mut payload = code.to_be_bytes().to_vec();
        payload.extend_from_slice(reason.as_bytes());
        Frame { opcode: CLOSE, payload }
    }

    pub fn code(&self) -> u16 {
        u16::from_be_bytes([self.payload[0], self.payload[1]])
    }

    pub fn reason(&self) -> String {
        String::from_utf8_lossy(&self.payload[2..]).into_owned()
    }
}

pub fn read_exact(stream: &mut TcpStream, count: usize) -> io::Result<Vec<u8>> {
    let mut buffer = vec![0; count];
    stream.read_exact(&mut buffer)?;
    Ok(buffer)
}

pub fn read_frame(stream: &mut TcpStream) -> io::Result<Frame> {
    let head = read_exact(stream, 2)?;
    let opcode = head[0] & 0x0f;
    let masked = head[1] & 0x80 != 0;
    assert!(masked, "client frames must be masked");
    assert_ne!(head[0] & 0x80, 0, "expected a complete client frame");
    assert_eq!(head[0] & 0x70, 0, "client must not set reserved bits");
    let mut length = usize::from(head[1] & 0x7f);
    match length {
        126 => {
            let extra = read_exact(stream, 2)?;
            length = u16::from_be_bytes([extra[0], extra[1]]) as usize;
        }
        127 => {
            let extra = read_exact(stream, 8)?;
            length = u64::from_be_bytes(extra.try_into().unwrap()) as usize;
        }
        _ => {}
    }
    let mask = read_exact(stream, 4)?;
    let mut payload = read_exact(stream, length)?;
    for (index, byte) in payload.iter_mut().enumerate() {
        *byte ^= mask[index % 4];
    }
    Ok(Frame { opcode, payload })
}

pub fn write_frame(
    stream: &mut TcpStream,
    fin: bool,
    opcode: u8,
    payload: &[u8],
) -> io::Result<()> {
    let mut head = vec![(if fin { 0x80 } else { 0 }) | opcode];
    let length = payload.len();
    if length < 126 {
        head.push(length as u8);
    } else if length <= u16::MAX as usize {
        head.push(126);
        head.extend_from_slice(&(length as u16).to_be_bytes());
    } else {
        head.push(127);
        head.extend_from_slice(&(length as u64).to_be_bytes());
    }
    stream.write_all(&head)?;
    stream.write_all(payload)
}

/// A handshook test server connection.
pub struct Session {
    pub stream: TcpStream,
    pub path: String,
    pub headers: String,
    pub protocol: String,
}

impl Session {
    pub fn read(&mut self) -> io::Result<Frame> {
        read_frame(&mut self.stream)
    }

    pub fn write(&mut self, opcode: u8, payload: &[u8]) -> io::Result<()> {
        write_frame(&mut self.stream, true, opcode, payload)
    }

    /// Assert the handshake request matches RFC 6455 and the runtime conventions.
    pub fn assert_handshake(&self) {
        assert!(
            self.path.starts_with('/'),
            "request target must be origin-form"
        );
        let has = |name: &str, value: &str| {
            self.headers.lines().any(|line| {
                line.split_once(':').is_some_and(|(field, found)| {
                    field == name && found.trim().eq_ignore_ascii_case(value)
                })
            })
        };
        assert!(has("upgrade", "websocket"));
        assert!(has("connection", "upgrade"));
        assert!(has("sec-websocket-version", "13"));
        assert!(
            self.headers
                .lines()
                .any(|line| line.starts_with("sec-websocket-key: ") && line.len() > 16)
        );
        assert!(has(
            "user-agent",
            concat!("avril/", env!("CARGO_PKG_VERSION"))
        ));
    }

    /// Echo data messages back until the peer closes, then echo the close frame.
    pub fn echo(mut self) -> io::Result<()> {
        self.assert_handshake();
        self.accept(None)?;
        loop {
            let frame = self.read()?;
            match frame.opcode {
                CLOSE => {
                    self.write(CLOSE, &frame.payload)?;
                    return Ok(());
                }
                TEXT | BINARY => self.write(frame.opcode, &frame.payload)?,
                _ => {}
            }
        }
    }

    pub fn accept(&mut self, protocol: Option<&str>) -> io::Result<()> {
        let mut response = format!(
            "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: {}\r\n",
            accept(&self.key())
        );
        if let Some(protocol) = protocol {
            response.push_str(&format!("Sec-WebSocket-Protocol: {protocol}\r\n"));
        }
        response.push_str("\r\n");
        self.stream.write_all(response.as_bytes())
    }

    pub fn reject(&mut self) -> io::Result<()> {
        self.stream
            .write_all(b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n")
    }

    pub fn key(&self) -> String {
        self.headers
            .lines()
            .find(|line| line.starts_with("sec-websocket-key:"))
            .map(|line| line.split_once(':').unwrap().1.trim().to_owned())
            .expect("Sec-WebSocket-Key header")
    }
}

/// Accept `count` connections and run `handler` for each.
pub fn server(
    count: usize,
    handler: impl Fn(Session) -> io::Result<()> + Send + 'static,
) -> (String, std::thread::JoinHandle<()>) {
    let (url, thread) = server_on("127.0.0.1", count, handler);
    (url, thread)
}

/// Accept one connection, panicking instead of blocking forever when the runtime never
/// connects: tests join the server thread, so a missed connection would otherwise hang.
pub fn incoming(listener: &TcpListener) -> TcpStream {
    listener.set_nonblocking(true).unwrap();
    let deadline = Instant::now() + Duration::from_secs(10);
    let stream = loop {
        match listener.accept() {
            Ok((stream, _)) => break stream,
            Err(error)
                if error.kind() == io::ErrorKind::WouldBlock && Instant::now() < deadline =>
            {
                std::thread::sleep(Duration::from_millis(5));
            }
            Err(error) => panic!("the runtime never connected: {error}"),
        }
    };
    stream.set_nonblocking(false).unwrap();
    stream
        .set_read_timeout(Some(Duration::from_secs(10)))
        .unwrap();
    stream
        .set_write_timeout(Some(Duration::from_secs(10)))
        .unwrap();
    stream
}

pub fn server_on(
    host: &str,
    count: usize,
    handler: impl Fn(Session) -> io::Result<()> + Send + 'static,
) -> (String, std::thread::JoinHandle<()>) {
    let listener = TcpListener::bind((host, 0)).unwrap();
    // SocketAddr serializes IPv6 literals with brackets.
    let address = listener.local_addr().unwrap().to_string();
    let thread = std::thread::spawn(move || {
        for _ in 0..count {
            let mut stream = incoming(&listener);
            let mut request = Vec::new();
            loop {
                let mut byte = [0; 1];
                match stream.read_exact(&mut byte) {
                    Ok(()) => request.push(byte[0]),
                    Err(_) => break,
                }
                if request.ends_with(b"\r\n\r\n") {
                    break;
                }
            }
            let request = String::from_utf8_lossy(&request);
            // Header names are matched case-insensitively; values are not.
            let request = request
                .lines()
                .map(|line| match line.split_once(':') {
                    Some((name, value)) => format!("{}:{value}", name.to_ascii_lowercase()),
                    None => line.to_owned(),
                })
                .collect::<Vec<_>>()
                .join("\n");
            let first = request.lines().next().unwrap_or_default().to_owned();
            let path = first.split(' ').nth(1).unwrap_or_default().to_owned();
            let headers = request.lines().skip(1).collect::<Vec<_>>().join("\n");
            let protocol = headers
                .lines()
                .find(|line| line.starts_with("sec-websocket-protocol:"))
                .map(|line| line.split_once(':').unwrap().1.trim().to_owned())
                .unwrap_or_default();
            let session = Session { stream, path, headers, protocol };
            handler(session).unwrap();
        }
    });
    (format!("ws://{address}"), thread)
}

pub fn connect(url: &str, body: &str) -> String {
    format!("const ws = new WebSocket('{url}');{body}")
}