use std::net::TcpStream;

use tungstenite::protocol::{ WebSocket, Message as WsMessage };

use crate::message::Message;

pub struct WebsocketWrapper {
    websocket: WebSocket<TcpStream>,
}

impl WebsocketWrapper {
    pub fn new(websocket: WebSocket<TcpStream>) -> Self {
        WebsocketWrapper { websocket }
    }
    pub fn send(&mut self, msg: Message) {
        self.websocket.write_message(WsMessage::Text(msg.to_string())).unwrap();
    }
    pub fn recv(&mut self) -> Option<Message> {
        match self.websocket.read_message() {
            Err(_) | Ok(WsMessage::Close(_)) => None,
            Ok(WsMessage::Text(text)) => match Message::parse(text) {
                Ok(msg) => Some(msg),
                Err(_) => {
                    self.websocket.write_message(WsMessage::Text("ERROR: bad_format".to_string()));
                    self.recv()
                }
            }
            _ => {
                self.websocket.write_message(WsMessage::Text("ERROR: bad_command".to_string()));
                self.recv()
            }
        }
    }
}