twgame 0.11.0

use crate::ids::TeamId;
use indexmap::IndexMap;
use rand::RngCore;
use rand_pcg::Lcg64Xsh32;
use std::cell::RefCell;
use std::collections::HashMap;
use std::rc::Rc;
use twgame_core::twsnap::time::{Duration, Instant, SnapTick};

type PrngCompat = HashMap<Instant, HashMap<TeamId, Vec<u32>>>;

#[derive(Debug)]
struct PrngInner {
    prng: Lcg64Xsh32,
    /// keep track of current tick compat data to only keep a record for ticks that have multiple
    /// teams accessing the Prng
    /// keep track of last use information to generate compat data
    // TODO: can I optimize the memory usage?
    uses: Option<(Instant, HashMap<TeamId, Vec<u32>>, PrngCompat)>,
    /// compat data from external source this only includes ticks where at least two teams access
    /// the rng source.
    external_info: Option<HashMap<Instant, HashMap<TeamId, Vec<u32>>>>,
}

impl PrngInner {
    fn random_bits(&mut self, team_id: TeamId, now: Instant, num: u32) -> u32 {
        // generate new random value
        let mut rand = self.prng.next_u32();

        if let Some(compat_data) = self.external_info.as_ref() {
            // check if we have compat data for this tick
            if let Some(cur_tick) = compat_data.get(&now) {
                // then we should have compat data for this team and num, panic for now if it doesn't
                // exist
                rand = cur_tick[&team_id][num as usize];
            }
        }
        if let Some((last_tick, last_compat, stored_compat_data)) = self.uses.as_mut() {
            if *last_tick != now {
                if last_compat.len() > 1 {
                    stored_compat_data.insert(*last_tick, last_compat.clone());
                }
                last_compat.clear();
            }
            *last_tick = now;
            let team_values = last_compat.entry(team_id).or_default();
            assert_eq!(num as usize, team_values.len());
            team_values.push(rand);
        }

        rand
    }

    /// None, if compat data already retrieved or game not set to collect compat data
    /// Empty if no compat needed to replay this with TwGame mode
    fn retrieve_compat_data(&mut self) -> Option<String> {
        if let Some((last_tick, last_compat, mut stored_compat_data)) = self.uses.take() {
            if last_compat.len() > 1 {
                stored_compat_data.insert(last_tick, last_compat);
            }
            if stored_compat_data.is_empty() {
                return Some(String::new());
            }

            // convert with stable output
            let mut compat_data: IndexMap<SnapTick, IndexMap<i32, Vec<u32>>> = IndexMap::new();
            let mut ticks: Vec<_> = stored_compat_data.keys().cloned().collect();
            ticks.sort();
            for tick in ticks {
                compat_data.insert(tick.snap_tick(), IndexMap::new());
                let mut team_ids: Vec<_> = stored_compat_data[&tick].keys().cloned().collect();
                team_ids.sort();
                for team_id in team_ids {
                    compat_data[&tick.snap_tick()].insert(
                        team_id.to_i32(),
                        stored_compat_data[&tick][&team_id].clone(),
                    );
                }
            }
            // serialize with intermediate struct
            Some(serde_json::to_string(&compat_data).unwrap())
        } else {
            None
        }
    }
}

/// Shared prng between team implementations for ddnet compatibility
#[derive(Debug)]
pub struct Prng {
    /// Allow sharing of Prng between multiple `GameState`s
    inner: Rc<RefCell<PrngInner>>,
    /// separate compat data by TeamId
    team_id: TeamId,
    /// store last_used to be able to store and use compat data
    last_used: Instant,
    /// number of uses within this tick for this team
    num_used: u32,
}

impl Prng {
    pub(super) fn new(team_id: TeamId, prng: Lcg64Xsh32) -> Self {
        Self {
            inner: Rc::new(RefCell::new(PrngInner {
                prng,
                uses: None,
                external_info: None,
            })),
            team_id,
            last_used: Default::default(),
            num_used: 0,
        }
    }

    /// Creates new Prng pointing to the same source
    pub(super) fn create_new(&self, team_id: TeamId) -> Self {
        Self {
            inner: self.inner.clone(),
            team_id,
            last_used: Instant::zero(),
            num_used: 0,
        }
    }

    /// Overwrite Prng description for configuration from Teehistorian header.
    pub(super) fn overwrite(&self, prng_description: Lcg64Xsh32) {
        self.inner.borrow_mut().prng = prng_description;
    }

    pub(super) fn supply_compat_data(&self, compat_data: &str) {
        if compat_data.is_empty() {
            self.inner.borrow_mut().external_info = Some(HashMap::new());
            return;
        }
        // deserialize with intermediate struct
        let compat_data: HashMap<SnapTick, HashMap<i32, Vec<u32>>> =
            serde_json::from_str(compat_data).unwrap();
        // convert to PrngCompat
        let compat_data: PrngCompat = compat_data
            .into_iter()
            .map(|(tick, teams)| {
                (
                    Instant::zero() + Duration::from_ticks(tick),
                    teams
                        .into_iter()
                        .map(|(team_id, values)| (TeamId::from_i32(team_id), values))
                        .collect(),
                )
            })
            .collect();
        self.inner.borrow_mut().external_info = Some(compat_data);
    }

    pub(super) fn enable_compat_data_collection(&self) {
        self.inner.borrow_mut().uses = Some((Instant::zero(), HashMap::new(), HashMap::new()));
    }

    pub(super) fn retrieve_compat_data(&self) -> Option<String> {
        self.inner.borrow_mut().retrieve_compat_data()
    }
}

impl Prng {
    // DDNet `CPrng::RandomBits`
    pub fn random_bits(&mut self, now: Instant) -> u32 {
        if self.last_used == now {
            self.num_used += 1;
        } else {
            self.last_used = now;
            self.num_used = 0;
        }
        self.inner
            .borrow_mut()
            .random_bits(self.team_id, now, self.num_used)
    }

    /// Only advances the Prng if a number >1 is passed to `below_this`.
    ///
    /// DDNet `CWorldCore::RandomOr0`.
    ///
    /// From DDNet source:
    /// > This makes the random number slightly biased if `BelowThis`
    /// > is not a power of two, but we have decided that this is not
    /// > significant for DDNet and favored the simple implementation.
    pub fn random_or_0(&mut self, now: Instant, below_this: u32) -> u32 {
        if below_this <= 1 {
            0
        } else {
            self.random_bits(now) % below_this
        }
    }
}