werewolves/werewolves-proto/src/game/settings.rs

use super::Result;
use core::num::NonZeroU8;
use std::collections::HashMap;

use serde::{Deserialize, Serialize};

use crate::{error::GameError, role::RoleTitle};

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct GameSettings {
    roles: HashMap<RoleTitle, NonZeroU8>,
}

impl Default for GameSettings {
    fn default() -> Self {
        Self {
            roles: [
                (RoleTitle::Werewolf, NonZeroU8::new(1).unwrap()),
                // (RoleTitle::Seer, NonZeroU8::new(1).unwrap()),
                // (RoleTitle::Militia, NonZeroU8::new(1).unwrap()),
                // (RoleTitle::Guardian, NonZeroU8::new(1).unwrap()),
                // (RoleTitle::Apprentice, NonZeroU8::new(1).unwrap()),
            ]
            .into_iter()
            .collect(),
        }
    }
}

impl GameSettings {
    pub fn spread(&self) -> Box<[RoleTitle]> {
        self.roles
            .iter()
            .flat_map(|(r, c)| [*r].repeat(c.get() as _))
            .collect()
    }
    pub fn wolves_count(&self) -> usize {
        self.roles
            .iter()
            .filter_map(|(r, c)| {
                if r.wolf() {
                    Some(c.get() as usize)
                } else {
                    None
                }
            })
            .sum()
    }

    pub fn village_roles_count(&self) -> usize {
        self.roles
            .iter()
            .filter_map(|(r, c)| {
                if !r.wolf() {
                    Some(c.get() as usize)
                } else {
                    None
                }
            })
            .sum()
    }

    pub fn roles(&self) -> Box<[(RoleTitle, NonZeroU8)]> {
        self.roles.iter().map(|(r, c)| (*r, *c)).collect()
    }

    pub fn villagers_needed_for_player_count(&self, players: usize) -> Result<usize> {
        let min = self.min_players_needed();
        if min > players {
            return Err(GameError::TooFewPlayers {
                got: players as _,
                need: min as _,
            });
        }
        Ok(players - self.roles.values().map(|c| c.get() as usize).sum::<usize>())
    }

    pub fn check(&self) -> Result<()> {
        if self.wolves_count() == 0 {
            return Err(GameError::NoWolves);
        }
        if self
            .roles
            .iter()
            .any(|(r, _)| matches!(r, RoleTitle::Apprentice))
            && self.roles.iter().filter(|(r, _)| r.is_mentor()).count() == 0
        {
            return Err(GameError::NoApprenticeMentor);
        }
        Ok(())
    }

    pub fn min_players_needed(&self) -> usize {
        let (wolves, villagers) = (self.wolves_count(), self.village_roles_count());

        if wolves > villagers {
            wolves + 1 + wolves
        } else if wolves < villagers {
            wolves + villagers
        } else {
            wolves + villagers + 1
        }
    }

    pub fn add(&mut self, role: RoleTitle) -> Result<()> {
        if role == RoleTitle::Villager {
            return Err(GameError::CantAddVillagerToSettings);
        }
        match self.roles.get_mut(&role) {
            Some(count) => *count = NonZeroU8::new(count.get() + 1).unwrap(),
            None => {
                self.roles.insert(role, NonZeroU8::new(1).unwrap());
            }
        }
        Ok(())
    }

    pub fn sub(&mut self, role: RoleTitle) {
        if let Some(count) = self.roles.get_mut(&role)
            && count.get() != 1
        {
            *count = NonZeroU8::new(count.get() - 1).unwrap();
        } else {
            self.roles.remove(&role);
        }
    }
}