PkmnLib_rs/tests/common/library_loader.rs

use chrono::Duration;
use std::convert::TryFrom;
use std::fmt::Debug;
use std::fs::File;
use std::io::{BufReader, Read};
use std::sync::Arc;

use hashbrown::{HashMap, HashSet};
use num_traits::PrimInt;
use project_root::get_project_root;
use serde_json::Value;

use pkmn_lib::defines::LevelInt;
use pkmn_lib::dynamic_data::Gen7DamageLibrary;
use pkmn_lib::dynamic_data::Gen7MiscLibrary;
use pkmn_lib::dynamic_data::{DynamicLibrary, DynamicLibraryImpl};
use pkmn_lib::dynamic_data::{Gen7BattleStatCalculator, ScriptResolver};
use pkmn_lib::static_data::{
    AbilityImpl, AbilityLibrary, AbilityLibraryImpl, BattleItemCategory, DataLibrary, EvolutionData, EvolutionMethod,
    Form, FormImpl, Gender, GrowthRateLibrary, GrowthRateLibraryImpl, ItemImpl, ItemLibrary, ItemLibraryImpl,
    LearnableMoves, LearnableMovesImpl, LibrarySettingsImpl, LookupGrowthRate, MoveDataImpl, MoveLibrary,
    MoveLibraryImpl, NatureImpl, NatureLibrary, NatureLibraryImpl, Parameter, SecondaryEffect, SecondaryEffectImpl,
    SpeciesImpl, SpeciesLibrary, SpeciesLibraryImpl, StaticDataImpl, StaticStatisticSet, Statistic, TimeOfDay,
    TypeLibrary, TypeLibraryImpl,
};
use pkmn_lib::StringKey;

pub struct LoadResult {
    pub library: Arc<dyn DynamicLibrary>,
    pub types_load_time: Duration,
    pub natures_load_time: Duration,
    pub items_load_time: Duration,
    pub growth_rate_load_time: Duration,
    pub abilities_load_time: Duration,
    pub moves_load_time: Duration,
    pub species_load_time: Duration,
    pub wasm_load_time: Duration,
}

pub fn load_library() -> LoadResult {
    let mut path = get_project_root().unwrap();
    path.push("tests/data/");
    let path = path.to_str().unwrap().to_string();

    let t1 = chrono::Utc::now();
    let types = load_types(&path);
    let t2 = chrono::Utc::now();
    let types_load_time = t2 - t1;
    let natures = load_natures(&path);
    let t1 = chrono::Utc::now();
    let natures_load_time = t1 - t2;
    let items = load_items(&path);
    let t2 = chrono::Utc::now();
    let items_load_time = t2 - t1;
    let growth_rates = load_growth_rates(&path);
    let t1 = chrono::Utc::now();
    let growth_rate_load_time = t1 - t2;
    let abilities = load_abilities(&path);
    let t2 = chrono::Utc::now();
    let abilities_load_time = t2 - t1;
    let moves = load_moves(&path, &types);
    let t1 = chrono::Utc::now();
    let moves_load_time = t1 - t2;
    let species = load_species(&path, &types, &moves);
    let t2 = chrono::Utc::now();
    let species_load_time = t2 - t1;

    let data = StaticDataImpl::new(
        Arc::new(LibrarySettingsImpl::new(100, 100).unwrap()),
        species,
        moves,
        items,
        growth_rates,
        types,
        natures,
        abilities,
    );

    let t1 = chrono::Utc::now();
    let script_resolver = load_script_resolver(&path);
    let t2 = chrono::Utc::now();
    let wasm_load_time = t2 - t1;

    let library = Arc::new(DynamicLibraryImpl::new(
        Arc::new(data),
        Arc::new(Gen7BattleStatCalculator::new()),
        Arc::new(Gen7DamageLibrary::new(false)),
        Arc::new(Gen7MiscLibrary::new(Box::new(|| TimeOfDay::Day))),
        script_resolver,
    ));

    LoadResult {
        library,
        types_load_time,
        natures_load_time,
        items_load_time,
        growth_rate_load_time,
        abilities_load_time,
        moves_load_time,
        species_load_time,
        wasm_load_time,
    }
}

pub fn load_types(path: &String) -> Arc<dyn TypeLibrary> {
    let mut reader = csv::ReaderBuilder::new()
        .delimiter(b'|')
        .from_path(path.to_string() + "Types.csv")
        .unwrap();

    let type_library = TypeLibraryImpl::new(20);

    let headers = reader.headers().unwrap();
    for header in headers.iter().skip(1) {
        type_library.register_type(&header.into());
    }

    for record in reader.records() {
        let record = record.unwrap();
        let offensive_type = record.get(0).unwrap();
        let offensive_type_id = type_library.get_type_id(&offensive_type.into()).unwrap();

        for (i, v) in record.iter().skip(1).enumerate() {
            let effectiveness = v.parse::<f32>().unwrap();
            type_library
                .set_effectiveness(offensive_type_id, ((i + 1) as u8).into(), effectiveness)
                .unwrap();
        }
    }
    Arc::new(type_library)
}

pub fn load_natures(path: &String) -> Arc<dyn NatureLibrary> {
    let mut reader = csv::ReaderBuilder::new()
        .delimiter(b'|')
        .from_path(path.to_string() + "Natures.csv")
        .unwrap();

    let nature_library = NatureLibraryImpl::new(24);
    for record in reader.records() {
        let record = record.unwrap();
        let nature_name = record.get(0).unwrap().into();
        let increased_statistic_str = record.get(1).unwrap();
        let decreased_statistic_str = record.get(2).unwrap();
        if increased_statistic_str.is_empty() || decreased_statistic_str.is_empty() {
            nature_library.load_nature(nature_name, NatureImpl::new(Statistic::HP, Statistic::HP, 1.0, 1.0));
        } else {
            let increased_statistic = serde_plain::from_str(increased_statistic_str).unwrap();
            let decreased_statistic = serde_plain::from_str(decreased_statistic_str).unwrap();
            nature_library.load_nature(
                nature_name,
                NatureImpl::new(increased_statistic, decreased_statistic, 1.1, 0.9),
            );
        }
    }
    Arc::new(nature_library)
}

pub fn load_items(path: &String) -> Arc<dyn ItemLibrary> {
    let mut file = File::open(path.to_string() + "Items.json").unwrap();
    let mut data = String::new();
    file.read_to_string(&mut data).unwrap();
    let json: Value = serde_json::from_str(&data).unwrap();
    let json_array = json.as_array().unwrap();

    let item_library = ItemLibraryImpl::new(400);
    for v in json_array {
        let name = v.get("name").unwrap().as_str().unwrap().into();
        let category = serde_json::from_value(v.get("itemType").unwrap().clone()).unwrap();
        let mut battle_category = BattleItemCategory::None;
        if let Some(c) = v.get("battleType") {
            battle_category = serde_json::from_value(c.clone()).unwrap();
        }
        let price = v.get("price").unwrap().as_i64().unwrap();
        let mut flags = HashSet::new();
        if let Some(f) = v.get("flags") {
            let a = f.as_array().unwrap();
            for flag in a {
                flags.insert(flag.as_str().unwrap().into());
            }
        }

        item_library.add(
            &name,
            Arc::new(ItemImpl::new(&name, category, battle_category, price as i32, flags)),
        );
    }
    Arc::new(item_library)
}

pub fn load_growth_rates(path: &String) -> Arc<dyn GrowthRateLibrary> {
    let mut file = File::open(path.to_string() + "GrowthRates.json").unwrap();
    let mut data = String::new();
    file.read_to_string(&mut data).unwrap();
    let json: Value = serde_json::from_str(&data).unwrap();
    let o = json.as_object().unwrap();

    let growth_rate_library = GrowthRateLibraryImpl::new(10);
    for (key, value) in o {
        let name = StringKey::new(key);
        let experience_required_json = value.as_array().unwrap();
        let mut experience_required = Vec::with_capacity(experience_required_json.len());
        for v in experience_required_json {
            experience_required.push(v.as_i64().unwrap() as u32);
        }

        growth_rate_library.add_growth_rate(&name, Arc::new(LookupGrowthRate::new(experience_required)));
    }
    Arc::new(growth_rate_library)
}

pub fn load_abilities(path: &String) -> Arc<dyn AbilityLibrary> {
    let mut file = File::open(path.to_string() + "Abilities.json").unwrap();
    let mut data = String::new();
    file.read_to_string(&mut data).unwrap();
    let json: Value = serde_json::from_str(&data).unwrap();
    let o = json.as_object().unwrap();

    let ability_library = AbilityLibraryImpl::new(400);
    for (key, value) in o {
        let name = StringKey::new(key);
        let mut effect = StringKey::empty();
        if let Some(e) = value.get("effect") {
            effect = e.as_str().unwrap().into();
        }
        let mut parameters = HashMap::new();
        if let Some(pars) = value.get("parameters") {
            let pars = pars.as_object().unwrap();
            for par in pars {
                parameters.insert(par.0.as_str().into(), parse_parameter(par.1));
            }
        }

        ability_library.add(&name, Arc::new(AbilityImpl::new(&name, &effect, parameters)));
    }
    Arc::new(ability_library)
}

pub fn load_moves(path: &String, types: &Arc<dyn TypeLibrary>) -> Arc<dyn MoveLibrary> {
    let mut file = File::open(path.to_string() + "Moves.json").unwrap();
    let mut data = String::new();
    file.read_to_string(&mut data).unwrap();
    let json: Value = serde_json::from_str(&data).unwrap();
    let data = json.as_object().unwrap().get("data").unwrap().as_array().unwrap();
    let move_library = MoveLibraryImpl::new(600);
    for move_data in data {
        let move_data = move_data.as_object().unwrap();
        let move_name = move_data.get("name").unwrap().as_str().unwrap().into();
        let move_type = move_data.get("type").unwrap().as_str().unwrap().into();
        let move_type_id = types.get_type_id(&move_type).unwrap();
        let move_category = serde_json::from_value(move_data.get("category").unwrap().clone()).unwrap();
        let base_power = move_data.get("power").unwrap().as_i64().unwrap() as u8;
        let accuracy = move_data.get("accuracy").unwrap().as_i64().unwrap() as u8;
        let pp = move_data.get("pp").unwrap().as_i64().unwrap() as u8;
        let target = serde_json::from_value(move_data.get("target").unwrap().clone()).unwrap();
        let priority = move_data.get("priority").unwrap().as_i64().unwrap() as i8;
        let secondary_effect: Option<Arc<dyn SecondaryEffect>> = if let Some(v) = move_data.get("effect") {
            let mut chance = -1.0;
            if let Some(chance_value) = v.get("chance") {
                chance = chance_value.as_f64().unwrap() as f32;
            }
            let mut parameters = HashMap::new();
            if let Some(pars) = v.get("parameters") {
                let pars = pars.as_object().unwrap();
                for par in pars {
                    parameters.insert(par.0.as_str().into(), parse_parameter(par.1));
                }
            }

            Some(Arc::new(SecondaryEffectImpl::new(
                chance,
                v.get("name").unwrap().as_str().unwrap().into(),
                parameters,
            )))
        } else {
            None
        };

        let mut flags = HashSet::new();
        if let Some(f) = move_data.get("flags") {
            let f = f.as_array().unwrap();
            for flag in f {
                flags.insert(flag.as_str().unwrap().into());
            }
        }

        move_library.add(
            &move_name,
            Arc::new(MoveDataImpl::new(
                &move_name.clone(),
                move_type_id,
                move_category,
                base_power,
                accuracy,
                pp,
                target,
                priority,
                secondary_effect,
                flags,
            )),
        );
    }
    Arc::new(move_library)
}

pub fn load_species(
    path: &String,
    types: &Arc<dyn TypeLibrary>,
    moves: &Arc<dyn MoveLibrary>,
) -> Arc<dyn SpeciesLibrary> {
    let mut file = File::open(path.to_string() + "Pokemon.json").unwrap();
    let mut data = String::new();
    file.read_to_string(&mut data).unwrap();
    let json: Value = serde_json::from_str(&data).unwrap();
    let o = json.as_object().unwrap();

    let species_library = SpeciesLibraryImpl::new(800);
    for (key, value) in o.iter() {
        if key.starts_with('$') {
            continue;
        }
        let name = StringKey::new(key);
        let id = value.get("id").unwrap().as_i64().unwrap();
        let gender_rate = value.get("genderRatio").unwrap().as_f64().unwrap();
        let growth_rate_name = value.get("growthRate").unwrap().as_str().unwrap().into();
        let base_happiness = value.get("baseHappiness").unwrap().as_i64().unwrap();
        let catch_rate = value.get("catchRate").unwrap().as_i64().unwrap();
        let _color = value.get("color").unwrap().as_str().unwrap();
        // let egg_groups = value.get("eggGroups").unwrap()
        //     .as_array()
        //     .unwrap()
        //     .iter()
        //     .map(|&a| a.as_str().unwrap())
        //     .collect();
        let _egg_cycle = value.get("eggCycles").unwrap().as_i64().unwrap();
        // TODO: tags
        let evolutions = if let Some(evolutions) = value.get("evolutions") {
            evolutions.as_array().unwrap().iter().map(parse_evolution).collect()
        } else {
            Vec::new()
        };

        let forms = value.get("formes").unwrap().as_object().unwrap();
        let default_form_value = forms.get("default").unwrap();
        let default_form = parse_form("default".into(), default_form_value, types, moves);

        let species = SpeciesImpl::new(
            id as u16,
            &name,
            gender_rate as f32,
            &growth_rate_name,
            catch_rate as u8,
            base_happiness as u8,
            default_form,
            Default::default(),
            evolutions,
        );
        species_library.add(&name, Arc::new(species));
    }
    Arc::new(species_library)
}

fn parse_evolution(value: &Value) -> EvolutionData {
    let species = StringKey::new(value.get("species").unwrap().as_str().unwrap());
    let method = value.get("method").unwrap().as_str().unwrap();
    let method = match method {
        "level" => {
            let level = value.get("data").unwrap().as_i64().unwrap() as LevelInt;
            EvolutionMethod::Level { level }
        }
        "levelfemale" => {
            let level = value.get("data").unwrap().as_i64().unwrap() as LevelInt;
            EvolutionMethod::LevelGender {
                level,
                gender: Gender::Female,
            }
        }
        "levelmale" => {
            let level = value.get("data").unwrap().as_i64().unwrap() as LevelInt;
            EvolutionMethod::LevelGender {
                level,
                gender: Gender::Male,
            }
        }
        "happiness" => {
            let happiness = value.get("data").unwrap().as_i64().unwrap() as u8;
            EvolutionMethod::Happiness { happiness }
        }
        "happinessday" => {
            let happiness = value.get("data").unwrap().as_i64().unwrap() as u8;
            EvolutionMethod::HappinessDay { happiness }
        }
        "happinessnight" => {
            let happiness = value.get("data").unwrap().as_i64().unwrap() as u8;
            EvolutionMethod::HappinessNight { happiness }
        }
        "item" => {
            let item = StringKey::new(value.get("data").unwrap().as_str().unwrap());
            EvolutionMethod::Item { item }
        }
        "itemmale" => {
            let item = StringKey::new(value.get("data").unwrap().as_str().unwrap());
            EvolutionMethod::ItemGender {
                item,
                gender: Gender::Male,
            }
        }
        "itemfemale" => {
            let item = StringKey::new(value.get("data").unwrap().as_str().unwrap());
            EvolutionMethod::ItemGender {
                item,
                gender: Gender::Female,
            }
        }
        "holditem" => {
            let item = StringKey::new(value.get("data").unwrap().as_str().unwrap());
            EvolutionMethod::HoldItem { item }
        }
        "dayholditem" => {
            let item = StringKey::new(value.get("data").unwrap().as_str().unwrap());
            EvolutionMethod::DayHoldItem { item }
        }
        "nightholditem" => {
            let item = StringKey::new(value.get("data").unwrap().as_str().unwrap());
            EvolutionMethod::NightHoldItem { item }
        }
        "hasmove" => {
            let move_name = StringKey::new(value.get("data").unwrap().as_str().unwrap());
            EvolutionMethod::HasMove { move_name }
        }
        "trade" => EvolutionMethod::Trade,
        "tradespecies" => {
            let species = StringKey::new(value.get("data").unwrap().as_str().unwrap());
            EvolutionMethod::TradeSpecies { species }
        }
        "tradeitem" => {
            let item = StringKey::new(value.get("data").unwrap().as_str().unwrap());
            EvolutionMethod::TradeItem { item }
        }
        "location" => {
            let location = StringKey::new(value.get("data").unwrap().as_str().unwrap());
            EvolutionMethod::Location { location }
        }
        "custom" => {
            let data = value.get("data").unwrap().as_array().unwrap();
            let name = StringKey::new(data[0].as_str().unwrap());
            let params = data[1..].iter().map(parse_parameter).collect();
            EvolutionMethod::Custom { name, params }
        }
        _ => panic!("Unknown evolution method: {}", method),
    };
    EvolutionData::new(method, species)
}

#[cfg(not(any(feature = "wasm", feature = "rune")))]
fn load_script_resolver(path: &String) -> Arc<dyn ScriptResolver> { Arc::new(EmptyScriptResolver::default()) }

#[cfg(feature = "wasm")]
fn load_script_resolver(path: &String) -> Arc<dyn ScriptResolver> {
    let resolver = pkmn_lib::script_implementations::wasm::script_resolver::WebAssemblyScriptResolver::new();
    let file = File::open(path.to_string() + "gen7_scripts.wasm").unwrap();
    let mut reader = BufReader::new(file);
    let mut buffer = Vec::new();
    reader.read_to_end(&mut buffer).unwrap();
    resolver.load_wasm_from_bytes(&buffer).unwrap();
    resolver.finalize().unwrap();
    resolver
}

#[cfg(feature = "rune")]
fn load_script_resolver(path: &String) -> Arc<dyn ScriptResolver> {
    let mut builder =
        pkmn_lib::script_implementations::rune::script_resolver::RuneScriptResolverBuilder::new().unwrap();
    // Recursively load all scripts in the scripts folder
    for entry in walkdir::WalkDir::new(path.to_string() + "scripts/") {
        let entry = entry.unwrap();
        let path = entry.path();
        if path.is_file() {
            let file = File::open(&path).unwrap();
            let mut reader = BufReader::new(file);
            let mut buffer = Vec::new();
            reader.read_to_end(&mut buffer).unwrap();
            builder
                .insert_script(path, &path.to_string_lossy(), &String::from_utf8(buffer).unwrap())
                .unwrap();
        }
    }
    let resolver = builder.build().unwrap();
    resolver
}

fn parse_form(
    name: StringKey,
    value: &Value,
    types: &Arc<dyn TypeLibrary>,
    moves: &Arc<dyn MoveLibrary>,
) -> Arc<dyn Form> {
    let mut abilities = Vec::new();
    for a in value.get("abilities").unwrap().as_array().unwrap() {
        abilities.push(a.as_str().unwrap().into());
    }
    let mut hidden_abilities = Vec::new();
    for a in value.get("hiddenAbilities").unwrap().as_array().unwrap() {
        hidden_abilities.push(a.as_str().unwrap().into());
    }
    let base_stats = parse_statistics(value.get("baseStats").unwrap());
    // TODO: ev reward
    let height = value.get("height").unwrap().as_f64().unwrap();
    let weight = value.get("weight").unwrap().as_f64().unwrap();
    let base_experience = value.get("baseExp").unwrap().as_u64().unwrap();
    let types = value
        .get("types")
        .unwrap()
        .as_array()
        .unwrap()
        .iter()
        .map(|a| types.get_type_id(&a.as_str().unwrap().into()).unwrap())
        .collect();

    let moves = parse_moves(value.get("moves").unwrap(), moves);

    Arc::new(FormImpl::new(
        &name,
        height as f32,
        weight as f32,
        base_experience as u32,
        types,
        base_stats,
        abilities,
        hidden_abilities,
        moves,
        Default::default(),
    ))
}

fn parse_statistics<T>(value: &Value) -> Arc<StaticStatisticSet<T>>
where
    T: PrimInt + TryFrom<u64>,
    <T as TryFrom<u64>>::Error: Debug,
{
    Arc::new(StaticStatisticSet::new(
        <T as TryFrom<u64>>::try_from(value.get("hp").unwrap_or(&Value::Number(0.into())).as_u64().unwrap()).unwrap(),
        <T as TryFrom<u64>>::try_from(
            value
                .get("attack")
                .unwrap_or(&Value::Number(0.into()))
                .as_u64()
                .unwrap(),
        )
        .unwrap(),
        <T as TryFrom<u64>>::try_from(
            value
                .get("defense")
                .unwrap_or(&Value::Number(0.into()))
                .as_u64()
                .unwrap(),
        )
        .unwrap(),
        <T as TryFrom<u64>>::try_from(
            value
                .get("specialAttack")
                .unwrap_or(&Value::Number(0.into()))
                .as_u64()
                .unwrap(),
        )
        .unwrap(),
        <T as TryFrom<u64>>::try_from(
            value
                .get("specialDefense")
                .unwrap_or(&Value::Number(0.into()))
                .as_u64()
                .unwrap(),
        )
        .unwrap(),
        <T as TryFrom<u64>>::try_from(value.get("speed").unwrap_or(&Value::Number(0.into())).as_u64().unwrap())
            .unwrap(),
    ))
}

fn parse_moves(value: &Value, move_library: &Arc<dyn MoveLibrary>) -> Arc<dyn LearnableMoves> {
    let moves = LearnableMovesImpl::new(100);

    let level_moves = value.get("levelMoves").unwrap().as_array().unwrap();
    for level_move in level_moves {
        let name = level_move.get("name").unwrap().as_str().unwrap().into();
        let level = level_move.get("level").unwrap().as_u64().unwrap() as LevelInt;
        assert!(move_library.get(&name).is_some());
        moves.add_level_move(level, &name).unwrap();
    }

    Arc::new(moves)
}

fn parse_parameter(value: &Value) -> Arc<Parameter> {
    Arc::new(match value {
        Value::Null => {
            panic!("Unexpected type")
        }
        Value::Bool(b) => (*b).into(),
        Value::Number(n) => {
            if n.is_f64() {
                (n.as_f64().unwrap() as f32).into()
            } else {
                n.as_i64().unwrap().into()
            }
        }
        Value::String(s) => StringKey::new(s.as_str()).into(),
        Value::Array(_) => {
            panic!("Unexpected type")
        }
        Value::Object(_) => {
            panic!("Unexpected type")
        }
    })
}

#[test]
#[cfg_attr(miri, ignore)]
fn test_type_library_loaded() {
    let mut path = get_project_root().unwrap();
    path.push("tests/data/");
    let types = load_types(&path.to_str().unwrap().to_string());

    assert_eq!(
        types
            .get_effectiveness(
                types.get_type_id(&"fire".into()).unwrap(),
                &[types.get_type_id(&"grass".into()).unwrap()],
            )
            .unwrap(),
        2.0
    );
}