use std::collections::BTreeMap;

use gc_arena::Collect;
use serde::Deserialize;
use wasm_bindgen::{closure::Closure, JsCast, JsValue};
use wasm_bindgen_futures::js_sys::Date;
use web_sys::{
    console,
    js_sys::{ArrayBuffer, JsString, Uint8Array},
    BinaryType, DomException, MessageEvent, WebSocket,
};

use crate::{
    logging::{self, LogEvent},
    lua_engine::muds::{handle_websocket_close_log_err, handle_websocket_output_log_err},
    GlobalMemoCell,
};

#[derive(PartialEq, PartialOrd, Eq, Ord, Debug, Clone, Collect)]
#[collect(require_static)]
pub struct WebSocketId(pub u64);
pub struct WebSocketData {
    pub connection: WebSocket,
    pub trusted: bool, // Is it allowed to submit Lua to be run?
    pub closed: bool,
    pub url: String,
    pub retained_closures: Option<(Closure<dyn FnMut(MessageEvent)>, Closure<dyn FnMut()>)>,
    pub log_dest: Option<String>,
}
pub type RegisteredWebSockets = BTreeMap<WebSocketId, WebSocketData>;

#[derive(Deserialize)]
pub enum MessageFromServer {
    RunLua(String),
    #[serde(other)]
    Unknown,
}

pub fn connect_websocket(
    trusted: bool,
    url: &str,
    sockets: &mut RegisteredWebSockets,
    // Only used for callbacks, expected sockets already borrowed mut!
    globals: &GlobalMemoCell,
) -> anyhow::Result<WebSocketId> {
    let mut data = WebSocketData {
        connection: WebSocket::new(url).map_err(|e| {
            anyhow::Error::msg(
                e.dyn_into::<DomException>()
                    .map(|e| e.message().to_string())
                    .unwrap_or("Unknown error connecting".to_owned()),
            )
        })?,
        trusted,
        closed: false,
        url: url.to_owned(),
        retained_closures: None,
        log_dest: None,
    };
    data.connection.set_binary_type(BinaryType::Arraybuffer);

    let new_id = sockets
        .last_key_value()
        .map_or(WebSocketId(0), |v| WebSocketId(v.0 .0 + 1));
    let data_globals = globals.clone();
    let data_new_id = new_id.clone();
    let data_closure: Closure<dyn FnMut(MessageEvent)> = Closure::new(move |ev: MessageEvent| {
        let data = ev.data();
        if data.has_type::<ArrayBuffer>() {
            let log_dest = data_globals
                .ws_registry
                .borrow()
                .get(&data_new_id)
                .and_then(|d| d.log_dest.clone());
            if let Some(log_dest) = &log_dest {
                let str_rendering =
                    String::from_utf8_lossy(&Uint8Array::new(&data).to_vec()).to_string();
                logging::log(
                    &data_globals,
                    &LogEvent {
                        log_stream: log_dest.clone(),
                        timestamp: Date::new_0(),
                        message: format!("R {}", &str_rendering),
                    },
                );
            }
            handle_websocket_output_log_err(
                &data_new_id,
                &data_globals,
                &Uint8Array::new(&data).to_vec(),
            );
        } else if data.has_type::<JsString>() {
            let msg: String = data.unchecked_into::<JsString>().into();
            if let Some(latest_data) = data_globals.ws_registry.borrow().get(&data_new_id) {
                if latest_data.trusted {
                    match data_globals.lua_engine.borrow_mut().execute(&msg) {
                        Ok(()) => {}
                        Err(e) => console::log_3(
                            &JsValue::from_str("Error executing Lua from websocket"),
                            &JsValue::from_str(&latest_data.url),
                            &JsValue::from_str(&e),
                        ),
                    }
                } else {
                    console::log_2(
                        &JsValue::from_str("Ignoring Lua from untrusted websocket."),
                        &JsValue::from_str(&latest_data.url),
                    );
                }
            }
        }
    });
    let close_globals = globals.clone();
    let close_id = new_id.clone();
    let close_closure: Closure<dyn FnMut()> =
        Closure::new(
            move || match close_globals.ws_registry.borrow_mut().get_mut(&close_id) {
                None => {}
                Some(closed_socket) => {
                    closed_socket.closed = true;
                    closed_socket.retained_closures = None;
                    handle_websocket_close_log_err(&close_id, &close_globals);
                }
            },
        );

    data.connection
        .add_event_listener_with_callback("message", data_closure.as_ref().unchecked_ref())
        .expect("Couldn't set message handler on WebSocket");
    data.connection
        .add_event_listener_with_callback("close", close_closure.as_ref().unchecked_ref())
        .expect("Couldn't set close handler on WebSocket");
    data.connection
        .add_event_listener_with_callback("error", close_closure.as_ref().unchecked_ref())
        .expect("Couldn't set error handler on WebSocket");
    data.retained_closures = Some((data_closure, close_closure));
    sockets.insert(new_id.clone(), data);
    Ok(new_id)
}

pub fn send_message_to_mud(
    socket: &WebSocketId,
    msg: &[u8],
    global: &GlobalMemoCell,
) -> anyhow::Result<()> {
    match global.ws_registry.borrow().get(socket) {
        None => Err(anyhow::Error::msg("MUD connection not found")),
        Some(sock_data) if sock_data.closed => Err(anyhow::Error::msg("MUD connection is closed")),
        Some(sock_data) => {
            if let Some(log_dest) = &sock_data.log_dest {
                logging::log(
                    global,
                    &LogEvent {
                        log_stream: log_dest.clone(),
                        timestamp: Date::new_0(),
                        message: format!("W {}", String::from_utf8_lossy(msg)),
                    },
                );
            }
            sock_data.connection.send_with_u8_array(msg).map_err(|e| {
                e.dyn_into::<DomException>()
                    .map(|e| anyhow::Error::msg(e.message()))
                    .unwrap_or_else(|_| anyhow::Error::msg("Unexpected exception while sending"))
            })?;
            Ok(())
        }
    }
}