cge_nes 0.1.2

//! Sprite evaluation and scanline preparation.
//!
//! Evaluates OAM to determine which sprites affect the current scanline and
//! collects sprite pattern data and start positions used during rendering.
use crate::ppu::oam::OamArray;
use crate::ppu::registers::Registers;
use crate::ppu::sprite::SpriteData;
use crate::ppu::{
    oam, PpuCartMemorySpace, SpriteDataArray, SpriteStartArray, LARGE_SPRITE_HEIGHT,
    SMALL_SPRITE_HEIGHT,
};

const SECOND_PATTERN_TABLE_BASE_ADDR: u16 = 0x1000;

/// Marker for a horizontal sprite edge within a scanline used to (re)start
/// sprite evaluation at specific X positions.
///
/// The renderer maintains a sorted list of these X positions (sprite starts and
/// ends) to know when to refresh the set of sprites affecting the current pixel.
#[derive(Copy, Clone, PartialEq, Eq, Debug)]
pub struct SpriteStartForScanline {
    pub x: u8,
}

/// Calculates the base address in pattern table memory for a sprite.
///
/// # Parameters
/// * `sprite` - The OAM entry containing the sprite's tile number and attributes
/// * `regs` - PPU registers containing sprite size and pattern table selection flags
///
/// # Returns
/// The base address in pattern table memory for this sprite's tile data
pub fn sprite_base_address_in_pattern_table(sprite: oam::Entry, regs: &Registers) -> u16 {
    let large_sprites = regs.large_sprites_flag();

    if large_sprites == false {
        let addr = if regs.sprite_second_table_selected() {
            SECOND_PATTERN_TABLE_BASE_ADDR
        } else {
            0
        };

        let tile_index: u8 = sprite.tile_num().into();
        let tile_index = tile_index as u16;

        addr | (tile_index << 4)
    } else {
        let tile_index_raw: u8 = sprite.tile_num().into();
        let bank = (tile_index_raw & 0x01) as u16;
        let tile_index = (tile_index_raw & 0xFE) as u16;

        (bank * SECOND_PATTERN_TABLE_BASE_ADDR) | (tile_index << 4)
    }
}

/// Calculates the base address in pattern table memory for a background tile.
///
/// # Parameters
/// * `tile_index` - The tile number from the name table
/// * `bg_second_table_selected` - Whether the second pattern table is selected for background tiles
///
/// # Returns
/// The base address in pattern table memory for this background tile's data
pub fn bg_base_address_in_pattern_table(tile_index: u8, bg_second_table_selected: bool) -> u16 {
    let addr = if bg_second_table_selected {
        SECOND_PATTERN_TABLE_BASE_ADDR
    } else {
        0
    };

    let tile_index = tile_index as u16;

    addr | (tile_index << 4)
}

/// Applies fine Y scroll to a pattern table address.
///
/// # Parameters
/// * `addr` - Base pattern table address
/// * `local_y` - Y offset within the tile (0-7)
/// * `vertical_flip` - Whether to flip the Y coordinate vertically
///
/// # Returns
/// The adjusted pattern table address including the Y offset
pub fn set_fine_y_offset_to_pattern_table_addr(
    addr: u16,
    local_y: u8,
    vertical_flip: bool,
    large_sprites: bool,
) -> u16 {
    // Each tile in the pattern table is 16 bytes.
    // Small sprites (8x8) use one tile.
    // Large sprites (8x16) use two consecutive tiles (0 and 1).
    // We mask out the low 4 bits (0..15) which represent the byte offset within a tile (or tile pair).
    let mut addr = addr & 0b_1111_1111_1111_0000;

    // Calculate tile offset (0 or 1 for 8x16) and fine Y (0-7 within the tile).
    let (tile_offset, fine_y) = if large_sprites {
        (local_y / 8, local_y % 8)
    } else {
        (0, local_y % 8)
    };

    // Apply vertical flip if requested.
    // For 8x8: fine_y 0 becomes 7, 1 becomes 6, etc.
    // For 8x16: the entire 16-pixel block is flipped.
    let fine_y = if vertical_flip { 7 - fine_y } else { fine_y };

    let tile_offset = if vertical_flip && large_sprites {
        // If flipping 8x16, tile 0 (top) becomes tile 1 (bottom) and vice-versa.
        tile_offset ^ 1
    } else {
        tile_offset
    };

    // Add the tile offset (each tile is 16 bytes, so we shift by 4).
    addr += (tile_offset as u16) << 4;
    // OR in the fine Y offset (0-7).
    addr | (fine_y as u16)
}

/// Reads the color index bytes for a pattern table row.
///
/// # Parameters
/// * `row_addr` - Address of the pattern table row
/// * `cart_mem` - Memory interface for reading pattern table data
///
/// # Returns
/// A tuple of (high byte, low byte) containing the color pattern data
pub fn get_color_index_row_bytes_from_pattern_table(
    row_addr: u16,
    cart_mem: &mut impl PpuCartMemorySpace,
) -> (u8, u8) {
    let low_byte = cart_mem.read(row_addr);
    let high_byte = cart_mem.read(row_addr | 0b_1000);

    (high_byte, low_byte)
}

/// Prepares sprite data for rendering the next scanline.
///
/// # Parameters
/// * `scanline_y` - The Y coordinate of the scanline to prepare
/// * `large_sprites` - Whether 8x16 sprites are enabled
/// * `oam` - Object Attribute Memory containing sprite data
/// * `sprites` - Array to store the processed sprite data
///
/// # Returns
/// `true` if sprite overflow occurred (more than 8 sprites on the scanline)
pub fn prepare_sprites_for_scanline(
    scanline_y: u8,
    large_sprites: bool,
    oam: &OamArray,
    sprites: &mut SpriteDataArray,
) -> bool {
    let mut sprite_overflow = false;

    // Determine sprite height based on size flag
    let sprite_height = if large_sprites {
        LARGE_SPRITE_HEIGHT
    } else {
        SMALL_SPRITE_HEIGHT
    };

    sprites.clear();
    for (i, entry) in oam.iter().enumerate() {
        // Check if sprite is visible on this scanline
        if scanline_y > (entry.y()) {
            if scanline_y < (entry.y().saturating_add(1 + sprite_height)) {
                // Try to add sprite if we haven't hit the per-scanline limit
                if sprites.len() < sprites.capacity() {
                    sprites.push(SpriteData::with_oam_index(entry, i as u8));
                } else {
                    // Set overflow flag if we exceed sprite limit. The PPU
                    // keeps scanning OAM after finding 8 in-range sprites
                    // (so a later out-of-range sprite can stop the scan),
                    // but on a real NES the overflow flag is *set* as soon
                    // as the 9th in-range sprite is found.
                    sprite_overflow = true;
                }
            }
        }
    }

    sprite_overflow
}

/// Sorts sprite data by X coordinate for a scanline.
///
/// # Parameters
/// * `sprites` - Array of sprite data to process
/// * `sorted_sprites` - Output array for sorted sprite positions
pub fn sort_sprites_by_col(sprites: &[SpriteData], sorted_sprites: &mut SpriteStartArray) {
    let sprite_count = sprites.len();
    assert!(sprite_count <= sorted_sprites.capacity());

    sorted_sprites.clear();
    for sprite in sprites {
        // Add sprite start X position if not already present
        // SAFETY: Asserted at the beginning that we have at most 16 elems to push
        // and the arrayvec have just been cleared
        let sprite_start = SpriteStartForScanline::new(sprite.obj_attributes().x());
        if !sorted_sprites.contains(&sprite_start) {
            unsafe {
                sorted_sprites.push_unchecked(sprite_start);
            }
        }

        // Add sprite end X position if not already present
        let sprite_end = SpriteStartForScanline::new(sprite.obj_attributes().x().saturating_add(8));
        if !sorted_sprites.contains(&sprite_end) {
            unsafe {
                sorted_sprites.push_unchecked(sprite_end);
            }
        }
    }

    sort_sprites_for_scanline(sorted_sprites);
}

/// Decodes the 2-bit palette index from an attribute table byte.
///
/// Attribute bytes select a palette for each 2x2 tile block. Given the tile
/// coordinates within the nametable, this function extracts the correct 2-bit
/// palette index for that quadrant.
///
/// # Parameters
/// * `coords_in_table` - (coarse_x, coarse_y) tile coordinates within the 32x30 nametable
/// * `attribute_byte` - Raw attribute byte covering a 4x4 tile area (2x2 blocks)
///
/// # Returns
/// The palette index in the range 0..=3 for the tile at `coords_in_table`.
pub fn palette_index_in_name_table3(coords_in_table: (u16, u16), attribute_byte: u8) -> u8 {
    let x_in_attr = (coords_in_table.0 / 2) % 2;
    let y_in_attr = (coords_in_table.1 / 2) % 2;

    let attribute_byte = attribute_byte >> (x_in_attr * 2 + y_in_attr * 4);

    attribute_byte & 0b_11
}

/// Sorts an array of sprite start positions by X coordinate.
///
/// Uses an insertion sort algorithm for small arrays (<=8 elements).
///
/// # Parameters
/// * `slice` - Array of sprite positions to sort in-place
fn sort_sprites_for_scanline(slice: &mut [SpriteStartForScanline]) {
    for index_to_insert in 1..slice.len() {
        let mut i = index_to_insert;
        // SAFETY: We know that all indices are always going to be less than slice.len() and >= 0
        unsafe {
            while (i > 0) && (slice.get_unchecked(i - 1).x > slice.get_unchecked(i).x) {
                let tmp = *slice.get_unchecked(i);
                *slice.get_unchecked_mut(i) = *slice.get_unchecked(i - 1);
                *slice.get_unchecked_mut(i - 1) = tmp;
                // SAFETY: non-wrapping because i > 0
                i -= 1;
            }
        }
    }
}

impl SpriteStartForScanline {
    /// Creates a new sprite edge marker at the given X coordinate.
    ///
    /// Used to identify where sprite coverage on a scanline starts or ends so
    /// the renderer can refresh the set of active sprites at those X positions.
    fn new(x: u8) -> Self {
        Self { x }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ppu::Registers;
    use arrayvec::ArrayVec;

    struct MockMemorySpace;

    impl PpuCartMemorySpace for MockMemorySpace {
        fn read(&mut self, addr: u16) -> u8 {
            if addr < 0x2000 {
                if (addr & 0b_1000) == 0 {
                    // lower plane
                    return 0x00;
                } else {
                    // higher plane
                    return 0xFF;
                }
            }

            // Outside pattern table
            0xCA
        }

        fn write(&mut self, _data: u8, _addr: u16) {}
    }

    #[test]
    fn base_addresses_correct_small_sprites() {
        for tile in 0..u8::MAX {
            let mut sprite = oam::Entry::default();
            sprite.set_tile_num(tile.into());

            let mut registers = Registers::default();

            let base_addr = sprite_base_address_in_pattern_table(sprite, &registers);
            let expected_address: u8 = sprite.tile_num().into();
            let expected_address: u16 = (expected_address as u16) << 4;
            assert_eq!(base_addr, expected_address);

            registers.set_sprite_second_table_selected();

            let base_addr = sprite_base_address_in_pattern_table(sprite, &registers);
            let expected_address = expected_address | SECOND_PATTERN_TABLE_BASE_ADDR;
            assert_eq!(base_addr, expected_address);
        }
    }

    #[test]
    fn fine_offset_no_flip() {
        const FINE_Y_MASK: u16 = 0b_0000_0000_0000_0111;
        const BASE_ADDR_MASK: u16 = 0b_1111_1111_1111_1000;

        for tile in 0..u8::MAX {
            for y in 0..8 {
                let mut sprite = oam::Entry::default();
                sprite.set_tile_num(tile.into());

                let registers = Registers::default();

                let base_addr = sprite_base_address_in_pattern_table(sprite, &registers);
                let final_addr =
                    set_fine_y_offset_to_pattern_table_addr(base_addr, y, false, false);

                assert_eq!(final_addr & BASE_ADDR_MASK, base_addr);
                assert_eq!(final_addr & FINE_Y_MASK, y as u16);
            }
        }
    }

    #[test]
    fn fine_offset_flip() {
        const FINE_Y_MASK: u16 = 0b_0000_0000_0000_0111;
        const BASE_ADDR_MASK: u16 = 0b_1111_1111_1111_1000;

        for tile in 0..u8::MAX {
            for y in 0..8 {
                let mut sprite = oam::Entry::default();
                sprite.set_tile_num(tile.into());

                let registers = Registers::default();

                let base_addr = sprite_base_address_in_pattern_table(sprite, &registers);
                let final_addr = set_fine_y_offset_to_pattern_table_addr(base_addr, y, true, false);
                let expected_flipped_y = 7 - y as u16;

                assert_eq!(final_addr & BASE_ADDR_MASK, base_addr);
                assert_eq!(final_addr & FINE_Y_MASK, expected_flipped_y);
            }
        }
    }

    #[test]
    fn get_color_index() {
        let mut cart = MockMemorySpace;

        for addr in 0..0x2000 {
            let color_indices = get_color_index_row_bytes_from_pattern_table(addr, &mut cart);

            if (addr & 0b_1000) == 0 {
                assert_eq!(color_indices, (0xFF, 0x00));
            } else {
                assert_eq!(color_indices, (0xFF, 0xFF));
            }
        }
    }

    #[test]
    fn sort() {
        let mut v = vec![
            SpriteStartForScanline::new(33),
            SpriteStartForScanline::new(159),
            SpriteStartForScanline::new(82),
            SpriteStartForScanline::new(241),
            SpriteStartForScanline::new(107),
            SpriteStartForScanline::new(47),
            SpriteStartForScanline::new(84),
            SpriteStartForScanline::new(6),
        ];

        sort_sprites_for_scanline(&mut v[..]);

        let expected_v = vec![
            SpriteStartForScanline::new(6),
            SpriteStartForScanline::new(33),
            SpriteStartForScanline::new(47),
            SpriteStartForScanline::new(82),
            SpriteStartForScanline::new(84),
            SpriteStartForScanline::new(107),
            SpriteStartForScanline::new(159),
            SpriteStartForScanline::new(241),
        ];

        assert_eq!(v.len(), expected_v.len());
        for (elem, expected_elem) in v.into_iter().zip(expected_v) {
            assert_eq!(elem, expected_elem);
        }
    }

    #[test]
    fn sorting_sprites_by_col() {
        let mut v = vec![
            SpriteData::with_oam_index(oam::Entry::with_data(0, 0, Default::default(), 33), 0),
            SpriteData::with_oam_index(oam::Entry::with_data(1, 0, Default::default(), 87), 1),
            SpriteData::with_oam_index(oam::Entry::with_data(2, 0, Default::default(), 82), 2),
            SpriteData::with_oam_index(oam::Entry::with_data(3, 0, Default::default(), 251), 3),
            SpriteData::with_oam_index(oam::Entry::with_data(4, 0, Default::default(), 107), 4),
            SpriteData::with_oam_index(oam::Entry::with_data(5, 0, Default::default(), 47), 5),
            SpriteData::with_oam_index(oam::Entry::with_data(6, 0, Default::default(), 84), 6),
            SpriteData::with_oam_index(oam::Entry::with_data(7, 0, Default::default(), 6), 7),
        ];

        v[2].set_draw_front();
        v[4].set_draw_front();
        v[5].set_draw_front();

        let mut sorted = ArrayVec::new();

        sort_sprites_by_col(&v[..], &mut sorted);

        let sorted_ordered = vec![
            SpriteStartForScanline::new(6),
            SpriteStartForScanline::new(14),
            SpriteStartForScanline::new(33),
            SpriteStartForScanline::new(41),
            SpriteStartForScanline::new(47),
            SpriteStartForScanline::new(55),
            SpriteStartForScanline::new(82),
            SpriteStartForScanline::new(84),
            SpriteStartForScanline::new(87),
            SpriteStartForScanline::new(90),
            SpriteStartForScanline::new(92),
            SpriteStartForScanline::new(95),
            SpriteStartForScanline::new(107),
            SpriteStartForScanline::new(115),
            SpriteStartForScanline::new(251),
            SpriteStartForScanline::new(255),
        ];

        assert_eq!(sorted.len(), sorted_ordered.len());
        for (elem, expected_elem) in sorted.into_iter().zip(sorted_ordered) {
            assert_eq!(elem, expected_elem);
        }
    }

    struct MockPpuMemPatternTable;

    impl PpuCartMemorySpace for MockPpuMemPatternTable {
        fn read(&mut self, addr: u16) -> u8 {
            if addr == 0b_0001_1010_0101_0010 {
                // Second table, tile row 10, tile col 5, pixel row 2, lower plane
                0xFF
            } else if addr == 0b_0001_1010_0101_1010 {
                // Second table, tile row 10, tile col 5, pixel row 2, higher plane
                0xAA
            } else {
                0
            }
        }

        fn write(&mut self, _data: u8, _addr: u16) {
            unreachable!()
        }
    }

    #[test]
    fn color_index_row_bytes_from_pattern_table() {
        let mut mem = MockPpuMemPatternTable;

        let row_addr = 0b_0001_1010_0101_0010;

        let colors = get_color_index_row_bytes_from_pattern_table(row_addr, &mut mem);

        let expected_colors = (0xAA, 0xFF);

        assert_eq!(colors, expected_colors);
    }

    #[test]
    fn fine_offset_large_sprites() {
        // Test 8x16 sprites (large_sprites = true)
        let base_addr = 0x0000; // Bank 0, Tile 0

        // Normal (no flip)
        // Top tile (local_y = 0..8)
        assert_eq!(
            set_fine_y_offset_to_pattern_table_addr(base_addr, 0, false, true),
            0x0000
        );
        assert_eq!(
            set_fine_y_offset_to_pattern_table_addr(base_addr, 7, false, true),
            0x0007
        );
        // Bottom tile (local_y = 8..16)
        assert_eq!(
            set_fine_y_offset_to_pattern_table_addr(base_addr, 8, false, true),
            0x0010
        );
        assert_eq!(
            set_fine_y_offset_to_pattern_table_addr(base_addr, 15, false, true),
            0x0017
        );

        // Vertical flip
        // When flipped, local_y = 0 should be the bottom-most row of the bottom tile
        // local_y = 0 -> tile_offset = 1, fine_y = 7
        assert_eq!(
            set_fine_y_offset_to_pattern_table_addr(base_addr, 0, true, true),
            0x0017
        );
        // local_y = 7 -> tile_offset = 1, fine_y = 0
        assert_eq!(
            set_fine_y_offset_to_pattern_table_addr(base_addr, 7, true, true),
            0x0010
        );
        // local_y = 8 -> tile_offset = 0, fine_y = 7
        assert_eq!(
            set_fine_y_offset_to_pattern_table_addr(base_addr, 8, true, true),
            0x0007
        );
        // local_y = 15 -> tile_offset = 0, fine_y = 0
        assert_eq!(
            set_fine_y_offset_to_pattern_table_addr(base_addr, 15, true, true),
            0x0000
        );
    }

    #[test]
    fn test_fine_y_mapping_small() {
        let base = 0x200u16;
        // local_y: 0 and 7
        let cases = [(0u8, 0u16), (7u8, 7u16)];
        for (local_y, expected_no) in cases {
            let got_no = set_fine_y_offset_to_pattern_table_addr(base, local_y, false, false);
            let expected = base | expected_no;
            assert_eq!(got_no, expected, "no-flip small sprite local_y={}", local_y);

            let got_flip = set_fine_y_offset_to_pattern_table_addr(base, local_y, true, false);
            let expected_flip = base | ((7u8 - local_y) as u16);
            assert_eq!(
                got_flip, expected_flip,
                "flip small sprite local_y={}",
                local_y
            );
        }
    }

    #[test]
    fn test_fine_y_mapping_large() {
        let base = 0x200u16;
        // local_y cases: (local_y, expected_tile_offset, expected_fine_y)
        let cases = [
            (0u8, 0u16, 0u16),
            (7u8, 0u16, 7u16),
            (8u8, 1u16, 0u16),
            (15u8, 1u16, 7u16),
        ];
        for (local_y, tile_off, fine) in cases {
            let got_no = set_fine_y_offset_to_pattern_table_addr(base, local_y, false, true);
            let expected_no = base + (tile_off << 4) | fine;
            assert_eq!(got_no, expected_no, "no-flip large local_y={}", local_y);

            let got_flip = set_fine_y_offset_to_pattern_table_addr(base, local_y, true, true);
            // flipped: tile_offset' = 1 - tile_off, fine' = 7 - fine
            let expected_flip = base + ((1u16 - tile_off) << 4) | ((7u16 - fine) as u16);
            assert_eq!(got_flip, expected_flip, "flip large local_y={}", local_y);
        }
    }
}