use bevy::{
    asset::RenderAssetUsages,
    prelude::*,
    render::render_resource::{Extent3d, TextureDimension, TextureFormat},
    sprite_render::{TileData, TilemapChunk, TilemapChunkTileData},
};

use crate::{
    world::{TILE_SIZE, Terrain, WorldMap, hash2},
    worldgen::WorldGenerated,
};

pub const CHUNK_SIZE: u32 = 64;
const TEXTURE_SIZE: u32 = 16;

pub struct TilemapPlugin;

impl Plugin for TilemapPlugin {
    fn build(&self, app: &mut App) {
        app.add_systems(Startup, spawn_chunks)
            .add_systems(Update, refresh_chunks.run_if(on_message::<WorldGenerated>));
    }
}

#[derive(Component)]
struct Chunk {
    origin: UVec2,
}

fn spawn_chunks(mut commands: Commands, map: Res<WorldMap>, mut images: ResMut<Assets<Image>>) {
    let tileset = images.add(build_tileset());
    let chunks = (map.size + CHUNK_SIZE - 1) / CHUNK_SIZE;
    for cy in 0..chunks.y {
        for cx in 0..chunks.x {
            let origin = UVec2::new(cx, cy) * CHUNK_SIZE;
            let center = (origin.as_vec2() + CHUNK_SIZE as f32 / 2.0) * TILE_SIZE;
            commands.spawn((
                Name::new(format!("Chunk {cx},{cy}")),
                Chunk { origin },
                TilemapChunk {
                    chunk_size: UVec2::splat(CHUNK_SIZE),
                    tile_display_size: UVec2::splat(TILE_SIZE as u32),
                    tileset: tileset.clone(),
                    ..default()
                },
                TilemapChunkTileData(chunk_tiles(&map, origin)),
                Transform::from_translation(center.extend(0.0)),
            ));
        }
    }
}

fn refresh_chunks(map: Res<WorldMap>, mut chunks: Query<(&Chunk, &mut TilemapChunkTileData)>) {
    for (chunk, mut data) in &mut chunks {
        data.0 = chunk_tiles(&map, chunk.origin);
    }
}

pub fn chunk_tiles(map: &WorldMap, origin: UVec2) -> Vec<Option<TileData>> {
    (0..CHUNK_SIZE)
        .flat_map(|y| (0..CHUNK_SIZE).map(move |x| origin + UVec2::new(x, y)))
        .map(|pos| {
            map.get(pos.as_ivec2()).map(|terrain| TileData {
                tileset_index: terrain.index(),
                color: tile_tint(pos),
                ..default()
            })
        })
        .collect()
}

fn tile_tint(pos: UVec2) -> Color {
    let v = 0.92 + (hash2(pos.x, pos.y) % 9) as f32 * 0.01;
    Color::srgb(v, v, v)
}

pub fn tileset_layer_count(terrains: u32) -> u32 {
    match terrains {
        0 | 1 => 2,
        n if n % 6 == 0 => n + 1,
        n => n,
    }
}

fn build_tileset() -> Image {
    let layers = tileset_layer_count(Terrain::ALL.len() as u32);
    let data = (0..layers)
        .flat_map(|layer| {
            let base = Terrain::ALL
                .get(layer as usize)
                .map_or(Srgba::NONE, |t| t.color());
            (0..TEXTURE_SIZE * TEXTURE_SIZE).flat_map(move |i| {
                let jitter = (hash2(i, layer) % 16) as f32 / 255.0 - 8.0 / 255.0;
                let c = Srgba::new(
                    base.red + jitter,
                    base.green + jitter,
                    base.blue + jitter,
                    1.0,
                );
                c.to_u8_array()
            })
        })
        .collect();

    let mut image = Image::new(
        Extent3d {
            width: TEXTURE_SIZE,
            height: TEXTURE_SIZE * layers,
            depth_or_array_layers: 1,
        },
        TextureDimension::D2,
        data,
        TextureFormat::Rgba8UnormSrgb,
        RenderAssetUsages::RENDER_WORLD,
    );
    image
        .reinterpret_stacked_2d_as_array(layers)
        .expect("tileset is a vertical stack of square layers");
    image
}

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

    #[test]
    fn chunk_tiles_beyond_map_edge_are_empty() {
        let map = WorldMap::filled(UVec2::new(CHUNK_SIZE + 2, 3), Terrain::Stone);
        let tiles = chunk_tiles(&map, UVec2::new(CHUNK_SIZE, 0));
        assert_eq!(tiles.len(), (CHUNK_SIZE * CHUNK_SIZE) as usize);
        assert!(tiles[0].is_some());
        assert!(tiles[1].is_some());
        assert!(tiles[2].is_none());
        assert!(tiles[CHUNK_SIZE as usize * 3].is_none());
    }

    // Chunk entities live for the whole session. World changes (regeneration now, digging and
    // building later) rewrite their tile data in place; TilemapChunk re-uploads it through change
    // detection, so no entities, meshes or materials are recreated.
    #[test]
    fn regeneration_rewrites_tile_data_of_existing_chunks() {
        let mut app = App::new();
        app.insert_resource(WorldMap::filled(UVec2::splat(CHUNK_SIZE), Terrain::Grass))
            .add_message::<WorldGenerated>()
            .add_systems(Update, refresh_chunks.run_if(on_message::<WorldGenerated>));
        let map = app.world().resource::<WorldMap>().clone();
        let chunk = app
            .world_mut()
            .spawn((
                Chunk {
                    origin: UVec2::ZERO,
                },
                TilemapChunkTileData(chunk_tiles(&map, UVec2::ZERO)),
            ))
            .id();

        *app.world_mut().resource_mut::<WorldMap>() =
            WorldMap::filled(UVec2::splat(CHUNK_SIZE), Terrain::Stone);
        app.world_mut().write_message(WorldGenerated);
        app.update();

        let data = app.world().get::<TilemapChunkTileData>(chunk).unwrap();
        assert!(
            data.0
                .iter()
                .all(|t| t.unwrap().tileset_index == Terrain::Stone.index())
        );
    }

    #[test]
    fn tileset_has_a_layer_for_every_terrain() {
        let image = build_tileset();
        assert!(image.texture_descriptor.size.depth_or_array_layers >= Terrain::ALL.len() as u32);
    }

    // wgpu's GLES backend (WebGL2) guesses the texture view dimension from the layer count:
    // 1 layer => D2, 6 => Cube, multiples of 6 => CubeArray. The tilemap shader needs D2Array,
    // so the tileset is padded with an unused layer to dodge those counts.
    #[test]
    fn tileset_layer_count_avoids_webgl_view_dimension_heuristics() {
        for n in 0..50 {
            let layers = tileset_layer_count(n);
            assert!(layers >= n.max(2));
            assert!(layers % 6 != 0, "{n} terrains -> {layers} layers");
        }
    }
}