dorf/src/world/tiles/tilemap_chunk.rs

use bevy::{
    prelude::*,
    sprite_render::{AlphaMode2d, TileData, TilemapChunk, TilemapChunkTileData},
};
use bevy_platform::collections::{HashMap, HashSet};
use bevy_platform::time::Instant;

use crate::constants::{ITILE_SIZE, TILE_PIXELS, TILE_SIZE};
use crate::game::ZIndex;
use crate::world::chunks::{CHUNK_SIZE, CHUNK_SIZE_TILE, Z_BELOW, Z_TOTAL};
use crate::world::textures::TilemapTileset;
use crate::world::tiles::{FloorTileData, TileMap};

pub const LAYER_FLOOR: u8 = 0;

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Component)]
pub struct ChunkLayerKey {
    pub chunk_pos: IVec2,
    pub z_index: usize,
    pub layer: u8,
}

#[derive(Resource)]
pub struct TilemapChunkRegistry {
    pub entities: HashMap<ChunkLayerKey, Entity>,
    pub dirty_keys: HashSet<ChunkLayerKey>,
    pub z_change_keys: HashSet<ChunkLayerKey>,
}

impl Default for TilemapChunkRegistry {
    fn default() -> Self {
        Self {
            entities: Default::default(),
            dirty_keys: Default::default(),
            z_change_keys: Default::default(),
        }
    }
}

#[derive(Resource, Default)]
pub struct TilemapChunkStates {
    pub solid_bits: HashMap<ChunkLayerKey, u64>,
    pub underground_bits: HashMap<ChunkLayerKey, u64>,
}

#[derive(Resource)]
pub struct TilemapChunkSpawner {
    pub pending: Vec<ChunkLayerKey>,
    pub started: bool,
}

impl Default for TilemapChunkSpawner {
    fn default() -> Self {
        Self {
            pending: Vec::new(),
            started: false,
        }
    }
}

pub const MAX_POPULATE_PER_FRAME: usize = 512;

#[derive(Resource)]
pub struct PreviousZIndex(pub i32);

impl Default for PreviousZIndex {
    fn default() -> Self {
        Self(0)
    }
}

impl TilemapChunkSpawner {
    pub fn queue_chunk(&mut self, chunk_pos: IVec2) {
        for z in 0..=(Z_TOTAL as usize) {
            let key = ChunkLayerKey {
                chunk_pos,
                z_index: z,
                layer: LAYER_FLOOR,
            };
            if !self.pending.iter().any(|k| k == &key) {
                self.pending.push(key);
            }
        }
    }
}

pub const FIXTURE_ROW_OFFSET: u16 = 6;
pub const LAYER_ALPHA: f32 = 41_f32 / 255_f32;

fn tile_id_to_tileset_index(tile_id: u32) -> Option<u16> {
    if tile_id == 0 {
        return Some(0);
    }
    if tile_id <= 5 {
        return Some(tile_id as u16);
    }
    let fixture_id = tile_id.saturating_sub(crate::world::textures::FIXTURE_ID_OFFSET);
    if fixture_id >= 1 && fixture_id <= 5 {
        Some(FIXTURE_ROW_OFFSET + fixture_id as u16 - 1)
    } else {
        None
    }
}

fn tile_for_depth(real_index: u16, z_diff: i32) -> TileData {
    const SKY_R: f32 = 133_f32 / 255_f32;
    const SKY_G: f32 = 167_f32 / 255_f32;
    const SKY_B: f32 = 178_f32 / 255_f32;

    if z_diff <= 0 {
        return TileData {
            tileset_index: real_index,
            color: Color::WHITE,
            visible: true,
        };
    }

    let n = (z_diff as f32).min(8_f32);
    let t = 1_f32 - (1_f32 - LAYER_ALPHA).powf(n);

    if t >= 0.98 {
        return TileData {
            tileset_index: 0,
            color: Color::srgb(SKY_R, SKY_G, SKY_B),
            visible: true,
        };
    }

    let r = (1_f32 - t) + SKY_R * t;
    let g = (1_f32 - t) + SKY_G * t;
    let b = (1_f32 - t) + SKY_B * t;
    TileData {
        tileset_index: real_index,
        color: Color::srgb(r, g, b),
        visible: true,
    }
}

fn is_tile_visible_at_z(tile_data: &FloorTileData, z_index: usize) -> bool {
    if z_index >= 256 {
        return false;
    }
    let word = z_index / 32;
    let bit = z_index % 32;
    (tile_data.visible_range[word] & (1 << bit)) != 0
}

fn camera_z_for_z_index(z_index: usize) -> i32 {
    z_index as i32 - Z_BELOW as i32
}

fn populate_chunk_tiles(
    tilemap: &TileMap,
    chunk_pos: IVec2,
    z_index: usize,
    camera_z: i32,
) -> (Vec<Option<TileData>>, u64) {
    let z_diff = i32::max(camera_z - camera_z_for_z_index(z_index), 0);
    let mut tiles = Vec::with_capacity((CHUNK_SIZE * CHUNK_SIZE) as usize);
    let mut underground_bits: u64 = 0;

    for local_y in 0..CHUNK_SIZE {
        for local_x in 0..CHUNK_SIZE {
            let slot_idx = local_y * CHUNK_SIZE + local_x;
            let world_pos = IVec3::new(
                chunk_pos.x * CHUNK_SIZE_TILE + local_x * ITILE_SIZE,
                chunk_pos.y * CHUNK_SIZE_TILE + local_y * ITILE_SIZE,
                camera_z_for_z_index(z_index) * ITILE_SIZE,
            );
            let tile = tilemap.get_floor(&world_pos);
            match tile {
                Some(t) => {
                    let idx = tile_id_to_tileset_index(t.id as u32);
                    if let Some(tileset_idx) = idx {
                        let is_visible = is_tile_visible_at_z(t, z_index);
                        let is_underground = !is_visible && t.id != 0;
                        if is_underground {
                            underground_bits |= 1u64 << slot_idx;
                        }
                        let tile_data = if is_visible {
                            tile_for_depth(tileset_idx, z_diff)
                        } else if t.id == 0 {
                            tile_for_depth(0, z_diff)
                        } else {
                            TileData {
                                tileset_index: 0,
                                color: Color::BLACK,
                                visible: true,
                            }
                        };
                        tiles.push(Some(tile_data));
                    } else {
                        tiles.push(None);
                    }
                }
                _ => tiles.push(None),
            }
        }
    }
    (tiles, underground_bits)
}

fn recolor_chunk_for_depth(
    tile_data: &mut TilemapChunkTileData,
    z_diff: i32,
    solid_bits: u64,
    underground_bits: u64,
) {
    for (i, slot) in tile_data.0.iter_mut().enumerate() {
        let Some(td) = slot else { continue };
        let is_underground = (underground_bits >> i) & 1 == 1;
        td.color = if is_underground {
            Color::BLACK
        } else {
            tile_fade_color(td.tileset_index, z_diff)
        };
    }
}

fn tile_fade_color(tileset_idx: u16, z_diff: i32) -> Color {
    if z_diff <= 0 {
        return Color::WHITE;
    }
    let n = (z_diff as f32).min(8_f32);
    let t = 1_f32 - (1_f32 - LAYER_ALPHA).powf(n);
    let sky_r = 133_f32 / 255_f32;
    let sky_g = 167_f32 / 255_f32;
    let sky_b = 178_f32 / 255_f32;
    if t >= 0.98 || tileset_idx == 0 {
        return Color::srgb(sky_r, sky_g, sky_b);
    }
    let r = (1_f32 - t) + sky_r * t;
    let g = (1_f32 - t) + sky_g * t;
    let b = (1_f32 - t) + sky_b * t;
    Color::srgb(r, g, b)
}

pub fn spawn_tilemap_chunks(
    mut commands: Commands,
    tilemap: Res<TileMap>,
    tileset: Res<TilemapTileset>,
    mut registry: ResMut<TilemapChunkRegistry>,
    mut states: ResMut<TilemapChunkStates>,
    mut spawner: ResMut<TilemapChunkSpawner>,
    z_index: Res<ZIndex>,
    mut bench: ResMut<super::TilemapBenchmark>,
) {
    if spawner.pending.is_empty() {
        if !spawner.started && !registry.entities.is_empty() {
            registry.entities.clear();
            registry.dirty_keys.clear();
            registry.z_change_keys.clear();
            states.solid_bits.clear();
            states.underground_bits.clear();
            spawner.started = true;
        }
        return;
    }
    spawner.started = true;
    let pending_count = spawner.pending.len();
    let now = Instant::now();
    let camera_z = z_index.0 as i32;
    let mut spawned_this_call = 0usize;

    registry.entities.reserve(pending_count);

    for key in spawner.pending.drain(..) {
        let z_diff = i32::max(camera_z - camera_z_for_z_index(key.z_index), 0);

        let (tile_data, underground_bits) = if key.layer == LAYER_FLOOR {
            let (tiles, ub) = populate_chunk_tiles(&tilemap, key.chunk_pos, key.z_index, camera_z);
            let solid_bits = compute_solid_bits(&tiles);
            states.solid_bits.insert(key, solid_bits);
            states.underground_bits.insert(key, ub);
            (tiles, ub)
        } else {
            (vec![None; (CHUNK_SIZE * CHUNK_SIZE) as usize], 0u64)
        };

        let half_chunk = (CHUNK_SIZE_TILE as f32) / 2.0 - TILE_SIZE / 2.0;
        let world_x = (key.chunk_pos.x as f32) * (CHUNK_SIZE_TILE as f32) + half_chunk;
        let world_y = (key.chunk_pos.y as f32) * (CHUNK_SIZE_TILE as f32) + half_chunk;
        let z_depth = -(Z_BELOW - key.z_index as f32) * TILE_SIZE;

        let visible = z_diff >= 0 && z_diff <= 8;

        let entity = commands
            .spawn((
                key,
                TilemapChunk {
                    chunk_size: UVec2::splat(CHUNK_SIZE as u32),
                    tile_display_size: UVec2::splat(TILE_PIXELS),
                    tileset: tileset.handle.clone(),
                    alpha_mode: AlphaMode2d::Blend,
                },
                TilemapChunkTileData(tile_data),
                Transform::from_xyz(world_x, world_y, z_depth),
                if visible {
                    Visibility::Visible
                } else {
                    Visibility::Hidden
                },
            ))
            .id();

        registry.entities.insert(key, entity);
        bench.tilemap_chunk_count += 1;
        spawned_this_call += 1;
    }

    println!(
        "[spawn_tilemap_chunks] spawned={} (skipped={}) total={}",
        spawned_this_call,
        pending_count - spawned_this_call,
        registry.entities.len()
    );
    bench.populate_ms += now.elapsed().as_secs_f64() * 1000.0;
}

fn compute_solid_bits(tiles: &[Option<TileData>]) -> u64 {
    let mut bits: u64 = 0;
    for (i, slot) in tiles.iter().enumerate() {
        if let Some(td) = slot {
            if td.tileset_index != 0 {
                bits |= 1u64 << i;
            }
        }
    }
    bits
}

pub fn update_tilemap_chunk_visibility(
    z_index: Res<ZIndex>,
    mut query: Query<(&ChunkLayerKey, &mut Visibility)>,
    mut bench: ResMut<super::TilemapBenchmark>,
) {
    let camera_z = z_index.0 as i32;
    let mut visible_count = 0u32;

    for (key, mut visibility) in query.iter_mut() {
        let z_diff = camera_z - camera_z_for_z_index(key.z_index);
        let should_show = z_diff >= 0 && z_diff <= 8;
        *visibility = if should_show {
            visible_count += 1;
            Visibility::Visible
        } else {
            Visibility::Hidden
        };
    }
    bench.tilemap_visible_count = visible_count;
}

pub fn populate_tilemap_chunk_data(
    mut registry: ResMut<TilemapChunkRegistry>,
    mut states: ResMut<TilemapChunkStates>,
    tilemap: Res<TileMap>,
    z_index: Res<ZIndex>,
    mut chunk_data_query: Query<(&mut TilemapChunkTileData, &ChunkLayerKey)>,
    mut bench: ResMut<super::TilemapBenchmark>,
) {
    let now = Instant::now();
    let camera_z = z_index.0 as i32;
    let mut processed = 0usize;

    if !registry.z_change_keys.is_empty() {
        let current_z_index = (camera_z as f32 + Z_BELOW) as usize;
        let mut z_keys: Vec<ChunkLayerKey> = registry.z_change_keys.drain().collect();
        z_keys.sort_by_key(|k| (k.z_index as i32 - current_z_index as i32).unsigned_abs());
        z_keys.truncate(MAX_POPULATE_PER_FRAME);

        for key in &z_keys {
            registry.z_change_keys.remove(key);
        }

        bench.last_z_change_dirty_count = registry.entities.len();

        for key in z_keys {
            let Some(&entity) = registry.entities.get(&key) else {
                continue;
            };
            let Ok((mut tile_data, _)) = chunk_data_query.get_mut(entity) else {
                continue;
            };
            if key.layer != LAYER_FLOOR {
                continue;
            }

            let z_diff = i32::max(camera_z - camera_z_for_z_index(key.z_index), 0);
            if z_diff > 8 {
                continue;
            }

            let solid_bits = states.solid_bits.get(&key).copied().unwrap_or(0);
            let underground_bits = states.underground_bits.get(&key).copied().unwrap_or(0);
            recolor_chunk_for_depth(&mut tile_data, z_diff, solid_bits, underground_bits);
            processed += 1;
        }

        let elapsed = now.elapsed().as_secs_f64() * 1000.0;
        bench.last_z_change_populate_ms = elapsed;
        bench.z_change_history_ms.push(elapsed);
        if bench.z_change_history_ms.len() > 100 {
            bench.z_change_history_ms.remove(0);
        }
    }

    if !registry.dirty_keys.is_empty() {
        let current_z_index = (camera_z as f32 + Z_BELOW) as usize;
        let mut to_process: Vec<ChunkLayerKey> = registry.dirty_keys.iter().cloned().collect();
        to_process.sort_by_key(|k| (k.z_index as i32 - current_z_index as i32).unsigned_abs());
        to_process.truncate(MAX_POPULATE_PER_FRAME);

        for key in &to_process {
            registry.dirty_keys.remove(key);
        }

        for key in to_process {
            let Some(&entity) = registry.entities.get(&key) else {
                continue;
            };
            let Ok((mut tile_data, _)) = chunk_data_query.get_mut(entity) else {
                continue;
            };
            if key.layer != LAYER_FLOOR {
                continue;
            }

            let z_diff = i32::max(camera_z - camera_z_for_z_index(key.z_index), 0);
            if z_diff > 8 {
                continue;
            }

            let (new_tiles, underground_bits) =
                populate_chunk_tiles(&tilemap, key.chunk_pos, key.z_index, camera_z);
            let solid_bits = compute_solid_bits(&new_tiles);
            states.solid_bits.insert(key, solid_bits);
            states.underground_bits.insert(key, underground_bits);
            tile_data.0 = new_tiles;
            processed += 1;
        }
    }

    let elapsed = now.elapsed().as_secs_f64() * 1000.0;
    bench.last_populate_ms = elapsed;
    bench.populate_ms += elapsed;
    bench.dirty_keys_last = processed;
    bench.dirty_keys_total += processed;
    bench.tile_data_mb = (registry.entities.len() as f64 * (CHUNK_SIZE * CHUNK_SIZE) as f64 * 4.0)
        / (1024.0 * 1024.0);
}

pub fn on_camera_z_changed(
    z_index: Res<ZIndex>,
    mut registry: ResMut<TilemapChunkRegistry>,
    mut prev_z: ResMut<PreviousZIndex>,
    mut bench: ResMut<super::TilemapBenchmark>,
) {
    if !z_index.is_changed() {
        return;
    }

    let old_camera_z = prev_z.0;
    let new_camera_z = z_index.0 as i32;
    prev_z.0 = new_camera_z;

    let now = Instant::now();

    let old_entered_idx = (old_camera_z as f32 + Z_BELOW) as usize;
    let old_exited_idx = ((old_camera_z - 8) as f32 + Z_BELOW) as usize;
    let new_entered_idx = (new_camera_z as f32 + Z_BELOW) as usize;
    let new_exited_idx = ((new_camera_z - 8) as f32 + Z_BELOW) as usize;

    let mut count = 0;

    let mut new_keys: Vec<ChunkLayerKey> = Vec::new();

    for key in registry.entities.keys() {
        if key.z_index == old_entered_idx
            || key.z_index == new_entered_idx
            || key.z_index == old_exited_idx
            || key.z_index == new_exited_idx
        {
            new_keys.push(*key);
            count += 1;
        }
    }

    if !new_keys.is_empty() {
        registry.z_change_keys.clear();
        for key in new_keys {
            registry.z_change_keys.insert(key);
        }
    }

    bench.last_z_change_dirty_ms = now.elapsed().as_secs_f64() * 1000.0;
    bench.last_z_change_dirty_count = count;
}

pub fn despawn_tilemap_chunks(
    mut commands: Commands,
    chunk_map: Res<super::super::chunks::ChunkMap>,
    mut registry: ResMut<TilemapChunkRegistry>,
    mut states: ResMut<TilemapChunkStates>,
    query: Query<(Entity, &ChunkLayerKey)>,
) {
    let loaded: HashSet<IVec2> = chunk_map.loaded_chunks.keys().cloned().collect();

    for (entity, key) in query.iter() {
        if !loaded.contains(&key.chunk_pos) {
            commands.entity(entity).despawn();
            registry.entities.remove(key);
            registry.dirty_keys.remove(key);
            registry.z_change_keys.remove(key);
            states.solid_bits.remove(key);
            states.underground_bits.remove(key);
        }
    }
}