w-gui 0.21.8

use std::net::{TcpListener, TcpStream};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::sync::mpsc;
use std::thread;
use std::time::Duration;

use indexmap::IndexMap;

use crate::element::ElementDecl;
use crate::protocol::{ClientMsg, ServerMsg};

const HTML_TEMPLATE: &str = include_str!("frontend.html");

fn html_escape(s: &str) -> String {
    s.replace('&', "&")
        .replace('<', "&lt;")
        .replace('>', "&gt;")
        .replace('"', "&quot;")
}

/// Spawn the HTTP server thread. Returns the join handle.
pub fn spawn_http(
    shutdown: Arc<AtomicBool>,
    port: u16,
    bind_addr: &str,
    title: &str,
    favicon: Option<Vec<u8>>,
) -> thread::JoinHandle<()> {
    let bind = bind_addr.to_string();
    let html = HTML_TEMPLATE.replace("__WGUI_TITLE__", &html_escape(title));
    thread::Builder::new()
        .name("wgui-http".into())
        .spawn(move || run_http(shutdown, port, &bind, html, favicon))
        .expect("failed to spawn wgui HTTP thread")
}

/// Spawn the WebSocket server thread. Returns the join handle.
pub fn spawn_ws(
    ws_rx: mpsc::Receiver<Vec<ServerMsg>>,
    edit_tx: mpsc::Sender<(String, crate::element::Value)>,
    port: u16,
    bind_addr: &str,
) -> thread::JoinHandle<()> {
    let bind = bind_addr.to_string();
    thread::Builder::new()
        .name("wgui-ws".into())
        .spawn(move || run_ws(ws_rx, edit_tx, port, &bind))
        .expect("failed to spawn wgui WS thread")
}

fn run_http(shutdown: Arc<AtomicBool>, port: u16, bind_addr: &str, html: String, favicon: Option<Vec<u8>>) {
    let server = tiny_http::Server::http(format!("{bind_addr}:{port}"))
        .unwrap_or_else(|e| panic!("wgui: failed to bind HTTP on port {port}: {e}"));

    log::info!("wgui: serving UI at http://{bind_addr}:{port}");

    loop {
        if shutdown.load(Ordering::Acquire) {
            break;
        }

        // Use a timeout so we can check shutdown periodically
        match server.recv_timeout(Duration::from_millis(200)) {
            Ok(Some(request)) => {
                match request.url() {
                    "/" | "/index.html" => {
                        let response = tiny_http::Response::from_string(&html)
                            .with_header(
                                "Content-Type: text/html; charset=utf-8"
                                    .parse::<tiny_http::Header>()
                                    .unwrap(),
                            );
                        let _ = request.respond(response);
                    }
                    "/favicon.png" => {
                        if let Some(ref icon) = favicon {
                            let response = tiny_http::Response::from_data(icon.clone())
                                .with_header(
                                    "Content-Type: image/png"
                                        .parse::<tiny_http::Header>()
                                        .unwrap(),
                                );
                            let _ = request.respond(response);
                        } else {
                            let _ = request.respond(
                                tiny_http::Response::from_string("404 Not Found")
                                    .with_status_code(404),
                            );
                        }
                    }
                    _ => {
                        let _ = request.respond(
                            tiny_http::Response::from_string("404 Not Found")
                                .with_status_code(404),
                        );
                    }
                }
            }
            Ok(None) => {} // timeout, loop again
            Err(e) => {
                log::error!("wgui HTTP error: {e}");
                break;
            }
        }
    }
}

fn run_ws(
    ws_rx: mpsc::Receiver<Vec<ServerMsg>>,
    edit_tx: mpsc::Sender<(String, crate::element::Value)>,
    port: u16,
    bind_addr: &str,
) {
    let listener = TcpListener::bind(format!("{bind_addr}:{port}"))
        .unwrap_or_else(|e| panic!("wgui: failed to bind WS on port {port}: {e}"));

    listener
        .set_nonblocking(true)
        .expect("failed to set WS listener non-blocking");

    log::info!("wgui: WebSocket listening on ws://{bind_addr}:{port}");

    // Local mirror of UI state, updated from incoming ServerMsg
    let mut mirror: IndexMap<String, ElementDecl> = IndexMap::new();
    let mut clients: Vec<tungstenite::WebSocket<TcpStream>> = Vec::new();

    loop {
        // Accept new connections
        loop {
            match listener.accept() {
                Ok((stream, _addr)) => {
                    stream.set_nonblocking(false).ok();
                    stream.set_read_timeout(Some(Duration::from_millis(1))).ok();
                    match tungstenite::accept(stream) {
                        Ok(mut ws) => {
                            log::info!("wgui: new WebSocket client connected (total: {})", clients.len() + 1);
                            // Send snapshot to new client
                            let snapshot = ServerMsg::Snapshot {
                                elements: mirror.values().cloned().collect(),
                            };
                            if let Ok(json) = serde_json::to_string(&snapshot) {
                                let _ = ws.send(tungstenite::Message::Text(json.into()));
                            }
                            ws.get_ref().set_read_timeout(Some(Duration::from_millis(1))).ok();
                            clients.push(ws);
                        }
                        Err(e) => log::error!("wgui: WS handshake failed: {e}"),
                    }
                }
                Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => break,
                Err(e) => {
                    log::error!("wgui WS accept error: {e}");
                    break;
                }
            }
        }

        // Drain pending messages from game loop, apply to mirror, broadcast
        let mut got_msgs = false;
        loop {
            match ws_rx.try_recv() {
                Ok(msgs) => {
                    got_msgs = true;
                    apply_messages_to_mirror(&mut mirror, &msgs);
                    for msg in &msgs {
                        if let Ok(json) = serde_json::to_string(msg) {
                            let text = tungstenite::Message::Text(json.into());
                            clients.retain_mut(|ws| {
                                ws.send(text.clone()).is_ok()
                            });
                        }
                    }
                }
                Err(mpsc::TryRecvError::Empty) => break,
                Err(mpsc::TryRecvError::Disconnected) => {
                    for ws in &mut clients {
                        let _ = ws.close(None);
                    }
                    return;
                }
            }
        }

        // Read incoming messages from all clients
        clients.retain_mut(|ws| {
            loop {
                match ws.read() {
                    Ok(tungstenite::Message::Text(text)) => {
                        if let Ok(client_msg) = serde_json::from_str::<ClientMsg>(&text) {
                            match client_msg {
                                ClientMsg::Set { id, value } => {
                                    let _ = edit_tx.send((id, value));
                                }
                            }
                        }
                    }
                    Ok(tungstenite::Message::Close(_)) => {
                        let _ = ws.close(None);
                        log::info!("wgui: WebSocket client disconnected");
                        return false;
                    }
                    Ok(_) => {} // ping/pong/binary
                    Err(tungstenite::Error::Io(ref e))
                        if e.kind() == std::io::ErrorKind::WouldBlock
                            || e.kind() == std::io::ErrorKind::TimedOut =>
                    {
                        return true; // No more data, keep client
                    }
                    Err(tungstenite::Error::ConnectionClosed | tungstenite::Error::AlreadyClosed) => {
                        log::info!("wgui: WebSocket client disconnected");
                        return false;
                    }
                    Err(e) => {
                        log::warn!("wgui: WS read error: {e}");
                        return false;
                    }
                }
            }
        });

        if !got_msgs && clients.is_empty() {
            thread::sleep(Duration::from_millis(50));
        } else if !got_msgs {
            thread::sleep(Duration::from_millis(8));
        }
    }
}

/// Apply a batch of ServerMsg to the local mirror.
fn apply_messages_to_mirror(mirror: &mut IndexMap<String, ElementDecl>, msgs: &[ServerMsg]) {
    for msg in msgs {
        match msg {
            ServerMsg::Snapshot { elements } => {
                mirror.clear();
                for elem in elements {
                    mirror.insert(elem.id.clone(), elem.clone());
                }
            }
            ServerMsg::Add { element } => {
                mirror.insert(element.id.clone(), element.clone());
            }
            ServerMsg::Update { id, value, label, meta } => {
                if let Some(elem) = mirror.get_mut(id) {
                    elem.value = value.clone();
                    if let Some(l) = label {
                        elem.label = l.clone();
                    }
                    if let Some(m) = meta {
                        elem.meta = m.clone();
                    }
                }
            }
            ServerMsg::Remove { id } => {
                mirror.shift_remove(id);
            }
            ServerMsg::Reorder { ids, .. } => {
                let mut new_mirror = IndexMap::with_capacity(mirror.len());
                for id in ids {
                    if let Some(elem) = mirror.get(id) {
                        new_mirror.insert(id.clone(), elem.clone());
                    }
                }
                // Keep any elements not in the reorder list at the end
                for (id, elem) in mirror.iter() {
                    if !new_mirror.contains_key(id) {
                        new_mirror.insert(id.clone(), elem.clone());
                    }
                }
                *mirror = new_mirror;
            }
        }
    }
}

/// Find an available port pair (http, ws) starting from `start`.
/// HTTP gets the even port, WS gets the odd port.
pub fn find_port_pair(start: u16, bind_addr: &str) -> (u16, u16) {
    for base in (start..start + 100).step_by(2) {
        let http_ok = TcpListener::bind(format!("{bind_addr}:{base}")).is_ok();
        let ws_ok = TcpListener::bind(format!("{bind_addr}:{}", base + 1)).is_ok();
        if http_ok && ws_ok {
            return (base, base + 1);
        }
    }
    panic!("wgui: could not find available port pair starting from {start}");
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::element::{ElementDecl, ElementKind, ElementMeta, Value};
    use std::sync::Arc;

    fn make_decl(id: &str, value: Value) -> ElementDecl {
        ElementDecl {
            id: id.to_string(),
            kind: ElementKind::Label,
            label: id.to_string(),
            value,
            meta: ElementMeta::default(),
            window: Arc::from("test"),
        }
    }

    fn into_mirror(decls: Vec<ElementDecl>) -> IndexMap<String, ElementDecl> {
        decls.into_iter().map(|d| (d.id.clone(), d)).collect()
    }

    #[test]
    fn mirror_add() {
        let mut mirror = IndexMap::new();
        let msgs = vec![ServerMsg::Add {
            element: make_decl("a", Value::Bool(true)),
        }];
        apply_messages_to_mirror(&mut mirror, &msgs);
        assert_eq!(mirror.len(), 1);
        assert_eq!(mirror["a"].id, "a");
    }

    #[test]
    fn mirror_update() {
        let mut mirror = into_mirror(vec![make_decl("a", Value::Bool(true))]);
        let msgs = vec![ServerMsg::Update {
            id: "a".to_string(),
            value: Value::Bool(false),
            label: None,
            meta: None,
        }];
        apply_messages_to_mirror(&mut mirror, &msgs);
        assert_eq!(mirror["a"].value, Value::Bool(false));
    }

    #[test]
    fn mirror_remove() {
        let mut mirror = into_mirror(vec![make_decl("a", Value::Bool(true))]);
        let msgs = vec![ServerMsg::Remove {
            id: "a".to_string(),
        }];
        apply_messages_to_mirror(&mut mirror, &msgs);
        assert!(mirror.is_empty());
    }

    #[test]
    fn mirror_snapshot() {
        let mut mirror = into_mirror(vec![make_decl("a", Value::Bool(true))]);
        let msgs = vec![ServerMsg::Snapshot {
            elements: vec![make_decl("b", Value::Bool(false))],
        }];
        apply_messages_to_mirror(&mut mirror, &msgs);
        assert_eq!(mirror.len(), 1);
        assert_eq!(mirror["b"].id, "b");
    }

    #[test]
    fn html_escape_special_chars() {
        assert_eq!(html_escape("<script>"), "&lt;script&gt;");
        assert_eq!(html_escape("a&b"), "a&amp;b");
        assert_eq!(html_escape("\"hi\""), "&quot;hi&quot;");
        assert_eq!(html_escape("plain"), "plain");
    }
}