werewolves/werewolves-server/src/connection.rs

use core::num::NonZeroU8;
use std::{collections::HashMap, sync::Arc};

use tokio::{
    sync::{
        Mutex,
        broadcast::{Receiver, Sender},
    },
    time::Instant,
};
use werewolves_proto::{
    error::GameError,
    message::{PublicIdentity, ServerMessage},
    player::PlayerId,
};

use crate::LogError;

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct ConnectionId(PlayerId, Instant);

impl ConnectionId {
    pub fn new(player_id: PlayerId) -> Self {
        Self(player_id, Instant::now())
    }

    pub const fn player_id(&self) -> &PlayerId {
        &self.0
    }
}

#[derive(Debug)]
pub struct JoinedPlayer {
    sender: Sender<ServerMessage>,
    receiver: Receiver<ServerMessage>,
    active_connection: ConnectionId,
    in_game: bool,
    pub name: String,
    pub number: Option<NonZeroU8>,
    pub pronouns: Option<String>,
}

impl JoinedPlayer {
    pub const fn new(
        sender: Sender<ServerMessage>,
        receiver: Receiver<ServerMessage>,
        active_connection: ConnectionId,
        name: String,
        number: Option<NonZeroU8>,
        pronouns: Option<String>,
    ) -> Self {
        Self {
            name,
            number,
            sender,
            pronouns,
            receiver,
            active_connection,
            in_game: false,
        }
    }
    pub fn resubscribe_reciever(&self) -> Receiver<ServerMessage> {
        self.receiver.resubscribe()
    }
}

#[derive(Debug, Clone)]
pub struct JoinedPlayers {
    players: Arc<Mutex<HashMap<PlayerId, JoinedPlayer>>>,
    connect_state_sender: Sender<(PlayerId, bool)>,
}

impl JoinedPlayers {
    pub fn new(connect_state_sender: Sender<(PlayerId, bool)>) -> Self {
        Self {
            connect_state_sender,
            players: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    pub async fn send_all_lobby(&self, in_lobby: Box<[PublicIdentity]>, in_lobby_ids: &[PlayerId]) {
        let players: tokio::sync::MutexGuard<'_, HashMap<PlayerId, JoinedPlayer>> =
            self.players.lock().await;
        let senders = players
            .iter()
            .map(|(pid, p)| (pid.clone(), p.sender.clone()))
            .collect::<Box<[_]>>();
        core::mem::drop(players);
        for (pid, send) in senders {
            send.send(ServerMessage::LobbyInfo {
                joined: in_lobby_ids.contains(&pid),
                players: in_lobby.clone(),
            })
            .log_debug();
        }
    }

    pub async fn is_connected(&self, player_id: &PlayerId) -> bool {
        self.players.lock().await.contains_key(player_id)
    }

    pub async fn update(&self, player_id: &PlayerId, f: impl FnOnce(&mut JoinedPlayer)) {
        if let Some(p) = self
            .players
            .lock()
            .await
            .iter_mut()
            .find_map(|(pid, p)| (player_id == pid).then_some(p))
        {
            f(p)
        }
    }

    pub async fn get_player_identity(&self, player_id: &PlayerId) -> Option<PublicIdentity> {
        self.players.lock().await.iter().find_map(|(id, p)| {
            (id == player_id).then(|| PublicIdentity {
                name: p.name.clone(),
                pronouns: p.pronouns.clone(),
                number: p.number,
            })
        })
    }

    /// Disconnect the player
    ///
    /// Will not disconnect if the player is currently in a game, allowing them to reconnect
    pub async fn disconnect(&self, connection: &ConnectionId) -> Option<JoinedPlayer> {
        let mut map = self.players.lock().await;

        self.connect_state_sender
            .send((connection.0.clone(), false))
            .log_warn();

        if map
            .get(connection.player_id())
            .map(|p| p.active_connection == *connection && !p.in_game)
            .unwrap_or_default()
        {
            return map.remove(connection.player_id());
        }

        None
    }

    pub async fn start_game_with(&self, players: &[PlayerId]) -> Result<InGameToken, GameError> {
        let mut map = self.players.lock().await;

        for player in players {
            if let Some(player) = map.get_mut(player) {
                player.in_game = true;
            };
        }

        Ok(InGameToken::new(
            self.clone(),
            players.iter().cloned().collect(),
        ))
    }

    pub async fn release_from_game(&self, players: &[PlayerId]) {
        self.players
            .lock()
            .await
            .iter_mut()
            .filter(|(p, _)| players.contains(*p))
            .for_each(|(_, p)| p.in_game = false)
    }

    pub async fn get_sender(&self, player_id: &PlayerId) -> Option<Sender<ServerMessage>> {
        self.players
            .lock()
            .await
            .get(player_id)
            .map(|c| c.sender.clone())
    }

    pub async fn insert_or_replace(
        &self,
        player_id: PlayerId,
        player: JoinedPlayer,
    ) -> Receiver<ServerMessage> {
        let mut map = self.players.lock().await;

        if let Some(old) = map.insert(player_id.clone(), player) {
            let old_map_entry = unsafe { map.get_mut(&player_id).unwrap_unchecked() };
            old_map_entry.receiver = old.resubscribe_reciever();

            old.receiver
        } else {
            self.connect_state_sender
                .send((player_id.clone(), true))
                .log_warn();
            unsafe { map.get(&player_id).unwrap_unchecked() }
                .receiver
                .resubscribe()
        }
    }
}

pub struct InGameToken {
    joined_players: JoinedPlayers,
    players_in_game: Option<Box<[PlayerId]>>,
}
impl InGameToken {
    const fn new(joined_players: JoinedPlayers, players_in_game: Box<[PlayerId]>) -> Self {
        Self {
            joined_players,
            players_in_game: Some(players_in_game),
        }
    }
}
impl Drop for InGameToken {
    fn drop(&mut self) {
        let joined_players = self.joined_players.clone();
        if let Some(players) = self.players_in_game.take() {
            tokio::spawn(async move {
                let players_in_game = players;
                joined_players.release_from_game(&players_in_game).await;
            });
        }
    }
}