enemize-rs/enemize/src/rom.rs

use std::collections::BTreeMap;
use std::iter;
use std::ops::{Index, IndexMut, Range};

use anyhow::bail;

use crate::asar::{pc_to_snes_address, snes_to_pc_address, Symbols};
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 {
    asar_symbols: Symbols,
    spoiler: String,
    pub(crate) rom_data: Vec<u8>,
    patch_data: BTreeMap<usize, u8>,
    pub(crate) seed: i32,
}

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

    pub fn reset_enemizer(&mut self) -> i32 {
        use crate::constants::*;

        let seed = self.get_enemizer_seed();

        self.rom_data[ROOM_HEADER_BANK_LOCATION] = 0x04;

        let dungeon_header_range =
            DUNGEON_HEADER_POINTER_TABLE..(DUNGEON_HEADER_POINTER_TABLE + 640);
        self.rom_data[dungeon_header_range].copy_from_slice(&ORIGINAL_ROOM_POINTERS);

        let room_range = 0x5b97..(0x5b97 + 576);
        self.rom_data[room_range].copy_from_slice(&ORIGINAL_ROOM_BLOCKS);

        let ow_gfx_range = OVERWORLD_AREA_GRAPHICS_BLOCK..(OVERWORLD_AREA_GRAPHICS_BLOCK + 272);
        self.rom_data[ow_gfx_range].copy_from_slice(&ORIGINAL_OVERWORLD_BLOCKS);

        seed
    }

    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
    }

    pub fn get_rom_bytes(&self) -> &[u8] {
        &self.rom_data
    }

    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.asar_symbols["enemizer_info_table"] + ENEMIZER_INFO_SEED_OFFSET]
                == b'E'
            && self.rom_data
                [self.asar_symbols["enemizer_info_table"] + ENEMIZER_INFO_SEED_OFFSET + 1]
                == b'N'
    }

    pub fn is_randomizer(&self) -> bool {
        let acceptable = [
            b"VT", // item rando
            b"ER", // entrance rando
            b"DR", // door rando
            b"BM", // Berserker's multiworld
            b"BD", // Berserker's multiworld doors
            b"AP", // Archipelago
            b"AD", // Archipelago with door rando
        ];

        acceptable
            .iter()
            .any(|abbr| &abbr[..] == &self.rom_data[0x7fc0..0x7fc2])
            || (self.rom_data.len() >= 0x20_0000
                && &self.rom_data[0x7fc0..0x7fcf] == b"ZELDANODENSETSU")
    }

    pub fn is_race(&self) -> bool {
        self.is_randomizer()
            && (&self.rom_data[0x180213..0x180214] == &[1, 0]
                || &self.rom_data[0x7fc0..0x7fca] == b"VT TOURNEY")
    }

    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.asar_symbols["enemizer_info_table"] + 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.asar_symbols["enemizer_info_table"] + 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 expand_rom(&mut self) {
        self.rom_data.resize(0x40_0000, 0);
        // Update header length.
        self.rom_data[0x7fd7] = 0x0c;
        self.set_patch_bytes(0x7fd7..0x7fd8);

        self.set_enemizer_version("6.0.32".to_owned());
    }

    pub fn get_enemizer_version(&self) -> anyhow::Result<&str> {
        if self.is_enemizer() {
            let version_start =
                self.asar_symbols["enemizer_info_table"] + 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.asar_symbols["enemizer_info_table"] + 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.asar_symbols["enemizer_info_table"] + 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.asar_symbols["enemizer_info_table"] + 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 new(asar_symbols: Symbols, rom_data: Vec<u8>) -> Self {
        Self {
            asar_symbols,
            spoiler: String::new(),
            rom_data: rom_data,
            patch_data: BTreeMap::new(),
            seed: -1,
        }
    }

    fn get_flag(&self, offset: usize) -> bool {
        let flag_idx = self.asar_symbols["EnemizerFlags"] + offset;
        self.rom_data[flag_idx] == 1
    }

    fn set_flag(&mut self, offset: usize, val: bool) {
        let flag_idx = self.asar_symbols["EnemizerFlags"] + offset;
        self.rom_data[flag_idx] = if val { 1 } else { 0 };
        self.set_patch_bytes(flag_idx..(flag_idx + 1));
    }

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

    pub fn move_room_headers(&mut self) {
        let table_base = DUNGEON_HEADER_POINTER_TABLE;
        let header_base = self.asar_symbols["room_header_table"];

        // Change room header bank (at 0xb5e7) to 0x04.
        let new_room_bank =
            self.rom_data[self.asar_symbols["moved_room_header_bank_value_address"]];
        self.rom_data[ROOM_HEADER_BANK_LOCATION] = new_room_bank;

        // Copy header table.
        for i in 0..320 {
            // Get i'th room's pointer.
            // Pointers are 16bits, with a hard-coded bank.
            let room_pointer = [
                self.rom_data[table_base + (i * 2)],
                self.rom_data[table_base + (i * 2) + 1],
                4,
                0,
            ];
            let snes_address = u32::from_le_bytes(room_pointer);
            let pc_address = snes_to_pc_address(snes_address);

            // Copy i'th room's headers to new room_header_table.
            let header_start = header_base + (i * 14);
            self.rom_data
                .copy_within(pc_address..(pc_address + 14), header_start);
        }

        // Repoint the pointer table to the new header table.
        for i in 0..320 {
            let snes = pc_to_snes_address(header_base + (i * 14)).to_le_bytes();

            assert_eq!(snes[2], new_room_bank, "We changed banks in the middle of moving the room headers! This should have been caught by dev team, unless you were playing with files you shouldn't touch.");

            self.rom_data[(table_base + (i * 2))..(table_base + (i * 2) + 1)]
                .copy_from_slice(&snes[0..1]);
        }
    }
}

impl Index<usize> for RomData {
    type Output = u8;

    fn index(&self, index: usize) -> &u8 {
        &self.rom_data[index]
    }
}

impl IndexMut<usize> for RomData {
    fn index_mut(&mut self, index: usize) -> &mut u8 {
        &mut self.rom_data[index]
    }
}