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>,
}

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

#[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,
        }
    }
}

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>> {
    let z_diff = i32::max(camera_z - camera_z_for_z_index(z_index), 0);
    if z_diff > 8 {
        return vec![None; (CHUNK_SIZE * CHUNK_SIZE) as usize];
    }

    let mut tiles = Vec::with_capacity((CHUNK_SIZE * CHUNK_SIZE) as usize);
    for local_y in 0..CHUNK_SIZE {
        for local_x in 0..CHUNK_SIZE {
            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 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
}

pub fn spawn_tilemap_chunks(
    mut commands: Commands,
    tilemap: Res<TileMap>,
    tileset: Res<TilemapTileset>,
    mut registry: ResMut<TilemapChunkRegistry>,
    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();
            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 = if key.layer == LAYER_FLOOR {
            populate_chunk_tiles(&tilemap, key.chunk_pos, key.z_index, camera_z)
        } else {
            vec![None; (CHUNK_SIZE * CHUNK_SIZE) as usize]
        };

        // Fixed
        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;
}

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>,
    tilemap: Res<TileMap>,
    z_index: Res<ZIndex>,
    mut chunk_data_query: Query<(&mut TilemapChunkTileData, &ChunkLayerKey)>,
    mut bench: ResMut<super::TilemapBenchmark>,
) {
    if registry.dirty_keys.is_empty() {
        return;
    }
    let now = Instant::now();
    let camera_z = z_index.0 as i32;
    let current_z = (camera_z as f32 + Z_BELOW) as usize;

    let dirty: Vec<ChunkLayerKey> = registry.dirty_keys.drain().collect();
    let mut processed = 0usize;

    for key in dirty {
        let Some(&entity) = registry.entities.get(&key) else {
            continue;
        };
        let Ok((mut tile_data, chunk_key)) = 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 = populate_chunk_tiles(&tilemap, key.chunk_pos, key.z_index, camera_z);
        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>) {
    if z_index.is_changed() {
        let keys: Vec<ChunkLayerKey> = registry.entities.keys().cloned().collect();
        for key in keys {
            registry.dirty_keys.insert(key);
        }
    }
}

pub fn despawn_tilemap_chunks(
    mut commands: Commands,
    chunk_map: Res<super::super::chunks::ChunkMap>,
    mut registry: ResMut<TilemapChunkRegistry>,
    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);
        }
    }
}