werewolves/werewolves-server/src/lobby.rs

use core::ops::{Deref, DerefMut};

use tokio::sync::broadcast::Sender;
use werewolves_proto::{
    error::GameError,
    game::{Game, GameSettings},
    message::{
        ClientMessage, Identification, PlayerState, ServerMessage,
        host::{HostLobbyMessage, HostMessage, ServerToHostMessage},
    },
    player::PlayerId,
};

use crate::{
    LogError,
    communication::lobby::LobbyComms,
    connection::JoinedPlayers,
    game::GameRunner,
    runner::{IdentifiedClientMessage, Message},
};

pub struct Lobby {
    players_in_lobby: LobbyPlayers,
    settings: GameSettings,
    joined_players: JoinedPlayers,
    comms: Option<LobbyComms>,
}

impl Lobby {
    pub fn new(joined_players: JoinedPlayers, comms: LobbyComms) -> Self {
        Self {
            joined_players,
            comms: Some(comms),
            settings: GameSettings::default(),
            players_in_lobby: LobbyPlayers(Vec::new()),
        }
    }

    pub fn set_players_in_lobby(&mut self, players_in_lobby: LobbyPlayers) {
        self.players_in_lobby = players_in_lobby
    }

    const fn comms(&mut self) -> Result<&mut LobbyComms, GameError> {
        match self.comms.as_mut() {
            Some(comms) => Ok(comms),
            None => Err(GameError::InactiveGameObject),
        }
    }

    pub async fn send_lobby_info_to_clients(&mut self) {
        let players = self
            .players_in_lobby
            .iter()
            .map(|(id, _)| id.public.clone())
            .collect::<Box<[_]>>();
        self.joined_players
            .send_all_lobby(
                players,
                &self
                    .players_in_lobby
                    .iter()
                    .map(|(id, _)| id.player_id.clone())
                    .collect::<Box<[_]>>(),
            )
            .await;
    }

    async fn get_lobby_player_list(&self) -> Box<[PlayerState]> {
        let mut players = Vec::new();
        for (player, _) in self.players_in_lobby.iter() {
            players.push(PlayerState {
                identification: player.clone(),
                connected: self.joined_players.is_connected(&player.player_id).await,
            });
        }

        players.into_boxed_slice()
    }

    pub async fn send_lobby_info_to_host(&mut self) -> Result<(), GameError> {
        let players = self.get_lobby_player_list().await;
        self.comms()?
            .host()
            .send(ServerToHostMessage::Lobby(players))
            .map_err(|err| GameError::GenericError(err.to_string()))
    }

    pub async fn next(&mut self) -> Option<GameRunner> {
        let msg = self
            .comms()
            .unwrap()
            .next_message()
            .await
            .expect("get next message");

        match self.next_inner(msg.clone()).await.map_err(|err| (msg, err)) {
            Ok(None) => {}
            Ok(Some(mut game)) => {
                game.role_reveal().await;
                match game.host_message(HostMessage::GetState) {
                    Ok(msg) => {
                        log::info!("initial message after night reveal: {msg:?}");
                        game.comms().host().send(msg).log_warn();
                    }
                    Err(err) => {
                        log::error!("processing get_state after role reveal to host: {err}")
                    }
                }

                return Some(game);
            }
            Err((Message::Host(_), err)) => self
                .comms()
                .unwrap()
                .host()
                .send(ServerToHostMessage::Error(err))
                .log_warn(),
            Err((
                Message::Client(IdentifiedClientMessage {
                    identity:
                        Identification {
                            player_id,
                            public: _,
                        },
                    message: _,
                }),
                GameError::InvalidMessageForGameState,
            )) => {
                let _ = self
                    .players_in_lobby
                    .send_if_present(&player_id, ServerMessage::InvalidMessageForGameState);
            }
            Err((
                Message::Client(IdentifiedClientMessage {
                    identity: Identification { player_id, public },
                    message: _,
                }),
                err,
            )) => {
                log::error!("processing message from {public} [{player_id}]: {err}");
                let _ = self
                    .players_in_lobby
                    .send_if_present(&player_id, ServerMessage::Reset);
            }
            Err((Message::Connect(_), _)) | Err((Message::Disconnect(_), _)) => {}
        }
        None
    }

    async fn next_inner(&mut self, msg: Message) -> Result<Option<GameRunner>, GameError> {
        match msg {
            Message::Host(HostMessage::InGame(_))
            | Message::Host(HostMessage::ForceRoleAckFor(_)) => {
                return Err(GameError::InvalidMessageForGameState);
            }
            Message::Host(HostMessage::NewLobby) => self
                .comms()
                .unwrap()
                .host()
                .send(ServerToHostMessage::Error(
                    GameError::InvalidMessageForGameState,
                ))
                .log_warn(),
            Message::Host(HostMessage::Lobby(HostLobbyMessage::GetState))
            | Message::Host(HostMessage::GetState) => self.send_lobby_info_to_host().await?,
            Message::Host(HostMessage::Lobby(HostLobbyMessage::GetGameSettings)) => {
                let msg = ServerToHostMessage::GameSettings(self.settings.clone());
                let _ = self.comms().unwrap().host().send(msg);
            }
            Message::Host(HostMessage::Lobby(HostLobbyMessage::SetGameSettings(settings))) => {
                settings.check()?;
                self.settings = settings;
            }
            Message::Host(HostMessage::Lobby(HostLobbyMessage::SetPlayerNumber(pid, num))) => {
                self.joined_players
                    .update(&pid, |p| p.number = Some(num))
                    .await;
                if let Some(p) = self
                    .players_in_lobby
                    .iter_mut()
                    .find_map(|(c, _)| (c.player_id == pid).then_some(c))
                    && let Some(joined_id) = self.joined_players.get_player_identity(&pid).await
                {
                    p.public = joined_id;
                }
                self.send_lobby_info_to_clients().await;
                self.send_lobby_info_to_host().await.log_debug();
            }
            Message::Host(HostMessage::Lobby(HostLobbyMessage::Start)) => {
                if self.players_in_lobby.len() < self.settings.min_players_needed() {
                    return Err(GameError::TooFewPlayers {
                        got: self.players_in_lobby.len() as _,
                        need: self.settings.min_players_needed() as _,
                    });
                }
                let playing_players = self
                    .players_in_lobby
                    .iter()
                    .map(|(id, _)| id.clone())
                    .collect::<Box<[_]>>();

                let game = Game::new(&playing_players, self.settings.clone())?;
                assert_eq!(game.village().characters().len(), playing_players.len());

                let (comms, recv) = self.comms.take().unwrap().into_inner();
                return Ok(Some(GameRunner::new(
                    game,
                    comms,
                    self.players_in_lobby.clone(),
                    recv,
                    self.joined_players.clone(),
                )));
            }
            Message::Client(IdentifiedClientMessage {
                identity,
                message: ClientMessage::Hello,
            }) => {
                if self
                    .players_in_lobby
                    .iter_mut()
                    .any(|p| p.0.player_id == identity.player_id)
                {
                    // Already have the player
                    return Ok(None);
                }
                if let Some(sender) = self.joined_players.get_sender(&identity.player_id).await {
                    self.players_in_lobby.push((identity, sender.clone()));
                    self.send_lobby_info_to_clients().await;
                    self.send_lobby_info_to_host().await?;
                }
            }
            Message::Host(HostMessage::Lobby(HostLobbyMessage::Kick(player_id)))
            | Message::Client(IdentifiedClientMessage {
                identity:
                    Identification {
                        player_id,
                        public: _,
                    },
                message: ClientMessage::Goodbye,
            }) => {
                log::error!("we are in there");
                if let Some(remove_idx) = self
                    .players_in_lobby
                    .iter()
                    .enumerate()
                    .find_map(|(idx, p)| (p.0.player_id == player_id).then_some(idx))
                {
                    log::error!("removing player {player_id} at idx {remove_idx}");
                    self.players_in_lobby.swap_remove(remove_idx);
                    self.send_lobby_info_to_host().await?;
                    self.send_lobby_info_to_clients().await;
                }
            }
            Message::Client(IdentifiedClientMessage {
                identity:
                    Identification {
                        player_id,
                        public: _,
                    },
                message: ClientMessage::GetState,
            }) => {
                let msg = ServerMessage::LobbyInfo {
                    joined: self
                        .players_in_lobby
                        .iter()
                        .any(|(p, _)| p.player_id == player_id),
                    players: self
                        .players_in_lobby
                        .iter()
                        .map(|(id, _)| id.public.clone())
                        .collect(),
                };
                if let Some(sender) = self.joined_players.get_sender(&player_id).await {
                    sender.send(msg).log_debug();
                }
            }
            Message::Client(IdentifiedClientMessage {
                identity: _,
                message: ClientMessage::RoleAck,
            }) => return Err(GameError::InvalidMessageForGameState),
            Message::Client(IdentifiedClientMessage {
                identity: Identification { player_id, public },
                message: ClientMessage::UpdateSelf(_),
            }) => {
                self.joined_players
                    .update(&player_id, |p| {
                        p.name = public.name.clone();
                        p.number = public.number;
                        p.pronouns = public.pronouns.clone();
                    })
                    .await;
                if let Some(p) = self
                    .players_in_lobby
                    .iter_mut()
                    .find_map(|(c, _)| (c.player_id == player_id).then_some(c))
                {
                    p.public = public;
                }
                self.send_lobby_info_to_clients().await;
                self.send_lobby_info_to_host().await.log_debug();
            }
            Message::Connect(_) | Message::Disconnect(_) => self.send_lobby_info_to_host().await?,
            Message::Host(HostMessage::Echo(msg)) => {
                self.comms()?.host().send(msg).log_warn();
            }
        }

        Ok(None)
    }
}

#[derive(Clone)]
pub struct LobbyPlayers(Vec<(Identification, Sender<ServerMessage>)>);

impl Deref for LobbyPlayers {
    type Target = Vec<(Identification, Sender<ServerMessage>)>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for LobbyPlayers {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

impl LobbyPlayers {
    pub fn find(&self, player_id: &PlayerId) -> Option<&Sender<ServerMessage>> {
        self.iter()
            .find_map(|(id, s)| (&id.player_id == player_id).then_some(s))
    }

    pub fn send_if_present(
        &self,
        player_id: &PlayerId,
        message: ServerMessage,
    ) -> Result<(), GameError> {
        if let Some(sender) = self.find(player_id) {
            sender
                .send(message)
                .map(|_| ())
                .map_err(|err| GameError::GenericError(err.to_string()))
        } else {
            Ok(())
        }
    }
}