PkmnLib_rs/src/dynamic_data/models/battle_side.rs

use std::ffi::c_void;
use std::ops::Deref;
use std::sync::atomic::{AtomicBool, AtomicU8, Ordering};
use std::sync::{Arc, Weak};

use anyhow::{anyhow, Result};
use parking_lot::lock_api::RwLockReadGuard;
use parking_lot::{RawRwLock, RwLock};

use crate::dynamic_data::choices::TurnChoice;
use crate::dynamic_data::event_hooks::Event;
use crate::dynamic_data::models::battle::Battle;
use crate::dynamic_data::models::pokemon::Pokemon;
use crate::dynamic_data::script_handling::{ScriptSource, ScriptSourceData, ScriptWrapper};
use crate::dynamic_data::ScriptSet;
use crate::dynamic_data::VolatileScriptsOwner;
use crate::dynamic_data::{Script, WeakBattleReference};
use crate::{script_hook, StringKey, VecExt};

/// The data that is stored for a battle side.
#[derive(Debug)]
struct BattleSideData {
    /// The index of the side on the battle.
    index: u8,
    /// The number of Pokemon that can be on the side.
    pokemon_per_side: u8,
    /// A list of pokemon currently on the battlefield.
    pokemon: RwLock<Vec<Option<Pokemon>>>,
    /// The currently set choices for all Pokemon on the battlefield. Cleared when the turn starts.
    choices: RwLock<Vec<Option<Arc<TurnChoice>>>>,
    /// The slots on the side that can still be filled. Once all slots are set to false, this side
    /// has lost the battle.
    fillable_slots: Vec<AtomicBool>,
    /// The number of choices that are set.
    choices_set: AtomicU8,
    /// A reference to the battle we're part of.
    battle: WeakBattleReference,
    /// Whether or not this side has fled.
    has_fled_battle: AtomicBool,
    /// The volatile scripts that are attached to the side.
    volatile_scripts: Arc<ScriptSet>,

    /// Data required for this to be a script source.
    script_source_data: RwLock<ScriptSourceData>,
}

/// A side on a battle.
#[derive(Debug, Clone)]
pub struct BattleSide {
    /// The data that is stored for this side.
    data: Arc<BattleSideData>,
}

/// A non owning reference to a battle side.
#[derive(Debug, Clone)]
pub struct WeakBattleSideReference {
    /// A weak reference to the data of the battle side.
    data: Weak<BattleSideData>,
}

impl BattleSide {
    /// Instantiates a battle side.
    pub fn new(index: u8, pokemon_per_side: u8) -> Self {
        let mut pokemon = Vec::with_capacity(pokemon_per_side as usize);
        let mut choices = Vec::with_capacity(pokemon_per_side as usize);
        let mut fillable_slots = Vec::with_capacity(pokemon_per_side as usize);

        for _i in 0..pokemon_per_side {
            pokemon.push(None);
            choices.push(None);
            fillable_slots.push(AtomicBool::new(false));
        }
        let choices = RwLock::new(choices);
        let pokemon = RwLock::new(pokemon);

        Self {
            data: Arc::new(BattleSideData {
                index,
                pokemon_per_side,
                pokemon,
                choices,
                fillable_slots,
                choices_set: AtomicU8::new(0),
                battle: WeakBattleReference::default(),
                has_fled_battle: AtomicBool::new(false),
                volatile_scripts: Default::default(),
                script_source_data: Default::default(),
            }),
        }
    }

    /// Set the battle this side belongs to.
    pub(crate) fn set_battle(&self, battle: WeakBattleReference) {
        #[allow(clippy::unwrap_used)] // Can only be Some()
        unsafe {
            (self.data.deref() as *const BattleSideData as *mut BattleSideData)
                .as_mut()
                .unwrap()
                .battle = battle;
        }
    }

    /// The index of the side on the battle.
    pub fn index(&self) -> u8 {
        self.data.index
    }
    /// The number of Pokemon that can be on the side.
    pub fn pokemon_per_side(&self) -> u8 {
        self.data.pokemon_per_side
    }
    /// A list of pokemon currently on the battlefield.
    pub fn pokemon(&self) -> RwLockReadGuard<'_, RawRwLock, Vec<Option<Pokemon>>> {
        self.data.pokemon.read()
    }
    /// The currently set choices for all Pokemon on the battlefield. Cleared when the turn starts.
    pub fn choices(&self) -> &RwLock<Vec<Option<Arc<TurnChoice>>>> {
        &self.data.choices
    }
    /// The slots on the side that can still be filled. Once all slots are set to false, this side
    /// has lost the battle.
    pub fn fillable_slots(&self) -> &Vec<AtomicBool> {
        &self.data.fillable_slots
    }
    /// The number of choices that are set.
    pub fn choices_set(&self) -> u8 {
        self.data.choices_set.load(Ordering::SeqCst)
    }
    /// A reference to the battle we're part of.
    pub fn battle(&self) -> Result<Battle> {
        self.data
            .battle
            .upgrade()
            .ok_or(anyhow!("Battle was not set, but requested"))
    }
    /// Whether or not this side has fled.
    pub fn has_fled_battle(&self) -> bool {
        self.data.has_fled_battle.load(Ordering::SeqCst)
    }
    /// The volatile scripts that are attached to the side.
    pub fn volatile_scripts(&self) -> &Arc<ScriptSet> {
        &self.data.volatile_scripts
    }

    /// Whether every Pokemon on this side has its choices  
    pub fn all_choices_set(&self) -> bool {
        self.choices_set() == self.data.pokemon_per_side
    }

    /// Returns true if there are slots that need to be filled with a new pokemon, that have parties
    /// responsible for them. Returns false if all slots are filled with usable pokemon, or slots are
    /// empty, but can't be filled by any party anymore.
    pub fn all_slots_filled(&self) -> Result<bool> {
        for (i, pokemon) in self.data.pokemon.read().iter().enumerate() {
            match pokemon {
                Some(pokemon) => {
                    if !pokemon.is_usable() && self.battle()?.can_slot_be_filled(self.data.index, i as u8) {
                        return Ok(false);
                    }
                }
                None => {
                    if self.battle()?.can_slot_be_filled(self.data.index, i as u8) {
                        return Ok(false);
                    }
                }
            }
        }
        Ok(true)
    }

    /// Sets a choice for a Pokemon on this side.
    pub(crate) fn set_choice(&self, choice: Arc<TurnChoice>) -> Result<()> {
        for (index, pokemon_slot) in self.data.pokemon.read().iter().enumerate() {
            if let Some(pokemon) = pokemon_slot {
                if Pokemon::eq(pokemon, choice.user()) {
                    self.data.choices.write().get_mut_res(index)?.replace(choice);
                    self.data.choices_set.fetch_add(1, Ordering::SeqCst);
                    return Ok(());
                }
            }
        }
        Ok(())
    }

    /// Resets all choices on this side.
    pub fn reset_choices(&self) {
        let len = self.data.choices.read().len();
        for i in 0..len {
            self.data.choices.write().get_mut(i).take();
        }
    }

    /// Forcibly removes a Pokemon from the field.
    pub fn force_clear_pokemon(&mut self, index: u8) {
        self.data.pokemon.write().get_mut(index as usize).take();
    }

    /// Switches out a spot on the field for a different Pokemon.
    pub fn set_pokemon(&self, index: u8, pokemon: Option<Pokemon>) -> Result<()> {
        {
            let mut write_lock = self.data.pokemon.write();
            let old = write_lock.get_mut_res(index as usize)?;
            let old = match pokemon {
                Some(pokemon) => old.replace(pokemon),
                None => old.take(),
            };
            if let Some(old_pokemon) = old {
                // We need to drop the lock before calling the hook, as it might try to access the
                // side again.
                drop(write_lock);
                script_hook!(on_remove, old_pokemon,);
                old_pokemon.set_on_battlefield(false)?;
            }
        }

        let pokemon = {
            let read_lock = self.data.pokemon.read();
            &read_lock.get_res(index as usize)?.clone()
        };
        if let Some(pokemon) = pokemon {
            pokemon.set_battle_data(self.data.battle.clone(), self.data.index);
            pokemon.set_on_battlefield(true)?;
            pokemon.set_battle_index(index);

            let battle = self.battle()?;
            for side in battle.sides() {
                if side.index() == self.data.index {
                    continue;
                }
                for opponent in side.pokemon().iter().flatten() {
                    opponent.mark_opponent_as_seen(pokemon.weak());
                    pokemon.mark_opponent_as_seen(opponent.weak());
                }
            }
            battle.event_hook().trigger(Event::Switch {
                side_index: self.data.index,
                index,
                pokemon: Some(pokemon.clone()),
            });
            script_hook!(on_switch_in, pokemon, pokemon);
        } else {
            self.battle()?.event_hook().trigger(Event::Switch {
                side_index: self.data.index,
                index,
                pokemon: None,
            });
        }
        Ok(())
    }

    /// Checks whether a Pokemon is on the field in this side.
    pub fn is_pokemon_on_side(&self, pokemon: Pokemon) -> bool {
        for p in self.data.pokemon.read().iter().flatten() {
            if Pokemon::eq(p, &pokemon) {
                return true;
            }
        }
        false
    }

    /// Marks a slot as unfillable. This happens when no parties are able to fill the slot anymore.
    /// If this happens, the slot can not be used again.
    pub(crate) fn mark_slot_as_unfillable(&self, index: u8) -> Result<()> {
        self.data
            .fillable_slots
            .get_res(index as usize)?
            .store(false, Ordering::SeqCst);
        Ok(())
    }

    /// Checks whether a slot is unfillable or not.
    pub fn is_slot_unfillable(&self, pokemon: Pokemon) -> Result<bool> {
        for (i, slot) in self.data.pokemon.read().iter().enumerate() {
            if let Some(p) = slot {
                if Pokemon::eq(p, &pokemon) {
                    return Ok(self.data.fillable_slots.get_res(i)?.load(Ordering::Relaxed));
                }
            }
        }
        Ok(false)
    }

    /// Checks whether the side has been defeated.
    pub fn is_defeated(&self) -> bool {
        for fillable_slot in &self.data.fillable_slots {
            if fillable_slot.load(Ordering::Relaxed) {
                return false;
            }
        }
        true
    }

    /// Mark the side as fled.
    pub fn mark_as_fled(&mut self) {
        self.data.has_fled_battle.store(true, Ordering::SeqCst);
    }

    /// Gets a random Pokemon on the given side.
    pub fn get_random_creature_index(&self) -> Result<u8> {
        // TODO: Consider adding parameter to only get index for available creatures.
        Ok(self.battle()?.random().get_max(self.data.pokemon_per_side as i32)? as u8)
    }

    /// Swap two Pokemon on a single side around.
    pub fn swap_positions(&mut self, a: u8, b: u8) -> Result<bool> {
        let data = &self.data;
        // If out of range, don't allow swapping.
        if a >= data.pokemon_per_side || b >= data.pokemon_per_side {
            return Ok(false);
        }
        // If the two indices are the same, don't allow swapping.
        if a == b {
            return Ok(false);
        }

        // Fetch parties for the two indices.
        let mut party_a = None;
        let mut party_b = None;
        let battle = self.battle()?;
        for party in battle.parties() {
            if party.is_responsible_for_index(data.index, a) {
                party_a = Some(party);
            }
            if party.is_responsible_for_index(data.index, b) {
                party_b = Some(party);
            }
        }
        // If either of the parties does not exist, fail.
        let party_a = match party_a {
            Some(party) => party,
            None => return Ok(false),
        };
        let party_b = match party_b {
            Some(party) => party,
            None => return Ok(false),
        };
        // Don't allow swapping if different parties are responsible for the indices.
        if !Arc::ptr_eq(party_a, party_b) {
            return Ok(false);
        }

        data.pokemon.write().swap(a as usize, b as usize);
        self.battle()?.event_hook().trigger(Event::Swap {
            side_index: data.index,
            index_a: a,
            index_b: b,
        });
        Ok(true)
    }

    /// Gets a weak reference to the side.
    pub fn weak(&self) -> WeakBattleSideReference {
        WeakBattleSideReference {
            data: Arc::downgrade(&self.data),
        }
    }

    pub fn as_ptr(&self) -> *const c_void {
        Arc::as_ptr(&self.data) as *const c_void
    }
}

impl PartialEq for BattleSide {
    fn eq(&self, other: &Self) -> bool {
        Arc::ptr_eq(&self.data, &other.data)
    }
}

impl WeakBattleSideReference {
    /// Upgrades the weak reference to a strong reference, returning `None` if the side has been
    /// dropped.
    pub fn upgrade(&self) -> Option<BattleSide> {
        self.data.upgrade().map(|data| BattleSide { data })
    }

    /// Gets the underlying pointer to the data of the side.
    pub(crate) fn as_ptr(&self) -> *const c_void {
        self.data.as_ptr() as *const c_void
    }
}

unsafe impl Send for WeakBattleSideReference {}

unsafe impl Sync for WeakBattleSideReference {}

impl PartialEq for WeakBattleSideReference {
    fn eq(&self, other: &Self) -> bool {
        self.data.ptr_eq(&other.data)
    }
}

impl VolatileScriptsOwner for BattleSide {
    fn volatile_scripts(&self) -> &Arc<ScriptSet> {
        &self.data.volatile_scripts
    }

    fn load_volatile_script(&self, key: &StringKey) -> Result<Option<Arc<dyn Script>>> {
        self.battle()?
            .library()
            .load_script(self.into(), crate::ScriptCategory::Side, key)
    }
}

impl ScriptSource for BattleSide {
    fn get_script_count(&self) -> Result<usize> {
        Ok(self.battle()?.get_script_count()? + 1)
    }

    fn get_script_source_data(&self) -> &RwLock<ScriptSourceData> {
        &self.data.script_source_data
    }

    fn get_own_scripts(&self, scripts: &mut Vec<ScriptWrapper>) {
        scripts.push((&self.data.volatile_scripts).into());
    }

    fn collect_scripts(&self, scripts: &mut Vec<ScriptWrapper>) -> Result<()> {
        self.get_own_scripts(scripts);
        self.battle()?.collect_scripts(scripts)
    }
}