enemize-rs/enemize/src/rom.rs

use std::collections::BTreeMap;
use std::iter;
use std::ops::Range;

use anyhow::bail;

use crate::constants::*;
use crate::option_flags::OptionFlags;
use crate::Patch;

pub const ENEMIZER_INFO_SEED_OFFSET: usize = 0;
pub const ENEMIZER_INFO_SEED_LENGTH: usize = 12;
pub const ENEMIZER_INFO_VERSION_OFFSET: usize = ENEMIZER_INFO_SEED_OFFSET + ENEMIZER_INFO_SEED_LENGTH;
pub const ENEMIZER_INFO_VERSION_LENGTH: usize = 8;
pub const ENEMIZER_INFO_FLAGS_OFFSET: usize =
    ENEMIZER_INFO_VERSION_OFFSET + ENEMIZER_INFO_VERSION_LENGTH;
pub const ENEMIZER_INFO_FLAGS_LENGTH: usize = 0x50;

pub const RANDOMIZER_HIDDEN_ENEMIES_FLAG: usize = 0x00;
pub const CLOSE_BLIND_DOOR_FLAG: usize = 0x01;
pub const MOLDORM_EYES_FLAG: usize = 0x02;
pub const RANDOM_SPRITE_FLAG: usize = 0x03;
pub const AGAHNIM_BOUNCE_BALLS_FLAG: usize = 0x04;
pub const ENABLE_MIMIC_OVERRIDE_FLAG: usize = 0x05;
pub const ENABLE_TERRORPIN_AI_FIX_FLAG: usize = 0x06;
pub const CHECKSUM_COMPLIMENT_ADDRESS: usize = 0x7FDC;
pub const CHECKSUM_ADDRESS: usize = 0x7FDE;
pub const RANDOMIZER_MODE_FLAG: usize = 0x180032;

pub struct RomData {
    pub enemizer_info_table_base_address: usize,
    pub enemizer_option_flags_base_address: usize,
    spoiler: String,
    rom_data: Vec<u8>,
    patch_data: BTreeMap<usize, u8>,
    seed: i32,
}

impl RomData {
    pub fn spoiler(&self) -> &str {
        self.spoiler.as_str()
    }

    pub fn generate_patch(&self) -> Vec<Patch> {
        let mut patches = vec![];

        for (&addr, &byte) in self.patch_data.iter() {
            match patches.last_mut().filter(|p: &&mut Patch| p.address + 1 == addr) {
                None => patches.push(Patch { address: addr, patch_data: vec![byte] }),
                Some(patch) => patch.patch_data.push(byte)
            }
        }

        patches
    }

    fn set_rom_bytes(&mut self, bytes: &[u8], range: Range<usize>) {
        self.rom_data.splice(range, bytes.into_iter().map(|&b| b));
    }

    fn set_patch_bytes(&mut self, range: Range<usize>) {
        let slice = &self.rom_data[range.clone()];
        self.patch_data.extend(iter::zip(range, slice.into_iter().map(|&b| b)));
    }

    pub fn is_enemizer(&self) -> bool {
        self.rom_data.len() == ENEMIZER_FILE_LENGTH
             && self.rom_data[self.enemizer_info_table_base_address + ENEMIZER_INFO_SEED_OFFSET] == b'E'
             && self.rom_data[self.enemizer_info_table_base_address + ENEMIZER_INFO_SEED_OFFSET + 1] == b'N'
    }

    fn assert_rom_length(&self) {
        assert!(self.rom_data.len() >= ENEMIZER_FILE_LENGTH, "You need to expand the rom before you can use Enemizer features.");
    }

    pub fn get_enemizer_seed(&self) -> i32 {
        self.seed
    }

    pub fn derive_enemizer_seed(&mut self) -> anyhow::Result<i32> {
        if self.seed < 0 && self.is_enemizer() {
            let seed_start = self.enemizer_info_table_base_address + ENEMIZER_INFO_SEED_OFFSET;
            let seed_end = seed_start + ENEMIZER_INFO_SEED_LENGTH;
            let seed_bytes = &self.rom_data[seed_start..seed_end];
            let seed_str = &std::str::from_utf8(seed_bytes)?.trim_end_matches('\0')[2..];

            self.seed = i32::from_str_radix(seed_str, 10)?;
        }

        Ok(self.seed)
    }

    pub fn set_enemizer_seed(&mut self, seed: i32) {
        self.assert_rom_length();

        let seed_str = format!("EN{:<10}", seed);
        let bytes = seed_str.as_bytes().to_owned();

        let seed_start = self.enemizer_info_table_base_address + ENEMIZER_INFO_SEED_OFFSET;
        let seed_end = seed_start + ENEMIZER_INFO_SEED_LENGTH;
        self.seed = seed;

        self.set_rom_bytes(&bytes, seed_start..seed_end);
        self.set_patch_bytes(seed_start..seed_end);
    }

    pub fn get_enemizer_version(&self) -> anyhow::Result<&str> {
        if self.is_enemizer() {
            let version_start = self.enemizer_info_table_base_address + ENEMIZER_INFO_VERSION_OFFSET;
            let version_end = version_start + ENEMIZER_INFO_VERSION_LENGTH;
            Ok(std::str::from_utf8(&self.rom_data[version_start..version_end])?)
        } else {
            bail!("Not Enemizer Rom")
        }
    }

    pub fn set_enemizer_version(&mut self, version: String) {
        self.assert_rom_length();

        let mut bytes = version.into_bytes();
        bytes.resize(ENEMIZER_INFO_VERSION_LENGTH, 0);

        let version_start = self.enemizer_info_table_base_address + ENEMIZER_INFO_VERSION_OFFSET;
        let version_end = version_start + ENEMIZER_INFO_VERSION_LENGTH;
        self.set_rom_bytes(&bytes, version_start..version_end);
        self.set_patch_bytes(version_start..version_end)
    }

    pub fn set_info_flags(&mut self, flags: OptionFlags) -> anyhow::Result<()>{
        let bytes = flags.into_bytes();
        if bytes.len() > 0x100 - ENEMIZER_INFO_FLAGS_OFFSET {
            bail!("Option flags is too long to fit in the space allocated. Need to move data/code in asm file.");
        }

        let flags_start = self.enemizer_info_table_base_address + ENEMIZER_INFO_FLAGS_OFFSET;
        let flags_end = flags_start + bytes.len();
        self.set_rom_bytes(&bytes, flags_start..flags_end);
        self.set_patch_bytes(flags_start..flags_end);

        Ok(())
    }

    pub fn get_info_flags(&self) -> Option<OptionFlags> {
        if self.is_enemizer() {
            let flags_start = self.enemizer_info_table_base_address + ENEMIZER_INFO_FLAGS_OFFSET;
            let flags_end = flags_start + ENEMIZER_INFO_FLAGS_LENGTH;
            OptionFlags::try_from_bytes(&self.rom_data[flags_start..flags_end]).ok()
        } else {
            None
        }
    }

    pub fn patch_bytes(&mut self, address: usize, patch_data: Vec<u8>) {
        self.rom_data
            .splice(address..(address + patch_data.len()), patch_data);
    }

    pub fn patch_data(&mut self, patch: Patch) {
        self.set_rom_bytes(
            &patch.patch_data,
            patch.address..(patch.address + patch.patch_data.len()),
        );
    }
}