shiritori/src/server.rs

use crate::client::{Client, ClientInfo};
use crate::database::{DatabaseCreationError, DatabaseSettings, DatabaseType};
use crate::response::MessageResponseData;
use crate::request::MessageRequest;
use crate::room::{Room, RoomSettings};

use derive_more::From;
use simple_websockets::{Event, EventHub, Message, Responder};
use std::cell::{RefCell, Ref};
use std::collections::HashMap;
use std::rc::{Rc, Weak};

#[derive(From, Debug)]
pub enum ServerCreationError {
    SimpleWebsocketsError(simple_websockets::Error),
    DatabaseCreationError(DatabaseCreationError),
}

#[derive(From, Debug)]
pub enum ServerError {
    DatabaseError(rusqlite::Error),
    RoomCreationError(RoomCreationError),
}

pub struct ServerSettings {
    pub port: u16,
    pub testing: bool,
}

impl Default for ServerSettings {
    fn default() -> Self {
        Self {
            port: 8080,
            testing: true,
        }
    }
}

pub struct Server {
    event_hub: EventHub,
    lobby: Rc<RefCell<Room>>,
    rooms: HashMap<String, Weak<RefCell<Room>>>,
    clients: HashMap<u64, RefCell<Client>>,
}

#[derive(From, Debug)]
pub enum RoomCreationError {
    DatabaseCreationError(DatabaseCreationError),
    NameConflict,
}

const LOBBY_NAME: &str = "lobby";

impl Server {
    pub fn new(settings: &ServerSettings) -> Result<Self, ServerCreationError> {
        let lobby = Rc::new(RefCell::new(Room::new(RoomSettings {
            name: LOBBY_NAME.to_string(),
            database_settings: DatabaseSettings {
                db_type: match settings.testing {
                    true => DatabaseType::InMemory,
                    false => DatabaseType::OnDisk("shiritori.sb".to_string()),
                },
            },
        })?));
        let lobby_weak = Rc::downgrade(&lobby);
        Ok(Self {
            event_hub: simple_websockets::launch(settings.port)?,
            lobby,
            rooms: {
                let mut rooms = HashMap::new();
                rooms.insert(LOBBY_NAME.to_string(), lobby_weak);
                rooms
            },
            clients: HashMap::new(),
        })
    }

    pub fn run(&mut self) {
        loop {
            match match self.event_hub.poll_event() {
                Event::Connect(client_id, responder) => {
                    let client = self.new_client(client_id, responder);
                    self.clients.insert(client_id, client);
                    let client = self.clients.get(&client_id).unwrap().borrow(); // moved, get it again
                    self.handle_connection(&client)
                }
                Event::Disconnect(client_id) => {
                    let client = self.clients.remove(&client_id).unwrap();
                    let client_ref = client.borrow();
                    self.handle_disconnection(&client_ref)
                },
                Event::Message(client_id, message) => {
                    println!("Received a message from client #{client_id}: {:?}", message);
                    let message = match self.process_message(message) {
                        Ok(message) => message,
                        Err(error) => {
                            let client = self.clients.get(&client_id).unwrap();
                            let client_ref = client.borrow();
                            client_ref.send(MessageResponseData::Error { message: error.to_string() }.into_message());
                            continue;
                        },
                    };
                    match message {
                        MessageRequest::Word { word } => {
                            let client = self.clients.get(&client_id).unwrap();
                            let client_ref = client.borrow();
                            match client_ref.room.borrow_mut().handle_query(&word, client_id) {
                                // Broadcast new words to all clients in same room
                                Ok(response) => self.announce_to_room(&client_ref.room, response),
                                // Send errors to only this client
                                Err(message) => {
                                    client_ref.send(
                                        MessageResponseData::Error {
                                            message: message.to_string(),
                                        }
                                        .into_message(),
                                    );
                                },
                            };
                            Ok(())
                        },
                        MessageRequest::ChangeRoom { name } => {
                            self.switch_rooms(client_id, name).unwrap();
                            Ok(())
                        },
                    }
                },
            } {
                Ok(()) => {}
                Err(error) => println!("{:?}", error),
            }
        }
    }

    fn new_client(&self, client_id: u64, responder: Responder) -> RefCell<Client> {
        RefCell::new(Client::new(
            ClientInfo::Ws {
                id: client_id,
                responder,
                discord_info: None,
            },
            self.lobby.clone(),
        ))
    }

    fn new_room(&mut self, name: &str) -> Result<Rc<RefCell<Room>>, RoomCreationError> {
        if self.rooms.contains_key(name) {
            return Err(RoomCreationError::NameConflict);
        }
        let room = Rc::new(RefCell::new(Room::new(RoomSettings {
            name: name.to_owned(),
            database_settings: DatabaseSettings {
                db_type: DatabaseType::InMemory,
            },
        })?));
        self.rooms.insert(name.to_owned(), Rc::downgrade(&room));
        Ok(room)
    }

    fn switch_rooms(&mut self, client_id: u64, room_name: String) -> Result<(), ServerError> {
        let room = self.rooms.get(&room_name)
            // upgrade Weak to Rc if exists
            .and_then(|weak| weak.upgrade()) // Option<Weak<RefCell<Room>>> -> Option<Option<Rc<RefCell<Room>>>>
            .map(Ok) // Option<Rc<RefCell<Room>>> -> Option<Result<Rc<RefCell<Room>>>>
            // if not exists OR failed to upgrade (value dropped), it will be None.
            // in that case, initialize a new room
            .unwrap_or_else(|| self.new_room(&room_name))?;
        let client = self.clients.get(&client_id).unwrap();
        let (old_room, room) = {
            let mut client_mut = client.borrow_mut();
            // Skip logic and return if going into same room
            if client_mut.room.borrow().name().eq(&room_name) {
                return Ok(());
            }
            client_mut.switch_rooms(room)
        };
        // Clean up old room to be dropped
        // However, lobby will never be dropped
        let old_room = if Rc::strong_count(&old_room) <= 1 {
            if !Rc::ptr_eq(&old_room, &self.lobby) {
                self.rooms.remove(old_room.borrow().name());
                println!("Removing room, {}!", old_room.borrow().name());
            };
            None
        } else {
            Some(old_room)
        };
        // broadcast reference count minus one
        // (We still have an Rc hanging around here)
        //self.broadcast_offseted_player_count(&old_room, -1);
        if let Some(old_room) = old_room {
            self.broadcast_player_count(&old_room);
        }
        self.broadcast_player_count(&room);

        self.room_welcome(&client.borrow())?;
        Ok(())
    }

    fn handle_connection(
        &self,
        client: &Ref<Client>,
    ) -> Result<(), ServerError> {
        // Debug
        println!("A client connected with id #{}", client.id());

        self.room_welcome(client)?;

        // Number of clients on Rc<RefCell<Room>> reference counter will be one more
        self.broadcast_player_count(&client.room);

        Ok(())
    }

    fn room_welcome(&self, client: &Ref<Client>) -> Result<(), ServerError> {
        // Get immutable access to room
        let room = client.room.borrow();

        // Send client greeting
        client.send(
            MessageResponseData::Greeting {
                id: client.id(),
                next_mora: room.next_mora().clone(),
            }
            .into_message(),
        );

        // Sent recent message history
        client.send(
            MessageResponseData::History {
                words: room.get_history()?,
            }
            .into_message(),
        );

        Ok(())
    }

    fn handle_disconnection(&self, client: &Ref<Client>) -> Result<(), ServerError> {
        let client_id = client.id();

        // Debug
        println!("Client #{client_id} disconnected.");

        // Number of clients on Rc<RefCell<Room>> reference counter will be one less
        self.broadcast_player_count(&client.room);

        Ok(())
    }

    fn for_client_in_room(&self, room: &Rc<RefCell<Room>>, mut closure: impl FnMut(Ref<Client>)) {
        for client in self
            .clients
            .iter()
            .filter(|(_id, client)| Rc::<RefCell<Room>>::ptr_eq(room, &client.borrow().room))
            .map(|(_id, refcell)| refcell.borrow())
        {
            closure(client);
        }
    }

    fn client_count_in_room(&self, room: &Rc<RefCell<Room>>) -> usize {
        let mut count: usize = 0;
        self.for_client_in_room(room, |_| count += 1);
        count
    }

    fn broadcast_player_count(&self, room: &Rc<RefCell<Room>>) {
        self.broadcast_offseted_player_count(room, 0);
    }

    fn broadcast_offseted_player_count(&self, room: &Rc<RefCell<Room>>, offset: i32) {
        let players = (self.client_count_in_room(room) as i32 + offset) as u64;
        println!("Broadcast player count {players} for room {}", room.borrow().name());
        let response = MessageResponseData::PlayerCount { players };
        self.announce_to_room(room, response);
    }

    fn announce_to_room(&self, room: &Rc<RefCell<Room>>, response_data: MessageResponseData) {
        let response = response_data.into_response();
        self.for_client_in_room(room, |client| {
            client.send(response.to_message());
        });
    }

    fn process_message(&self, message: Message) -> Result<MessageRequest, &str> {
        // Ignore binary messages
        let message: MessageRequest = match serde_json::from_str(&match message {
            Message::Text(message) => message,
            Message::Binary(_message) => return Err("Invalid request."),
        }) {
            Ok(message) => message,
            Err(_) => return Err("Invalid request."),
        };

        Ok(message)
    }
}