br-web-server 0.5.15

use crate::request::Request;
use crate::response::Response;
use crate::{Handler, HttpError};
use json::{object, JsonValue};
use std::sync::mpsc::{channel, Sender};
use std::sync::{Mutex};
use std::{io, thread};
use dashmap::DashMap;
use log::{info};

pub static USERS: std::sync::LazyLock<DashMap<String, Websocket>> = std::sync::LazyLock::new(DashMap::new);
pub static WS_NOTICE: std::sync::LazyLock<Mutex<Vec<NoticeMsg>>> = std::sync::LazyLock::new(|| Mutex::new(Vec::new()));
#[derive(Debug, Clone)]
pub struct Websocket {
    pub send: Option<Sender<Message>>,
    /// 接收
    pub key: String,
    pub user_user: String,
    pub org_org: String,
    version: String,
    request: Request,
    response: Response,
}


impl Websocket {
    #[must_use]
    pub fn http(request: Request, response: Response) -> Self {
        Self {
            send: None,
            request,
            key: String::new(),
            user_user: "".to_string(),
            org_org: "".to_string(),
            version: String::new(),
            response,
        }
    }
    pub fn new(request: Request, response: Response) -> Self {
        Self {
            send: None,
            request,
            key: String::new(),
            user_user: "".to_string(),
            org_org: String::new(),
            version: String::new(),
            response,
        }
    }
    // 发送数据
    pub fn send(&mut self, data: &JsonValue) {
        let msg = Message {
            mode: MessageMode::Server,
            message_type: MessageType::Text,
            payload: data.to_string().into_bytes(),
            text: data.to_string(),
            close: CloseCode::None,
            error: ErrorCode::None,
        };
        match self.send.clone().unwrap().send(msg) {
            Ok(()) => (),
            Err(_) => self.on_error(ErrorCode::SendingDataFailed),
        }
    }
    // 关闭连接
    pub fn close(&mut self, code: CloseCode, reason: &str) {
        let msg = Message {
            mode: MessageMode::Server,
            message_type: MessageType::Close,
            payload: reason.as_bytes().to_vec(),
            text: reason.to_string(),
            close: code,
            error: ErrorCode::None,
        };
        match self.send.clone().unwrap().send(msg) {
            Ok(()) => (),
            Err(_) => self.on_error(ErrorCode::SendingDataFailed),
        }
    }
    ///// 发送给指定对象
    //pub fn send_user(&mut self, user_user: &str, data: JsonValue) {
    //    if USERS.get(user_user).is_some() {
    //        for mut user in USERS.iter_mut() {
    //            if user.user_user == user_user {
    //                user.send(data.clone());
    //                return;
    //            }
    //        }
    //    }
    //}
    ///// 发送给指定的企业
    //pub fn send_org(&mut self, org_org: &str, data: JsonValue) {
    //    if USERS.get(org_org).is_some() {
    //        for mut user in USERS.iter_mut() {
    //            if user.org_org == org_org {
    //                user.send(data.clone());
    //                return;
    //            }
    //        }
    //    }
    //}
    /// 在线人数
    pub fn online_users(&mut self) -> usize {
        USERS.len()
    }
    pub fn handle(&mut self) -> Result<(), HttpError> {
        let (send, receive) = channel();
        self.send = Some(send);
        self.on_frame()?;
        let mut factory = (self.response.factory)(self.clone());
        USERS.insert(self.key.to_string(), self.clone());
        factory.on_open()?;
        let that = self.clone();
        let scheme = self.response.request.scheme.clone();

        let thr = thread::spawn(move || -> Result<(), HttpError> {
            loop {
                let msgs = scheme.lock().unwrap().read_ws_data();
                let msg = match msgs {
                    Ok(e) => e,
                    Err(_) => return Ok(())
                };
                match msg.message_type {
                    MessageType::TimeOut => continue,
                    _ => {
                        match that.send.clone().unwrap().send(msg) {
                            Ok(()) => continue,
                            Err(_) => return Ok(())
                        }
                    }
                }
            }
        });
        let that = self.clone();
        let scheme = self.response.request.scheme.clone();
        let key = self.key.clone();
        thread::spawn(move || -> io::Result<()> {
            let mut factory = (that.response.factory)(that.clone());
            loop {
                match receive.recv() {
                    Ok(msg) => {
                        match msg.message_type {
                            MessageType::TimeOut => continue,
                            MessageType::Close => {
                                if USERS.get(&key).is_some() {
                                    USERS.remove(&key);
                                }
                                factory.on_close(msg.close.clone(), &msg.text);
                                if let Ok(()) = scheme.lock().unwrap().write_all(&Message::send_close(CloseCode::ServerClose, "客户退出关闭")) {};
                                return Ok(());
                            }
                            MessageType::Pong => {
                                info!("接收到一个Pong: {:?} {:?} {:?}", msg.mode, msg.message_type, msg.payload);
                            }
                            MessageType::Binary | MessageType::Text => {
                                info!("接收到数据: {:?}", msg);
                                if let Ok(()) = factory.on_message(msg) {};
                            }
                            _ => {
                                info!("Client有数据: {:?} {:?} {}", msg.mode, msg.message_type, msg.text.clone());
                                return Ok(());
                            }
                        }
                    }
                    Err(_) => return Ok(()),
                }
            }
        });
        let _ = thr.join().unwrap();
        Ok(())
    }
}
impl Handler for Websocket {
    fn on_request(&mut self, _request: Request, _response: &mut Response) {}
    fn on_frame(&mut self) -> Result<(), HttpError> {
        self.key = self.request.header["sec-websocket-key"].as_str().unwrap_or("").to_string();
        self.version = self.request.header["sec-websocket-version"].as_str().unwrap_or("").to_string();
        self.response.header("Upgrade", "websocket");
        self.response.header("Connection", "Upgrade");
        let sec_websocket_accept = br_crypto::sha1::encrypt_base64(format!("{}258EAFA5-E914-47DA-95CA-C5AB0DC85B11", self.key).as_bytes());
        self.response.header("Sec-WebSocket-Accept", sec_websocket_accept.as_str());
        self.response.status(101).send()?;
        Ok(())
    }
}
#[derive(Debug, Clone)]
pub struct Message {
    pub mode: MessageMode,
    pub message_type: MessageType,
    pub payload: Vec<u8>, // 消息载荷，以字节向量形式表示
    pub text: String,
    pub close: CloseCode,
    pub error: ErrorCode,
}

impl Message {
    #[must_use]
    pub fn msg_error() -> Self {
        Message {
            mode: MessageMode::Client,
            message_type: MessageType::Error,
            payload: vec![],
            text: "长度不够".to_string(),
            close: CloseCode::None,
            error: ErrorCode::SendingDataFailed,
        }
    }
    // 解析WebSocket消息
    pub fn parse_message(data: &mut Vec<u8>) -> Message {
        println!("{data:?}");

        // 检查数据是否足够长来包含消息类型和载荷长度
        if data.len() < 2 {
            return Message {
                mode: MessageMode::Client,
                message_type: MessageType::Error,
                payload: vec![],
                text: "长度不够".to_string(),
                close: CloseCode::None,
                error: ErrorCode::SendingDataFailed,
            };
        }

        let header = data.drain(..2).collect::<Vec<u8>>();

        // 解析帧头
        let _fin = (header[0] & 0b1000_0000) != 0;
        let opcode = header[0] & 0b0000_1111;
        let masked = (header[1] & 0b1000_0000) != 0;
        let len_flag = header[1] & 0b0111_1111;
        let mut payload_data = Vec::new();
        let message_tpye = MessageType::from(opcode);
        println!("fin: {:#?} message_tpye: {:?} opcode: {} masked: {} len_flag: {}", _fin, message_tpye, opcode, masked, len_flag);
        match message_tpye {
            MessageType::Text => {
                let payload_length = match len_flag {
                    0..=125 => len_flag as usize,
                    126 => {
                        let ext = data.drain(..2).collect::<Vec<u8>>();
                        u16::from_be_bytes([ext[0], ext[1]]) as usize
                    }
                    127 => {
                        let ext = data.drain(..8).collect::<Vec<u8>>();
                        u64::from_be_bytes([
                            ext[0], ext[1], ext[2], ext[3],
                            ext[4], ext[5], ext[6], ext[7],
                        ]) as usize
                    }
                    _ => return Message {
                        mode: MessageMode::Client,
                        message_type: MessageType::Error,
                        payload: vec![],
                        text: "数据格式错误".to_string(),
                        close: CloseCode::None,
                        error: ErrorCode::SendingDataFailed,
                    }
                };
                if masked {
                    if data.len() < payload_length {
                        return Message {
                            mode: MessageMode::Client,
                            message_type: message_tpye,
                            payload: payload_data,
                            text: "继续加载".to_string(),
                            close: CloseCode::None,
                            error: ErrorCode::None,
                        };
                    }
                    let mask_key = data.drain(..4).collect::<Vec<u8>>();
                    let payload = &data[..payload_length];
                    for i in 0..payload.len() {
                        payload_data.push(payload[i] ^ mask_key[i % 4]);
                    }
                } else {
                    if data.len() < payload_length {
                        return Message {
                            mode: MessageMode::Client,
                            message_type: message_tpye,
                            payload: payload_data,
                            text: "继续加载".to_string(),
                            close: CloseCode::None,
                            error: ErrorCode::None,
                        };
                    }
                    let t = data.drain(..payload_length).collect::<Vec<u8>>();
                    payload_data.extend_from_slice(&t);
                }
                let text = unsafe { String::from_utf8_unchecked(payload_data.clone()) };
                Message {
                    mode: MessageMode::Client,
                    message_type: message_tpye,
                    payload: payload_data,
                    text: text.to_string(),
                    close: CloseCode::None,
                    error: ErrorCode::None,
                }
            }
            MessageType::Binary => Message {
                mode: MessageMode::Client,
                message_type: message_tpye,
                payload: payload_data,
                text: String::new(),
                close: CloseCode::None,
                error: ErrorCode::None,
            },
            MessageType::Continuation => Message {
                mode: MessageMode::Client,
                message_type: message_tpye,
                payload: payload_data,
                text: "继续加载".to_string(),
                close: CloseCode::None,
                error: ErrorCode::None,
            },
            MessageType::Close => Message {
                mode: MessageMode::Client,
                message_type: message_tpye,
                payload: payload_data,
                text: "客户端关闭".to_string(),
                close: CloseCode::ClientClose,
                error: ErrorCode::None,
            },
            MessageType::Ping => Message {
                mode: MessageMode::Client,
                message_type: message_tpye,
                payload: payload_data,
                text: "Ping".to_string(),
                close: CloseCode::None,
                error: ErrorCode::None,
            },
            MessageType::Pong => Message {
                mode: MessageMode::Client,
                message_type: message_tpye,
                payload: payload_data,
                text: "Pong".to_string(),
                close: CloseCode::None,
                error: ErrorCode::None,
            },
            MessageType::Error => {
                Message {
                    mode: MessageMode::Client,
                    message_type: message_tpye,
                    payload: vec![],
                    text: String::new(),
                    close: CloseCode::None,
                    error: ErrorCode::Unknown,
                }
            }
            MessageType::None => Message {
                mode: MessageMode::Client,
                message_type: message_tpye,
                payload: vec![],
                text: String::new(),
                close: CloseCode::None,
                error: ErrorCode::None,
            },
            MessageType::TimeOut => Message {
                mode: MessageMode::Client,
                message_type: message_tpye,
                payload: vec![],
                text: String::new(),
                close: CloseCode::None,
                error: ErrorCode::TimeOut,
            }
        }
    }
    pub fn send_message(&mut self) -> Vec<u8> {
        let mut frame = Vec::new();

        // 第1字节：FIN + RSV1-3 + OPCODE
        let opcode = self.clone().message_type.to_u8();
        let mut byte1 = opcode & 0x0F;
        byte1 |= 0x80; // FIN = 1

        frame.push(byte1);

        // 第2字节：MASK = 0（服务器发送），+ payload length
        let payload_len = self.payload.len();
        if payload_len < 126 {
            frame.push(payload_len as u8);
        } else if payload_len <= 65535 {
            frame.push(126);
            frame.extend_from_slice(&u16::try_from(payload_len).unwrap().to_be_bytes());
        } else {
            frame.push(127);
            frame.extend_from_slice(&(payload_len as u64).to_be_bytes());
        }
        frame.extend_from_slice(&self.payload);
        frame
    }
    #[must_use]
    pub fn send_close(code: CloseCode, reason: &str) -> Vec<u8> {
        let mut frame = Vec::new();
        frame.push(0x88);
        let payload_len = code.clone().to_u16().to_be_bytes().len() + reason.len();
        frame.push(u8::try_from(payload_len).unwrap());
        frame.extend(&code.to_u16().to_be_bytes());
        frame.extend(reason.as_bytes());
        frame
    }
}
#[derive(Debug, Clone)]
pub enum MessageType {
    /// 文本
    Text,
    Continuation,
    /// 客户端关闭
    Close,
    Binary,
    Ping,
    Pong,
    None,
    TimeOut,
    Error,
}

impl MessageType {
    #[must_use]
    pub fn from(types: u8) -> Self {
        match types {
            0x0 => Self::Continuation,
            0x1 => Self::Text,
            0x2 => Self::Binary,
            0x8 => Self::Close,
            0x9 => Self::Ping,
            0xa => Self::Pong,
            _ => Self::None,
        }
    }
    #[must_use]
    pub fn to_u8(self) -> u8 {
        match self {
            MessageType::Text => 0x1,
            MessageType::Continuation | MessageType::None | MessageType::Error | MessageType::TimeOut => 0x0,
            MessageType::Close => 0x8,
            MessageType::Binary => 0x2,
            MessageType::Ping => 0x9,
            MessageType::Pong => 0xa,
        }
    }
}
#[derive(Debug, Clone)]
pub enum CloseCode {
    /// 客户端主动关闭
    ClientClose,
    /// 服务端主动关闭
    ServerClose,
    /// 正常关闭
    NormalClosure,
    GoingAway,
    /// 协议错误
    ProtocolError,
    /// 其它错误
    Other,
    None,
}
impl CloseCode {
    #[must_use]
    pub fn from_err(_err: ErrorCode) -> CloseCode {
        CloseCode::None
    }
    #[must_use]
    pub fn str(&self) -> String {
        match self {
            CloseCode::ClientClose => "客户端主动关闭",
            CloseCode::ServerClose => "服务端主动关闭",
            CloseCode::None => "未知关闭",
            CloseCode::NormalClosure => "正常关闭",
            CloseCode::GoingAway => "对方离开",
            CloseCode::ProtocolError => "协议错误",
            CloseCode::Other => "其它错误",
        }.to_string()
    }
    #[must_use]
    pub fn to_u16(self) -> u16 {
        match self {
            CloseCode::NormalClosure => 1000,
            CloseCode::GoingAway => 1001,
            CloseCode::ProtocolError => 1002,
            CloseCode::ClientClose => 1003,
            CloseCode::ServerClose => 1004,
            CloseCode::Other => 1005,
            CloseCode::None => 1006,
        }
    }
}
#[derive(Debug, Clone, Copy)]
pub enum ErrorCode {
    /// 发送数据失败
    SendingDataFailed,
    /// Unknown request error
    Unknown,
    /// 线程异常
    ThreadException,
    /// 超时
    TimeOut,
    None,
}
#[derive(Debug, Clone)]
pub enum MessageMode {
    Client,
    Server,
}

pub struct NoticeMsg {
    /// 消息类型
    pub types: Types,
    /// 消息实体内容
    pub msg: JsonValue,
    /// 消息创建时间
    pub timestamp: i64,
    /// 监听通道名称
    pub channel: String,
    pub user: String,
    pub org: String,
}
impl NoticeMsg {
    pub fn json(&mut self) -> JsonValue {
        object! {
            type:"notice",
            channel: self.channel.clone(),
            msg: self.msg.clone(),
            timestamp: self.timestamp,
        }
    }
}

pub enum Types {
    /// 全体消息
    All,
    /// 指定用户消息
    User,
    /// 指定企业消息
    Org,
}