//! Mapper 056 - Kaiser KS202 (Pirate SMB3)
//!
//! Specifications:
//! - Main: <https://www.nesdev.org/wiki/INES_Mapper_056>
//!
//! Known Limitations:
//! - PRG-RAM is exposed at $6000-$7FFF but not backed by battery save.

use crate::nes::cartridge::BaseMapper;
use crate::nes::cartridge::NametableLayout;
use crate::nes::cartridge::mapper::{Mapper, MapperCapabilities};

/// Mapper 056 - Kaiser KS202
///
/// Hardware: KS202 ASIC (an upgrade to Konami's VRC3)
///
/// Specifications:
/// - Main: <https://www.nesdev.org/wiki/INES_Mapper_056>
/// - PRG-ROM: Up to 256 KiB (32 × 8 KiB banks via bank reg + PRG A17 bit)
/// - CHR: Up to 128 KiB (128 × 1 KiB banks, 7-bit registers)
/// - Mirroring: Programmable (H/V)
/// - IRQ: 16-bit CPU-cycle counter, VRC3-like
///
/// Register map (CPU address space):
/// - $8000-$8FFF: IRQ latch nibble 0 [3:0]
/// - $9000-$9FFF: IRQ latch nibble 1 [7:4]
/// - $A000-$AFFF: IRQ latch nibble 2 [11:8]
/// - $B000-$BFFF: IRQ latch nibble 3 [15:12]
/// - $C000-$CFFF: IRQ control (bit1=E: enable-now; bit0 is ignored, unlike VRC3)
/// - $D000-$DFFF: IRQ acknowledge
/// - $E000-$EFFF: Bank register select (bits [2:0], values 1/2/3 for $8000/$A000/$C000)
/// - $F000-$FFFF: Bank data and sub-registers (superimposed):
///   - $F000-$F3FF (mask $FC03): PRG A17 bit for banks 0-3 (bit 4 of written value)
///   - $F800-$FBFF (mask $FC00): Mirroring (bit 0: 0=H, 1=V)
///   - $FC00-$FC07 (mask $FC07): CHR 1KB banks 0-7 (7-bit value)
///   - All $F000-$FFFF: Bank data [3:0] → update last-selected PRG bank register
///
/// PRG effective 5-bit bank = (prg_a17[window] << 4) | prg_reg[window]
/// Power-on: prg_a17[0..3] all = 1
pub struct Mapper56 {
    base: BaseMapper,
    prg_reg: [u8; 3],  // 4-bit PRG bank selects for $8000/$A000/$C000
    prg_a17: [u8; 4],  // A17 extension bit for each 8KB window (0-3)
    chr_regs: [u8; 8], // 7-bit CHR 1KB bank selects
    bank_select: u8,   // Last value written to $E000
    irq_latch: u16,
    irq_counter: u16,
    irq_enabled: bool,
    irq_pending: bool,
}

impl Mapper56 {
    pub fn new(ctx: crate::nes::cartridge::mapper::MapperContext) -> Self {
        let capabilities = MapperCapabilities {
            has_irq: true,
            has_chr_banking: true,
            has_dynamic_mirroring: true,
            max_prg_ram_kb: 8,
            prg_bank_size_kb: 8,
            chr_bank_size_kb: 1,
            ..Default::default()
        };
        let mut base = BaseMapper::new(&ctx, capabilities);
        base.configure_prg_banking(0x2000);
        base.configure_chr_banking(0x0400);
        base.set_mirroring(NametableLayout::Vertical);
        let mut mapper = Self {
            base,
            prg_reg: [0; 3],
            prg_a17: [1; 4], // power-on: holding 1 per spec
            chr_regs: [0; 8],
            bank_select: 0,
            irq_latch: 0,
            irq_counter: 0,
            irq_enabled: false,
            irq_pending: false,
        };
        mapper.update_banks();
        mapper
    }

    fn update_banks(&mut self) {
        // PRG slots 0-2: switchable via bank register + A17 extension
        for i in 0..3 {
            let a17 = (self.prg_a17[i] as i16) & 1;
            let reg = (self.prg_reg[i] as i16) & 0x0F;
            self.base.select_prg_page(i, (a17 << 4) | reg);
        }
        // PRG slot 3: fixed to last bank within the A17 block
        let a17 = (self.prg_a17[3] as usize) & 1;
        let block_last = (a17 << 4) | 0x0F;
        let count = self.base.prg_bank_count();
        let bank = if count == 0 {
            0
        } else {
            block_last.min(count - 1)
        };
        self.base.select_prg_page(3, bank as i16);
        // CHR: 8 × 1KB slots
        for i in 0..8 {
            self.base.select_chr_page(i, self.chr_regs[i] as i16);
        }
    }
}

impl Mapper for Mapper56 {
    fn base(&self) -> &BaseMapper {
        &self.base
    }
    fn base_mut(&mut self) -> &mut BaseMapper {
        &mut self.base
    }

    fn read_prg(&self, addr: u16) -> u8 {
        match addr {
            0x6000..=0x7FFF => self.base.try_read_prg_ram(addr).unwrap_or(0),
            0x8000..=0xFFFF => self.base.read_prg_banked(addr),
            _ => 0,
        }
    }

    fn write_prg(&mut self, addr: u16, value: u8) {
        match addr {
            0x6000..=0x7FFF => {
                self.base.try_write_prg_ram(addr, value);
            }
            0x8000..=0x8FFF => {
                self.irq_latch = (self.irq_latch & 0xFFF0) | ((value as u16) & 0x0F);
            }
            0x9000..=0x9FFF => {
                self.irq_latch = (self.irq_latch & 0xFF0F) | (((value as u16) & 0x0F) << 4);
            }
            0xA000..=0xAFFF => {
                self.irq_latch = (self.irq_latch & 0xF0FF) | (((value as u16) & 0x0F) << 8);
            }
            0xB000..=0xBFFF => {
                self.irq_latch = (self.irq_latch & 0x0FFF) | (((value as u16) & 0x0F) << 12);
            }
            0xC000..=0xCFFF => {
                // IRQ control: bit1=E (enable now); bit0 is ignored on KS202
                // (unlike VRC3 where bit0 re-enables after acknowledge).
                // Spec (KS202/Mesen): reload counter from latch immediately when E=1;
                // always clear pending IRQ.
                self.irq_enabled = (value & 0x02) != 0;
                if self.irq_enabled {
                    self.irq_counter = self.irq_latch;
                }
                self.irq_pending = false;
            }
            0xD000..=0xDFFF => {
                // IRQ acknowledge: only clear pending — do NOT reload counter.
                self.irq_pending = false;
            }
            0xE000..=0xEFFF => {
                self.bank_select = value & 0x07;
            }
            0xF000..=0xFFFF => {
                // Bank data write (primary)
                match self.bank_select {
                    1 => self.prg_reg[0] = value & 0x0F,
                    2 => self.prg_reg[1] = value & 0x0F,
                    3 => self.prg_reg[2] = value & 0x0F,
                    _ => {}
                }

                // Superimposed register: PRG A17 ($F000-$F3FF, mask $FC03)
                // Spec: bit 4 ([...P ....] = 0x10) holds the A17 extension bit.
                if (addr & 0xFC00) == 0xF000 {
                    let slot = (addr & 0x0003) as usize;
                    if slot < 4 {
                        self.prg_a17[slot] = (value >> 4) & 0x01;
                    }
                }

                // Superimposed register: Mirroring ($F800-$FBFF, mask $FC00)
                if (addr & 0xFC00) == 0xF800 {
                    self.base.set_mirroring_hv((value & 0x01) == 0);
                }

                // Superimposed register: CHR banks ($FC00-$FC07, mask $FC07)
                if (addr & 0xFC00) == 0xFC00 {
                    let slot = (addr & 0x0007) as usize;
                    self.chr_regs[slot] = value & 0x7F;
                }

                self.update_banks();
            }
            _ => {}
        }
    }

    fn irq_pending(&self) -> bool {
        self.irq_pending
    }

    fn cpu_cycle(&mut self) {
        if !self.irq_enabled {
            return;
        }
        // Spec (KS202/Mesen): counter counts UP; IRQ fires when it reaches 0xFFFF.
        self.irq_counter = self.irq_counter.wrapping_add(1);
        if self.irq_counter == 0xFFFF {
            self.irq_pending = true;
            self.irq_enabled = false;
        }
    }

    fn registers_snapshot(&self) -> Vec<u8> {
        let mirror_byte = match self.base.mirroring() {
            NametableLayout::Vertical => 1u8,
            _ => 0u8,
        };
        let irq_flags = (self.irq_enabled as u8) | ((self.irq_pending as u8) << 1);
        let mut v = vec![
            self.prg_reg[0],
            self.prg_reg[1],
            self.prg_reg[2],
            self.prg_a17[0],
            self.prg_a17[1],
            self.prg_a17[2],
            self.prg_a17[3],
            mirror_byte,
            self.bank_select,
            irq_flags,
            (self.irq_latch & 0xFF) as u8,
            (self.irq_latch >> 8) as u8,
            (self.irq_counter & 0xFF) as u8,
            (self.irq_counter >> 8) as u8,
        ];
        v.extend_from_slice(&self.chr_regs);
        v
    }

    fn restore_registers(&mut self, data: &[u8]) {
        if data.len() < 22 {
            return;
        }
        self.prg_reg[0] = data[0];
        self.prg_reg[1] = data[1];
        self.prg_reg[2] = data[2];
        self.prg_a17[0] = data[3];
        self.prg_a17[1] = data[4];
        self.prg_a17[2] = data[5];
        self.prg_a17[3] = data[6];
        self.base.set_mirroring_hv(data[7] == 0);
        self.bank_select = data[8];
        self.irq_enabled = (data[9] & 1) != 0;
        self.irq_pending = (data[9] & 2) != 0;
        self.irq_latch = (data[10] as u16) | ((data[11] as u16) << 8);
        self.irq_counter = (data[12] as u16) | ((data[13] as u16) << 8);
        self.chr_regs.copy_from_slice(&data[14..22]);
        self.update_banks();
    }

    fn reset(&mut self) {
        self.prg_reg = [0; 3];
        self.prg_a17 = [1; 4];
        self.chr_regs = [0; 8];
        self.base.set_mirroring(NametableLayout::Vertical);
        self.bank_select = 0;
        self.irq_latch = 0;
        self.irq_counter = 0;
        self.irq_enabled = false;
        self.irq_pending = false;
        self.update_banks();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::nes::cartridge::mapper::{MapperContext, create_mapper};
    use crate::nes::cartridge::test_helpers::banked_data;

    const PRG_BANKS: usize = 16; // 128 KiB
    const CHR_BANKS: usize = 64; // 64 KiB

    fn make_mapper() -> Mapper56 {
        let prg = banked_data(8 * 1024, PRG_BANKS);
        let chr = banked_data(1024, CHR_BANKS);
        Mapper56::new(MapperContext::new_for_test(
            56,
            prg,
            chr,
            NametableLayout::Horizontal,
        ))
    }

    #[test]
    fn mapper_56_is_registered() {
        let result = create_mapper(MapperContext::new_for_test(
            56,
            banked_data(8 * 1024, PRG_BANKS),
            banked_data(1024, CHR_BANKS),
            NametableLayout::Vertical,
        ));
        assert!(result.is_ok(), "Mapper 56 must be registered");
    }

    // --- PRG banking via $E000/$F000 ---

    #[test]
    fn prg_bank_select_via_e000_and_f000() {
        let mut mapper = make_mapper();
        mapper.write_prg(0xE000, 1); // select bank reg for $8000
        mapper.write_prg(0xF000, 5); // set bank 5; prg_a17[0] = (5>>3)&1 = 0
        assert_eq!(mapper.prg_reg[0], 5 & 0x0F, "PRG reg0 must be 5");
        // Effective bank = (0 << 4) | 5 = 5
        assert_eq!(mapper.read_prg(0x8000), 5, "PRG $8000 must map to bank 5");
    }

    #[test]
    fn prg_a17_set_by_superimposed_f000_write() {
        let mut mapper = make_mapper();
        // PRG A17 register: addr $F000, bit 4 = A17 for bank 0
        // Writing value = 0x10 (bit4=1) to $F000 sets prg_a17[0]=1
        mapper.write_prg(0xE000, 1);
        mapper.write_prg(0xF000, 0x10); // prg_a17[0] = 1, prg_reg[0] = 0
        assert_eq!(mapper.prg_a17[0], 1, "PRG A17 bit must be 1");
        // bit 3 alone (0x08) must NOT set A17
        mapper.write_prg(0xF000, 0x08); // prg_a17[0] = 0, prg_reg[0] = 8
        assert_eq!(mapper.prg_a17[0], 0, "bit 3 alone must NOT set A17");
    }

    /// Spec: $F000 PRG A17 uses bit 4 ([...P ....]), not bit 3 ([....P...])
    #[test]
    fn prg_a17_extracted_from_bit4_not_bit3() {
        let mut mapper = make_mapper();
        // Value 0x10 has only bit 4 set → A17 must be 1
        mapper.write_prg(0xF000, 0x10);
        assert_eq!(mapper.prg_a17[0], 1, "value 0x10 (bit 4) must set A17=1");
        // Value 0x08 has only bit 3 set → A17 must be 0
        mapper.write_prg(0xF000, 0x08);
        assert_eq!(mapper.prg_a17[0], 0, "value 0x08 (bit 3) must NOT set A17");
    }

    #[test]
    fn prg_e000_fixed_last_bank() {
        let mapper = make_mapper();
        // prg_a17[3] = 1 on power-on; fixed last = highest in block starting at A17=1
        // Block starts at bank 16, size 16 → bank 31; but PRG_BANKS=16, so bank 15
        assert_eq!(
            mapper.read_prg(0xE000),
            (PRG_BANKS - 1) as u8,
            "$E000-$FFFF must be fixed to last bank"
        );
    }

    // --- CHR banking via $FC00-$FC07 ---

    #[test]
    fn chr_banks_via_fc00() {
        let mut mapper = make_mapper();
        for slot in 0..8u16 {
            mapper.write_prg(0xFC00 + slot, (slot * 7) as u8 & 0x3F);
        }
        for slot in 0..8u16 {
            let bank = (slot * 7) as usize & 0x3F;
            let expected = (bank % CHR_BANKS) as u8;
            assert_eq!(
                mapper.read_chr(slot * 1024),
                expected,
                "CHR slot {slot} wrong bank"
            );
        }
    }

    // --- Mirroring via $F800 ---

    #[test]
    fn mirroring_horizontal_via_f800() {
        let mut mapper = make_mapper();
        mapper.write_prg(0xF800, 0x00); // 0 = horizontal
        assert_eq!(mapper.get_mirroring(), NametableLayout::Horizontal);
    }

    #[test]
    fn mirroring_vertical_via_f800() {
        let mut mapper = make_mapper();
        mapper.write_prg(0xF800, 0x01); // 1 = vertical
        assert_eq!(mapper.get_mirroring(), NametableLayout::Vertical);
    }

    // --- IRQ ---

    #[test]
    fn irq_not_pending_by_default() {
        let mapper = make_mapper();
        assert!(!mapper.irq_pending());
    }

    #[test]
    fn irq_latch_nibbles() {
        let mut mapper = make_mapper();
        mapper.write_prg(0x8000, 0x01); // nibble 0 = 1
        mapper.write_prg(0x9000, 0x02); // nibble 1 = 2 → bits 7:4 = 0x20
        mapper.write_prg(0xA000, 0x03); // nibble 2 = 3 → bits 11:8 = 0x300
        mapper.write_prg(0xB000, 0x04); // nibble 3 = 4 → bits 15:12 = 0x4000
        assert_eq!(mapper.irq_latch, 0x4321, "IRQ latch must be 0x4321");
    }

    #[test]
    fn irq_fires_after_n_cycles() {
        let mut mapper = make_mapper();
        // Set latch to 0xFFFC so IRQ fires after exactly 3 cpu_cycle calls.
        // Counter counts UP and fires when it reaches 0xFFFF.
        mapper.write_prg(0x8000, 0x0C); // bits[3:0] = C
        mapper.write_prg(0x9000, 0x0F); // bits[7:4] = F
        mapper.write_prg(0xA000, 0x0F); // bits[11:8] = F
        mapper.write_prg(0xB000, 0x0F); // bits[15:12] = F → latch = 0xFFFC
        mapper.write_prg(0xC000, 0x02); // E=1: enable and reload counter to 0xFFFC
        for _ in 0..2 {
            assert!(!mapper.irq_pending());
            mapper.cpu_cycle(); // 0xFFFC→0xFFFD, 0xFFFD→0xFFFE
        }
        mapper.cpu_cycle(); // 0xFFFE→0xFFFF → IRQ fires
        assert!(
            mapper.irq_pending(),
            "IRQ must fire when counter reaches 0xFFFF"
        );
    }

    #[test]
    fn irq_acknowledge_clears_pending() {
        let mut mapper = make_mapper();
        // Set latch to 0xFFFE so IRQ fires after exactly 1 cpu_cycle call.
        mapper.write_prg(0x8000, 0x0E); // bits[3:0] = E
        mapper.write_prg(0x9000, 0x0F); // bits[7:4] = F
        mapper.write_prg(0xA000, 0x0F); // bits[11:8] = F
        mapper.write_prg(0xB000, 0x0F); // bits[15:12] = F → latch = 0xFFFE
        mapper.write_prg(0xC000, 0x02); // E=1: enable and reload counter to 0xFFFE
        mapper.cpu_cycle(); // 0xFFFE→0xFFFF → IRQ fires
        assert!(mapper.irq_pending());
        mapper.write_prg(0xD000, 0); // acknowledge
        assert!(!mapper.irq_pending());
    }

    /// Spec (Mesen/KS202): writing $C000 with E=1 immediately reloads counter from latch.
    #[test]
    fn irq_c000_reloads_counter_when_enabled() {
        let mut mapper = make_mapper();
        mapper.write_prg(0x8000, 5); // latch = 5
        assert_eq!(mapper.irq_counter, 0, "counter starts at 0");
        mapper.write_prg(0xC000, 0x02); // E=1: must reload counter from latch
        assert_eq!(
            mapper.irq_counter, 5,
            "counter must be reloaded from latch on E=1"
        );
    }

    /// KS202 (mapper 142 wiki): bit 0 of $C000 ("A") must NOT re-enable IRQ after
    /// acknowledge, unlike VRC3. Disabling (E=0) then re-writing E=1 is fine, but
    /// a plain $D000 acknowledge must leave the IRQ disabled.
    #[test]
    fn irq_bit0_of_c000_does_not_reenable_after_acknowledge() {
        let mut mapper = make_mapper();
        mapper.write_prg(0x8000, 0x0E); // latch = 0xFFFE
        mapper.write_prg(0x9000, 0x0F);
        mapper.write_prg(0xA000, 0x0F);
        mapper.write_prg(0xB000, 0x0F);
        // Write $C000 with bits A=1 and E=1
        mapper.write_prg(0xC000, 0x03);
        mapper.cpu_cycle(); // counter reaches 0xFFFF → IRQ fires, E cleared
        assert!(mapper.irq_pending());
        mapper.write_prg(0xD000, 0); // acknowledge
        // IRQ must remain disabled; running more cycles must NOT fire again
        for _ in 0..10 {
            mapper.cpu_cycle();
        }
        assert!(
            !mapper.irq_pending(),
            "bit0 (A) of $C000 must not re-enable IRQ after acknowledge"
        );
    }

    /// Spec: $D000 only acknowledges the IRQ; it must NOT reload the counter.
    #[test]
    fn irq_d000_does_not_reload_counter() {
        let mut mapper = make_mapper();
        mapper.write_prg(0x8000, 5); // latch = 5
        mapper.write_prg(0xC000, 0x02); // E=1: counter = 5
        mapper.cpu_cycle(); // counter = 6
        mapper.write_prg(0xD000, 0); // acknowledge — must not reload counter
        assert_eq!(mapper.irq_counter, 6, "$D000 write must not reload counter");
    }

    // -----------------------------------------------------------------------
    // PRG-RAM ($6000-$7FFF)
    // -----------------------------------------------------------------------

    #[test]
    fn prg_ram_readable_and_writable_at_6000() {
        let mut mapper = make_mapper();
        mapper.write_prg(0x6000, 0xAB);
        assert_eq!(mapper.read_prg(0x6000), 0xAB);
        mapper.write_prg(0x7FFF, 0xCD);
        assert_eq!(mapper.read_prg(0x7FFF), 0xCD);
    }

    #[test]
    fn wram_snapshot_round_trip() {
        let mut mapper = make_mapper();
        mapper.write_prg(0x6000, 0x11);
        mapper.write_prg(0x6100, 0x22);
        mapper.write_prg(0x7FFF, 0x33);
        let snapshot = mapper.wram_snapshot();
        assert_eq!(snapshot.len(), 8192);
        assert_eq!(snapshot[0x0000], 0x11);
        assert_eq!(snapshot[0x0100], 0x22);
        assert_eq!(snapshot[0x1FFF], 0x33);

        // Restore into a fresh mapper
        let mut mapper2 = make_mapper();
        mapper2.load_wram_snapshot(&snapshot);
        assert_eq!(mapper2.read_prg(0x6000), 0x11);
        assert_eq!(mapper2.read_prg(0x6100), 0x22);
        assert_eq!(mapper2.read_prg(0x7FFF), 0x33);
    }

    // ── Mapper 142 (Kaiser KS202, same hardware as mapper 56) ────────────────

    #[test]
    fn mapper_142_is_registered_in_factory() {
        let result = create_mapper(MapperContext::new_for_test(
            142,
            banked_data(8 * 1024, PRG_BANKS),
            banked_data(1024, CHR_BANKS),
            NametableLayout::Vertical,
        ));
        assert!(result.is_ok(), "Mapper 142 must be creatable via factory");
    }

    #[test]
    fn mapper_142_prg_bank_switching_works_like_mapper_56() {
        let mut mapper = create_mapper(MapperContext::new_for_test(
            142,
            banked_data(8 * 1024, PRG_BANKS),
            banked_data(1024, CHR_BANKS),
            NametableLayout::Vertical,
        ))
        .expect("mapper 142 must be registered");
        // Select bank register 1 ($A000 window), set bank 2
        mapper.write_prg(0xE000, 0x02); // select bank register 1
        mapper.write_prg(0xF000, 0x02); // set bank data = 2
        assert_eq!(
            mapper.read_prg(0xA000),
            2,
            "mapper 142 PRG bank 1 window set to bank 2"
        );
    }
}
