enemize-rs/enemize/src/rom.rs

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

use anyhow::bail;

use crate::asar::{Symbols, snes_to_pc_address, pc_to_snes_address};
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,
    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.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 = [
            // item rando
            b"VT",
            // entrance rando
            b"ER",
            // door rando
            b"DR",
            // Berserker's multiworld
            b"BM",
            // Berserker's multiworld doors
            b"BD",
            // Archipelago
            b"AP",
            // Archipelago with door rando
            b"AD",
        ];

        acceptable.iter().any(|abbr| &abbr[..] == &self.rom_data[0x7FC0..0x7Fc1]) ||
            (self.rom_data.len() >= 0x20_0000 &&
             &self.rom_data[0x7FC0..0x7FCE] == b"ZELDANODENSETSU")
    }

    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 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_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()),
        );
    }

    pub fn move_room_headers(&mut self) {
        let table_base = DUNGEON_HEADER_POINTER_TABLE;
        let header_base = self.asar_symbols["rom_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]);
        }
    }
}