//! The host side: a WebSocket listener on its own threads. Each connection
//! proves it has the key, then trades JSON messages with the workspace,
//! which answers on the UI thread through [`ServerEvent`]s.

use std::{
    collections::HashMap,
    io::ErrorKind,
    net::{TcpListener, TcpStream},
    sync::{
        Arc, Mutex,
        atomic::{AtomicBool, AtomicU64, Ordering},
        mpsc,
    },
    thread,
    time::Duration,
};

use tungstenite::{Message, WebSocket};

use super::{ClientMessage, PROTOCOL, ServerMessage, keys_match};

/// How long a connection may take to say hello.
const HELLO_TIMEOUT: Duration = Duration::from_secs(10);
/// How often a connection's thread looks for messages to send.
const POLL: Duration = Duration::from_millis(40);

pub enum ServerEvent {
    Connected { conn: u64, name: String, address: String },
    Message { conn: u64, message: ClientMessage },
    Disconnected { conn: u64 },
}

/// A client that's connected, for the settings to list.
#[derive(Clone, Debug)]
pub struct Peer {
    pub conn: u64,
    pub name: String,
    pub address: String,
}

pub struct RemoteServer {
    pub port: u16,
    stop: Arc<AtomicBool>,
    outboxes: Arc<Mutex<HashMap<u64, mpsc::Sender<String>>>>,
    pub events: smol::channel::Receiver<ServerEvent>,
}

impl RemoteServer {
    /// Starts listening on every interface. Connections must say `key`.
    pub fn start(port: u16, key: String, host: String) -> std::io::Result<Self> {
        let listener = TcpListener::bind(("0.0.0.0", port))?;
        listener.set_nonblocking(true)?;
        let stop = Arc::new(AtomicBool::new(false));
        let outboxes: Arc<Mutex<HashMap<u64, mpsc::Sender<String>>>> = Arc::default();
        let (events_tx, events) = smol::channel::unbounded();
        let next_conn = Arc::new(AtomicU64::new(1));
        {
            let stop = stop.clone();
            let outboxes = outboxes.clone();
            thread::Builder::new().name("remote-listener".into()).spawn(move || {
                while !stop.load(Ordering::Relaxed) {
                    match listener.accept() {
                        Ok((stream, address)) => {
                            let conn = next_conn.fetch_add(1, Ordering::Relaxed);
                            let (stop, outboxes, events, key, host) =
                                (stop.clone(), outboxes.clone(), events_tx.clone(), key.clone(), host.clone());
                            let _ = thread::Builder::new().name("remote-connection".into()).spawn(move || {
                                serve(conn, stream, address.ip().to_string(), &key, &host, &stop, &outboxes, &events);
                                outboxes.lock().unwrap_or_else(|poisoned| poisoned.into_inner()).remove(&conn);
                            });
                        }
                        Err(err) if err.kind() == ErrorKind::WouldBlock => thread::sleep(Duration::from_millis(150)),
                        Err(_) => thread::sleep(Duration::from_millis(500)),
                    }
                }
            })?;
        }
        Ok(Self { port, stop, outboxes, events })
    }

    pub fn send(&self, conn: u64, message: &ServerMessage) {
        self.sender(conn).send(message);
    }

    /// Sends to one client from any thread, for as long as it's connected.
    pub fn sender(&self, conn: u64) -> ConnectionSender {
        ConnectionSender { conn, outboxes: self.outboxes.clone() }
    }

    pub fn broadcast(&self, message: &ServerMessage) {
        let Ok(text) = serde_json::to_string(message) else { return };
        for outbox in self.outboxes.lock().unwrap_or_else(|poisoned| poisoned.into_inner()).values() {
            let _ = outbox.send(text.clone());
        }
    }

    /// Drops a client's connection.
    pub fn disconnect(&self, conn: u64) {
        // Its thread ends once its outbox is gone.
        self.outboxes.lock().unwrap_or_else(|poisoned| poisoned.into_inner()).remove(&conn);
    }
}

impl Drop for RemoteServer {
    fn drop(&mut self) {
        self.stop.store(true, Ordering::Relaxed);
    }
}

/// A way to one client's connection that works off the UI thread.
#[derive(Clone)]
pub struct ConnectionSender {
    conn: u64,
    outboxes: Arc<Mutex<HashMap<u64, mpsc::Sender<String>>>>,
}

impl ConnectionSender {
    pub fn send(&self, message: &ServerMessage) {
        let Ok(text) = serde_json::to_string(message) else { return };
        if let Some(outbox) = self.outboxes.lock().unwrap_or_else(|poisoned| poisoned.into_inner()).get(&self.conn) {
            let _ = outbox.send(text);
        }
    }
}

#[allow(clippy::too_many_arguments)]
fn serve(
    conn: u64,
    stream: TcpStream,
    address: String,
    key: &str,
    host: &str,
    stop: &AtomicBool,
    outboxes: &Mutex<HashMap<u64, mpsc::Sender<String>>>,
    events: &smol::channel::Sender<ServerEvent>,
) {
    let _ = stream.set_nonblocking(false);
    let _ = stream.set_nodelay(true);
    let _ = stream.set_read_timeout(Some(HELLO_TIMEOUT));
    // A client's Device panel streaming one of this device's simulators.
    if super::device_proxy::is_device_request(&stream) {
        super::device_proxy::serve(stream, key, || crate::device::service().hub_for_remote());
        return;
    }
    let Ok(mut socket) = tungstenite::accept(stream) else { return };

    let name = match read_message(&mut socket) {
        Some(ClientMessage::Hello { key: offered, name, protocol }) if keys_match(&offered, key) && protocol == PROTOCOL => name,
        Some(ClientMessage::Hello { protocol, .. }) if protocol != PROTOCOL => {
            deny(&mut socket, "This app's version speaks a different protocol; update both devices.");
            return;
        }
        _ => {
            // Slow down anyone guessing keys.
            thread::sleep(Duration::from_secs(1));
            deny(&mut socket, "Wrong key.");
            return;
        }
    };
    let welcome = ServerMessage::Welcome { host: host.to_string(), version: env!("CARGO_PKG_VERSION").into(), protocol: PROTOCOL };
    if socket.send(Message::text(serde_json::to_string(&welcome).unwrap_or_default())).is_err() {
        return;
    }
    let (outbox, inbox) = mpsc::channel::<String>();
    outboxes.lock().unwrap_or_else(|poisoned| poisoned.into_inner()).insert(conn, outbox);
    let _ = events.send_blocking(ServerEvent::Connected { conn, name, address });
    let _ = socket.get_mut().set_read_timeout(Some(POLL));

    'connection: while !stop.load(Ordering::Relaxed) {
        loop {
            match inbox.try_recv() {
                Ok(text) => {
                    if socket.send(Message::text(text)).is_err() {
                        break 'connection;
                    }
                }
                Err(mpsc::TryRecvError::Empty) => break,
                Err(mpsc::TryRecvError::Disconnected) => break 'connection,
            }
        }
        match socket.read() {
            Ok(Message::Text(text)) => {
                if let Ok(message) = serde_json::from_str::<ClientMessage>(text.as_str()) {
                    let _ = events.send_blocking(ServerEvent::Message { conn, message });
                }
            }
            Ok(Message::Close(_)) => break,
            Ok(_) => {}
            Err(tungstenite::Error::Io(err)) if matches!(err.kind(), ErrorKind::WouldBlock | ErrorKind::TimedOut) => {}
            Err(_) => break,
        }
    }
    let _ = socket.close(None);
    let _ = events.send_blocking(ServerEvent::Disconnected { conn });
}

fn read_message(socket: &mut WebSocket<TcpStream>) -> Option<ClientMessage> {
    loop {
        match socket.read().ok()? {
            Message::Text(text) => return serde_json::from_str(text.as_str()).ok(),
            Message::Ping(_) | Message::Pong(_) => continue,
            _ => return None,
        }
    }
}

fn deny(socket: &mut WebSocket<TcpStream>, reason: &str) {
    let denied = ServerMessage::Denied { reason: reason.into() };
    let _ = socket.send(Message::text(serde_json::to_string(&denied).unwrap_or_default()));
    let _ = socket.close(None);
}