//! EDRAM tile book + render-target key bookkeeping.
//!
//! Mirrors `xenia-canary/src/xenia/gpu/render_target_cache.h` at the data-
//! structure level. Xenos's 10 MiB EDRAM is divided into 2048 "tiles" of
//! 80×16 samples each; render targets claim a contiguous range of those
//! tiles based on `(base_tiles, pitch_tiles_at_32bpp, msaa_samples, format,
//! is_depth)`. Two render targets with overlapping tile ranges share the
//! underlying EDRAM — canary tracks this with per-tile "Host vs Shared"
//! ownership, which is what this module's `TileOwner` captures.
//!
//! P4 ships the **bookkeeping**. Actual host texture allocation per key (so
//! the host can draw into a wgpu texture matching the guest's RT) is left to
//! a future host-side cache built on top of this module; the same for
//! format-conversion compute shaders (the plan's P5 territory).

use std::collections::HashMap;

/// Number of EDRAM tiles on Xenos. Matches canary's `xenos::kEdramTileCount`.
pub const EDRAM_TILE_COUNT: usize = 2048;

/// MSAA sample count encoded into [`RenderTargetKey`]. Canary uses this as
/// `xenos::MsaaSamples` (1×/2×/4×).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum MsaaSamples {
    X1 = 0,
    X2 = 1,
    X4 = 2,
}

impl MsaaSamples {
    pub fn from_raw(raw: u32) -> Self {
        match raw & 0x3 {
            1 => MsaaSamples::X2,
            2 => MsaaSamples::X4,
            _ => MsaaSamples::X1,
        }
    }
    pub fn count(self) -> u32 {
        1u32 << (self as u32)
    }
}

/// The packed EDRAM render-target identity. Bit layout matches
/// `render_target_cache.h:251-321`'s `RenderTargetKey` union (26 bits used,
/// stored as a single `u32` so it hashes cheaply). `pitch_tiles_at_32bpp`
/// is always the 32bpp-equivalent pitch — 64bpp targets halve their tile
/// pitch from the nominal tile grid (canary's `GetPitchTiles()` handles
/// that).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct RenderTargetKey {
    pub base_tiles: u16,              // [0..2048)
    pub pitch_tiles_at_32bpp: u16,    // 0..=256 in practice
    pub msaa_samples: MsaaSamples,
    pub is_depth: bool,
    /// Color format: `xenos::ColorRenderTargetFormat` when !is_depth.
    /// Depth format: `xenos::DepthRenderTargetFormat` when is_depth.
    pub resource_format: u8,          // 4 bits
}

impl RenderTargetKey {
    /// Pack into canary's 26-bit layout. Useful for compact storage /
    /// hashing when we add a LRU cache later on.
    pub fn pack(&self) -> u32 {
        (self.base_tiles as u32 & 0x7FF)
            | (((self.pitch_tiles_at_32bpp as u32) & 0xFF) << 11)
            | (((self.msaa_samples as u32) & 0x3) << 19)
            | ((self.is_depth as u32) << 21)
            | (((self.resource_format as u32) & 0xF) << 22)
    }

    pub fn unpack(raw: u32) -> Self {
        Self {
            base_tiles: (raw & 0x7FF) as u16,
            pitch_tiles_at_32bpp: ((raw >> 11) & 0xFF) as u16,
            msaa_samples: MsaaSamples::from_raw((raw >> 19) & 0x3),
            is_depth: ((raw >> 21) & 1) != 0,
            resource_format: ((raw >> 22) & 0xF) as u8,
        }
    }

    /// How many EDRAM tiles the whole surface occupies (rough estimate; a
    /// real height-aware calc needs viewport info). We conservatively use
    /// `pitch_tiles_at_32bpp * 1` until a draw tells us otherwise; callers
    /// that know the height can call [`tile_footprint_with_height`].
    pub fn tile_pitch(&self) -> u16 {
        // 64bpp formats pack two 32bpp tiles into one 64bpp tile.
        if self.is_64bpp() {
            self.pitch_tiles_at_32bpp / 2
        } else {
            self.pitch_tiles_at_32bpp
        }
    }

    pub fn is_64bpp(&self) -> bool {
        if self.is_depth {
            false
        } else {
            // Canary: `ColorRenderTargetFormat::{k_16_16_16_16,
            // k_16_16_16_16_FLOAT, k_32_32_FLOAT}` are 64bpp; indices 4, 5, 7
            // in the enum. (Kept narrow because the enum is 4 bits wide.)
            matches!(self.resource_format, 4 | 5 | 7)
        }
    }

    /// Tiles claimed by this RT if its surface height is `rows_of_tiles`
    /// (i.e. `ceil(height_in_samples / 16)`).
    pub fn tile_footprint_with_height(&self, rows_of_tiles: u16) -> u16 {
        self.tile_pitch().saturating_mul(rows_of_tiles)
    }
}

/// Who currently owns a tile of EDRAM.
///
/// `None`: untouched; free to claim.
/// `Host(idx)`: a single RT has exclusive ownership.
/// `Shared(idx)`: two+ RT keys map to the same tile (usually after a
/// format change without an intervening clear); the named RT is the most
/// recent owner whose format should be honored for readback.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[derive(Default)]
pub enum TileOwner {
    #[default]
    None,
    Host(u32),
    Shared(u32),
}


/// Bookkeeping across the 2048 EDRAM tiles. Not a GPU resource by itself —
/// tracks which render target (by index) currently owns each tile.
pub struct EdramTileBook {
    tiles: Vec<TileOwner>,
}

impl Default for EdramTileBook {
    fn default() -> Self {
        Self::new()
    }
}

impl EdramTileBook {
    pub fn new() -> Self {
        Self {
            tiles: vec![TileOwner::None; EDRAM_TILE_COUNT],
        }
    }

    pub fn who_owns(&self, tile: u16) -> TileOwner {
        self.tiles
            .get(tile as usize)
            .copied()
            .unwrap_or(TileOwner::None)
    }

    /// Mark `[base, base+count)` as owned by `rt_idx`. Pre-existing owners
    /// in the range are demoted to `Shared` (format reinterpretation).
    /// Returns the number of tiles newly claimed (not previously the same
    /// owner).
    pub fn claim(&mut self, base: u16, count: u16, rt_idx: u32) -> u32 {
        let mut newly_claimed = 0u32;
        for i in 0..(count as usize) {
            let t = base as usize + i;
            if t >= self.tiles.len() {
                break;
            }
            let prev = self.tiles[t];
            let already_ours = matches!(
                prev,
                TileOwner::Host(idx) | TileOwner::Shared(idx) if idx == rt_idx
            );
            match prev {
                TileOwner::None => {
                    self.tiles[t] = TileOwner::Host(rt_idx);
                }
                TileOwner::Host(idx) if idx == rt_idx => {
                    // re-claim of same RT — no-op
                }
                _ => {
                    // Format change / shared range.
                    self.tiles[t] = TileOwner::Shared(rt_idx);
                }
            }
            if !already_ours {
                newly_claimed += 1;
            }
        }
        newly_claimed
    }

    /// Drop `rt_idx` from any tile it owns; tiles revert to `None` unless
    /// they were `Shared(rt_idx)` (in which case they also revert to
    /// `None`; the other sharer's ownership is lost — `release` is a
    /// coarse "this RT is gone" operation).
    pub fn release(&mut self, rt_idx: u32) {
        for t in self.tiles.iter_mut() {
            if matches!(
                *t,
                TileOwner::Host(idx) | TileOwner::Shared(idx) if idx == rt_idx
            ) {
                *t = TileOwner::None;
            }
        }
    }

    /// Count tiles currently assigned to any RT (Host or Shared).
    pub fn occupied_count(&self) -> u32 {
        self.tiles
            .iter()
            .filter(|o| !matches!(o, TileOwner::None))
            .count() as u32
    }
}

/// Minimal per-RT descriptor stored alongside the tile book. P5's texture
/// cache will expand this with the actual wgpu texture handle.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RtDescriptor {
    pub key: RenderTargetKey,
    /// Number of times this key has been bound since creation. Rough
    /// proxy for activity / hot-RT identification.
    pub bind_count: u32,
    /// Draw index on first bind — handy for debugging divergence.
    pub first_draw_index: u32,
}

/// Top-level cache: maps packed keys to small descriptors + the tile book.
pub struct RenderTargetCache {
    next_idx: u32,
    by_key: HashMap<u32, u32>,
    descriptors: HashMap<u32, RtDescriptor>,
    pub tiles: EdramTileBook,
}

impl Default for RenderTargetCache {
    fn default() -> Self {
        Self::new()
    }
}

impl RenderTargetCache {
    pub fn new() -> Self {
        Self {
            next_idx: 0,
            by_key: HashMap::new(),
            descriptors: HashMap::new(),
            tiles: EdramTileBook::new(),
        }
    }

    /// Look up or allocate an RT descriptor for `key`. `draw_index` is the
    /// current monotonic draw counter — recorded on first insert for
    /// provenance.
    pub fn bind(&mut self, key: RenderTargetKey, draw_index: u32) -> u32 {
        let packed = key.pack();
        if let Some(&idx) = self.by_key.get(&packed) {
            if let Some(d) = self.descriptors.get_mut(&idx) {
                d.bind_count += 1;
            }
            return idx;
        }
        let idx = self.next_idx;
        self.next_idx += 1;
        self.by_key.insert(packed, idx);
        self.descriptors.insert(
            idx,
            RtDescriptor {
                key,
                bind_count: 1,
                first_draw_index: draw_index,
            },
        );
        idx
    }

    pub fn descriptor(&self, idx: u32) -> Option<&RtDescriptor> {
        self.descriptors.get(&idx)
    }

    pub fn len(&self) -> usize {
        self.descriptors.len()
    }

    pub fn is_empty(&self) -> bool {
        self.descriptors.is_empty()
    }

    /// Claim tiles for the descriptor at `rt_idx`. `height_tiles` is
    /// `ceil(viewport_height_samples / 16)` — callers supply it because
    /// the key itself doesn't carry height.
    pub fn claim_tiles(&mut self, rt_idx: u32, height_tiles: u16) -> u32 {
        if let Some(d) = self.descriptors.get(&rt_idx) {
            let footprint = d.key.tile_footprint_with_height(height_tiles);
            self.tiles.claim(d.key.base_tiles, footprint, rt_idx)
        } else {
            0
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn render_target_key_pack_roundtrip() {
        let k = RenderTargetKey {
            base_tiles: 1600,
            pitch_tiles_at_32bpp: 80,
            msaa_samples: MsaaSamples::X4,
            is_depth: true,
            resource_format: 0b1010,
        };
        let packed = k.pack();
        let round = RenderTargetKey::unpack(packed);
        assert_eq!(round, k);
    }

    #[test]
    fn tile_book_claim_marks_owners() {
        let mut book = EdramTileBook::new();
        assert_eq!(book.occupied_count(), 0);
        let new_count = book.claim(100, 10, 42);
        assert_eq!(new_count, 10);
        assert_eq!(book.who_owns(100), TileOwner::Host(42));
        assert_eq!(book.who_owns(109), TileOwner::Host(42));
        assert_eq!(book.who_owns(110), TileOwner::None);
    }

    #[test]
    fn tile_book_claim_demotes_to_shared() {
        let mut book = EdramTileBook::new();
        book.claim(100, 10, 1);
        book.claim(105, 10, 2);
        // Overlap: tiles 105..110 should be Shared(2); 100..105 stay Host(1);
        // tiles 110..115 are fresh Host(2).
        assert_eq!(book.who_owns(104), TileOwner::Host(1));
        assert_eq!(book.who_owns(105), TileOwner::Shared(2));
        assert_eq!(book.who_owns(110), TileOwner::Host(2));
    }

    #[test]
    fn tile_book_release_frees_all() {
        let mut book = EdramTileBook::new();
        book.claim(0, 50, 7);
        book.release(7);
        assert_eq!(book.occupied_count(), 0);
    }

    #[test]
    fn rt_cache_bind_is_idempotent_by_key() {
        let mut cache = RenderTargetCache::new();
        let k = RenderTargetKey {
            base_tiles: 0,
            pitch_tiles_at_32bpp: 80,
            msaa_samples: MsaaSamples::X1,
            is_depth: false,
            resource_format: 0,
        };
        let a = cache.bind(k, 0);
        let b = cache.bind(k, 1);
        assert_eq!(a, b);
        let d = cache.descriptor(a).unwrap();
        assert_eq!(d.bind_count, 2);
        assert_eq!(d.first_draw_index, 0);
    }

    #[test]
    fn rt_cache_claim_tiles_tracks_footprint() {
        let mut cache = RenderTargetCache::new();
        let k = RenderTargetKey {
            base_tiles: 0,
            pitch_tiles_at_32bpp: 80, // 32bpp 1280-wide target
            msaa_samples: MsaaSamples::X1,
            is_depth: false,
            resource_format: 0,
        };
        let idx = cache.bind(k, 0);
        // 720 samples tall / 16 per tile = 45 rows → 80 * 45 = 3600 tiles;
        // caps out at 2048. Verify clamping.
        let newly = cache.claim_tiles(idx, 45);
        assert_eq!(newly, 2048);
        assert_eq!(cache.tiles.occupied_count(), 2048);
    }
}
