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

use crate::{
    light::LightMap,
    world::{MAX_HEIGHT, Plant, TILE_SIZE, Terrain, TileChanged, WorldMap, hash2},
    worldgen::WorldGenerated,
    zones::Zones,
};

pub const CHUNK_SIZE: u32 = 64;
const TEXTURE_SIZE: u32 = 16;
const RELIEF: f32 = 0.35;
const LIGHT: Vec3 = Vec3::new(-1.0, 1.0, 1.5);

type Texture = Vec<[u8; 4]>;

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(world_or_light_changed),
                refresh_changed_chunks,
                refresh_zone_chunks.run_if(resource_changed::<Zones>),
            ),
        );
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum Layer {
    Terrain,
    Zones,
    Roofs,
    Plants,
}

impl Layer {
    fn z(self) -> f32 {
        match self {
            Layer::Terrain => 0.0,
            Layer::Zones => 0.3,
            Layer::Roofs => 0.85,
            Layer::Plants => 0.5,
        }
    }

    fn tiles(
        self,
        map: &WorldMap,
        zones: &Zones,
        light: &LightMap,
        origin: UVec2,
    ) -> Vec<Option<TileData>> {
        match self {
            Layer::Terrain => chunk_tiles(map, light, origin),
            Layer::Zones => zone_tiles(zones, origin),
            Layer::Roofs => roof_tiles(map, origin),
            Layer::Plants => plant_tiles(map, light, origin),
        }
    }
}

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

fn spawn_chunks(
    mut commands: Commands,
    map: Res<WorldMap>,
    zones: Res<Zones>,
    light: Res<LightMap>,
    mut images: ResMut<Assets<Image>>,
) {
    let terrain = images.add(stack_tileset(
        Terrain::ALL.iter().map(|&t| terrain_texture(t)).collect(),
    ));
    let plants = images.add(stack_tileset(
        Plant::ALL.iter().map(|&p| plant_texture(p)).collect(),
    ));
    let zone = images.add(stack_tileset(vec![zone_texture()]));
    let roof = images.add(stack_tileset(vec![roof_texture()]));
    let chunks = (map.size + CHUNK_SIZE - 1) / CHUNK_SIZE;
    for (layer, tileset) in [
        (Layer::Terrain, terrain),
        (Layer::Zones, zone),
        (Layer::Plants, plants),
        (Layer::Roofs, roof),
    ] {
        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!("{layer:?} chunk {cx},{cy}")),
                    Chunk { origin, layer },
                    TilemapChunk {
                        chunk_size: UVec2::splat(CHUNK_SIZE),
                        tile_display_size: UVec2::splat(TILE_SIZE as u32),
                        tileset: tileset.clone(),
                        alpha_mode: AlphaMode2d::Blend,
                    },
                    TilemapChunkTileData(layer.tiles(&map, &zones, &light, origin)),
                    Transform::from_translation(center.extend(layer.z())),
                ));
            }
        }
    }
}

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

fn world_or_light_changed(
    mut generated: MessageReader<WorldGenerated>,
    light: Res<LightMap>,
) -> bool {
    generated.read().count() > 0 || light.is_changed()
}

fn refresh_zone_chunks(
    map: Res<WorldMap>,
    zones: Res<Zones>,
    light: Res<LightMap>,
    mut chunks: Query<(&Chunk, &mut TilemapChunkTileData)>,
) {
    for (chunk, mut data) in &mut chunks {
        if chunk.layer == Layer::Zones {
            data.0 = chunk.layer.tiles(&map, &zones, &light, chunk.origin);
        }
    }
}

fn refresh_changed_chunks(
    map: Res<WorldMap>,
    zones: Res<Zones>,
    light: Res<LightMap>,
    mut changes: MessageReader<TileChanged>,
    mut chunks: Query<(&Chunk, &mut TilemapChunkTileData)>,
) {
    let origins: Vec<UVec2> = changes
        .read()
        .map(|TileChanged(tile)| chunk_origin(*tile))
        .collect();
    for (chunk, mut data) in &mut chunks {
        if origins.contains(&chunk.origin) {
            data.0 = chunk.layer.tiles(&map, &zones, &light, chunk.origin);
        }
    }
}

pub fn chunk_origin(tile: UVec2) -> UVec2 {
    tile / CHUNK_SIZE * CHUNK_SIZE
}

fn chunk_positions(origin: UVec2) -> impl Iterator<Item = UVec2> {
    (0..CHUNK_SIZE).flat_map(move |y| (0..CHUNK_SIZE).map(move |x| origin + UVec2::new(x, y)))
}

pub fn chunk_tiles(map: &WorldMap, light: &LightMap, origin: UVec2) -> Vec<Option<TileData>> {
    chunk_positions(origin)
        .map(|pos| {
            map.get(pos.as_ivec2()).map(|terrain| TileData {
                tileset_index: terrain.index(),
                color: tint(jitter(pos) * shade(map, pos), light.color(pos)),
                ..default()
            })
        })
        .collect()
}

pub fn zone_tiles(zones: &Zones, origin: UVec2) -> Vec<Option<TileData>> {
    chunk_positions(origin)
        .map(|pos| {
            zones
                .is_stockpile(pos)
                .then(|| TileData::from_tileset_index(0))
        })
        .collect()
}

pub fn roof_tiles(map: &WorldMap, origin: UVec2) -> Vec<Option<TileData>> {
    chunk_positions(origin)
        .map(|pos| {
            map.is_roofed(pos.as_ivec2())
                .then(|| TileData::from_tileset_index(0))
        })
        .collect()
}

pub fn plant_tiles(map: &WorldMap, light: &LightMap, origin: UVec2) -> Vec<Option<TileData>> {
    chunk_positions(origin)
        .map(|pos| {
            map.plant(pos.as_ivec2()).map(|plant| TileData {
                tileset_index: plant.index(),
                color: tint(jitter(pos) * hillshade(map, pos), light.color(pos)),
                ..default()
            })
        })
        .collect()
}

fn tint(v: f32, light: Vec3) -> Color {
    Color::srgb(v * light.x, v * light.y, v * light.z)
}

fn jitter(pos: UVec2) -> f32 {
    0.92 + (hash2(pos.x, pos.y) % 9) as f32 * 0.01
}

fn shade(map: &WorldMap, pos: UVec2) -> f32 {
    match map.get(pos.as_ivec2()) {
        Some(t) if t.is_water() => 1.0,
        _ => hillshade(map, pos) * height_brightness(map, pos),
    }
}

fn height_brightness(map: &WorldMap, pos: UVec2) -> f32 {
    0.9 + 0.2 * map.height(pos.as_ivec2()).unwrap_or(0.0) / MAX_HEIGHT
}

pub fn hillshade(map: &WorldMap, pos: UVec2) -> f32 {
    let p = pos.as_ivec2();
    let center = map.height(p).unwrap_or(0.0);
    let h = |d: IVec2| map.height(p + d).unwrap_or(center);
    let dx = (h(IVec2::X) - h(IVec2::NEG_X)) / 2.0;
    let dy = (h(IVec2::Y) - h(IVec2::NEG_Y)) / 2.0;
    let normal = Vec3::new(-dx * RELIEF, -dy * RELIEF, 1.0).normalize();
    let light = LIGHT.normalize();
    (normal.dot(light) / light.z).clamp(0.6, 1.25)
}

fn terrain_texture(terrain: Terrain) -> Texture {
    let base = terrain.color();
    let layer = terrain.index() as u32;
    (0..TEXTURE_SIZE * TEXTURE_SIZE)
        .map(|i| {
            let speck = terrain
                .speck_color()
                .filter(|_| hash2(i / 2 + i / 32, layer) % 5 == 0);
            let jitter = (hash2(i, layer) % 16) as f32 / 255.0 - 8.0 / 255.0;
            let c = speck.unwrap_or(base);
            Srgba::new(c.red + jitter, c.green + jitter, c.blue + jitter, 1.0).to_u8_array()
        })
        .collect()
}

fn zone_texture() -> Texture {
    (0..TEXTURE_SIZE * TEXTURE_SIZE)
        .map(|i| {
            let (x, y) = (i % TEXTURE_SIZE, i / TEXTURE_SIZE);
            let edge = x == 0 || y == 0 || x == TEXTURE_SIZE - 1 || y == TEXTURE_SIZE - 1;
            [90, 160, 230, if edge { 150 } else { 55 }]
        })
        .collect()
}

fn roof_texture() -> Texture {
    (0..TEXTURE_SIZE * TEXTURE_SIZE)
        .map(|i| {
            let (x, y) = (i % TEXTURE_SIZE, i / TEXTURE_SIZE);
            let stripe = (x + y) % 4 == 0;
            [110, 80, 50, if stripe { 90 } else { 45 }]
        })
        .collect()
}

fn plant_texture(plant: Plant) -> Texture {
    match plant {
        Plant::Tree => tree_texture(),
    }
}

fn tree_texture() -> Texture {
    let leaves = Srgba::rgb_u8(38, 92, 42);
    let highlight = Srgba::rgb_u8(66, 126, 58);
    let trunk = Srgba::rgb_u8(92, 64, 40);
    let canopy_center = Vec2::new(7.5, 6.5);
    (0..TEXTURE_SIZE * TEXTURE_SIZE)
        .map(|i| {
            let (x, y) = (i % TEXTURE_SIZE, i / TEXTURE_SIZE);
            let p = Vec2::new(x as f32, y as f32);
            let from_center = p - canopy_center;
            if from_center.length() <= 6.0 {
                let lit = from_center.dot(Vec2::new(-1.0, -1.0)) > 3.0;
                if lit { highlight } else { leaves }.to_u8_array()
            } else if (7..=8).contains(&x) && y >= 12 && y <= 14 {
                trunk.to_u8_array()
            } else {
                [0, 0, 0, 0]
            }
        })
        .collect()
}

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

fn stack_tileset(textures: Vec<Texture>) -> Image {
    let layers = tileset_layer_count(textures.len() as u32);
    let empty = vec![[0, 0, 0, 0]; (TEXTURE_SIZE * TEXTURE_SIZE) as usize];
    let data = (0..layers as usize)
        .flat_map(|i| textures.get(i).unwrap_or(&empty).clone())
        .flatten()
        .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, &LightMap::default(), 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))
            .init_resource::<Zones>()
            .init_resource::<LightMap>()
            .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,
                    layer: Layer::Terrain,
                },
                TilemapChunkTileData(chunk_tiles(&map, &LightMap::default(), 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 darkness_tints_tiles_and_fires_warm_them() {
        let map = WorldMap::filled(UVec2::splat(CHUNK_SIZE), Terrain::Grass);
        let fire = (
            UVec2::new(10, 10),
            crate::structures::StructureKind::Campfire,
        );
        let night = LightMap::compute(map.size, 0.0, [fire]);
        let tiles = chunk_tiles(&map, &night, UVec2::ZERO);
        let lum = |pos: UVec2| {
            let c = tiles[(pos.y * CHUNK_SIZE + pos.x) as usize]
                .unwrap()
                .color
                .to_srgba();
            c.red + c.green + c.blue
        };
        assert!(lum(UVec2::new(10, 11)) > lum(UVec2::new(40, 40)) * 2.0);
    }

    #[test]
    fn roof_layer_covers_roofed_tiles_only() {
        let mut map = WorldMap::filled(UVec2::splat(4), Terrain::Grass);
        map.set_roof(UVec2::new(2, 1), true);
        let tiles = roof_tiles(&map, UVec2::ZERO);
        assert!(tiles[(CHUNK_SIZE + 2) as usize].is_some());
        assert_eq!(tiles.iter().filter(|t| t.is_some()).count(), 1);
    }

    #[test]
    fn chunk_origin_rounds_down_to_chunk_grid() {
        assert_eq!(chunk_origin(UVec2::new(0, 63)), UVec2::new(0, 0));
        assert_eq!(chunk_origin(UVec2::new(64, 130)), UVec2::new(64, 128));
    }

    #[test]
    fn tileset_has_a_layer_for_every_texture() {
        let textures = Terrain::ALL.iter().map(|&t| terrain_texture(t)).collect();
        let image = stack_tileset(textures);
        assert!(image.texture_descriptor.size.depth_or_array_layers >= Terrain::ALL.len() as u32);
    }

    #[test]
    fn plant_layer_only_has_tiles_where_plants_grow() {
        let mut map = WorldMap::filled(UVec2::splat(4), Terrain::Grass);
        map.set_plant(UVec2::new(1, 0), Some(Plant::Tree));
        let tiles = plant_tiles(&map, &LightMap::default(), UVec2::ZERO);
        assert!(tiles[0].is_none());
        assert_eq!(tiles[1].unwrap().tileset_index, Plant::Tree.index());
    }

    // Light comes from the north-west (LIGHT has -x, +y), the convention for hillshaded maps.
    #[test]
    fn slopes_facing_the_light_are_brighter_than_slopes_facing_away() {
        let mut map = WorldMap::filled(UVec2::new(5, 1), Terrain::Grass);
        for (x, h) in [0.0, 5.0, 10.0, 5.0, 0.0].into_iter().enumerate() {
            map.set_height(UVec2::new(x as u32, 0), h);
        }
        let flat = WorldMap::filled(UVec2::new(3, 1), Terrain::Grass);
        let facing_west = hillshade(&map, UVec2::new(1, 0));
        let facing_east = hillshade(&map, UVec2::new(3, 0));
        let level = hillshade(&flat, UVec2::new(1, 0));
        assert!(facing_west > level);
        assert!(facing_east < level);
    }

    // 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");
        }
    }
}