use std::{
    cmp::Reverse,
    collections::{BinaryHeap, VecDeque},
};

use bevy::prelude::*;

use crate::{
    actions::Action,
    noise::Noise,
    rng::SimRng,
    world::{MAX_HEIGHT, Plant, Terrain, WorldMap, hash2},
};

pub const WORLD_SIZE: UVec2 = UVec2::new(256, 256);
pub const DEFAULT_SEED: u32 = 1;

const ELEVATION_SCALE: f32 = 1.0 / 48.0;
const FERTILITY_SCALE: f32 = 1.0 / 16.0;
const DEEP_WATER_SHARE: f32 = 0.10;
const SHALLOW_WATER_SHARE: f32 = 0.06;
const SAND_SHARE: f32 = 0.04;
const STONE_SHARE: f32 = 0.14;
const SOIL_SHARE: f32 = 0.20;
const ORE_SCALE: f32 = 1.0 / 6.0;
const IRON_SHARE_OF_ROCK: f32 = 0.10;
const GOLD_SHARE_OF_ROCK: f32 = 0.03;
const FOREST_SCALE: f32 = 1.0 / 32.0;
const FOREST_SHARE: f32 = 0.35;
const FOREST_TREE_PERCENT: u32 = 45;
const LONE_TREE_PERCENT: u32 = 3;
const BERRY_BUSH_PERCENT: u32 = 2;
const HERB_PERCENT: u32 = 1;
const RIVER_COUNT: usize = 4;
const RIVER_CANDIDATES: usize = 40;
const RIVER_SOURCE_SPACING: f32 = 40.0;

pub struct WorldGenPlugin;

impl Plugin for WorldGenPlugin {
    fn build(&self, app: &mut App) {
        app.insert_resource(WorldSeed(DEFAULT_SEED))
            .insert_resource(generate(WORLD_SIZE, DEFAULT_SEED))
            .add_message::<WorldGenerated>()
            .add_systems(Startup, |mut generated: MessageWriter<WorldGenerated>| {
                generated.write(WorldGenerated);
            })
            .add_systems(Update, regenerate);
    }
}

#[derive(Resource, Debug, Clone, Copy, PartialEq)]
pub struct WorldSeed(pub u32);

#[derive(Message, Debug, Clone, Copy)]
pub struct WorldGenerated;

fn regenerate(
    mut actions: MessageReader<Action>,
    mut seed: ResMut<WorldSeed>,
    mut map: ResMut<WorldMap>,
    mut generated: MessageWriter<WorldGenerated>,
) {
    if actions.read().any(|a| *a == Action::NewWorld) {
        seed.0 = seed.0.wrapping_add(1);
        *map = generate(map.size, seed.0);
        generated.write(WorldGenerated);
    }
}

pub fn generate(size: UVec2, seed: u32) -> WorldMap {
    let mut map = WorldMap::filled(size, Terrain::Grass);
    let positions: Vec<UVec2> = map.positions().collect();
    let sample = |offset: u32, scale: f32, octaves: u32| -> Vec<f32> {
        let noise = Noise::new(seed.wrapping_add(offset));
        positions
            .iter()
            .map(|p| noise.fbm(p.as_vec2() * scale, octaves))
            .collect()
    };
    let elevation = sample(0, ELEVATION_SCALE, 5);
    let fertility = sample(0x1000, FERTILITY_SCALE, 3);
    let iron = sample(0x2000, ORE_SCALE, 2);
    let gold = sample(0x3000, ORE_SCALE, 2);
    let forest = sample(0x4000, FOREST_SCALE, 3);

    let (low, high) = (quantile(&elevation, 0.0), quantile(&elevation, 1.0));
    for (&pos, &e) in positions.iter().zip(&elevation) {
        map.set_height(pos, (e - low) / (high - low).max(f32::EPSILON) * MAX_HEIGHT);
    }

    let deep = quantile(&elevation, DEEP_WATER_SHARE);
    let shallow = quantile(&elevation, DEEP_WATER_SHARE + SHALLOW_WATER_SHARE);
    let sand = quantile(
        &elevation,
        DEEP_WATER_SHARE + SHALLOW_WATER_SHARE + SAND_SHARE,
    );
    let stone = quantile(&elevation, 1.0 - STONE_SHARE);
    let soil = quantile(&fertility, 1.0 - SOIL_SHARE);
    let is_rock = |i: usize| elevation[i] >= stone;
    let rock_values = |values: &[f32]| -> Vec<f32> {
        (0..values.len())
            .filter(|&i| is_rock(i))
            .map(|i| values[i])
            .collect()
    };
    let iron_min = quantile(&rock_values(&iron), 1.0 - IRON_SHARE_OF_ROCK);
    let gold_min = quantile(&rock_values(&gold), 1.0 - GOLD_SHARE_OF_ROCK);

    for (i, &pos) in positions.iter().enumerate() {
        let e = elevation[i];
        let terrain = if e < deep {
            Terrain::DeepWater
        } else if e < shallow {
            Terrain::ShallowWater
        } else if e < sand {
            Terrain::Sand
        } else if is_rock(i) && gold[i] >= gold_min {
            Terrain::GoldOre
        } else if is_rock(i) && iron[i] >= iron_min {
            Terrain::IronOre
        } else if is_rock(i) {
            Terrain::Stone
        } else if fertility[i] >= soil {
            Terrain::Soil
        } else {
            Terrain::Grass
        };
        map.set(pos, terrain);
    }

    for river in pick_rivers(&map, seed) {
        carve_river(&mut map, &river);
    }

    let forest_min = quantile(&forest, 1.0 - FOREST_SHARE);
    for (i, &pos) in positions.iter().enumerate() {
        let fertile = matches!(
            map.get(pos.as_ivec2()),
            Some(Terrain::Grass | Terrain::Soil)
        );
        map.set_forest(pos, forest[i] >= forest_min);
        let percent = if forest[i] >= forest_min {
            FOREST_TREE_PERCENT
        } else {
            LONE_TREE_PERCENT
        };
        let roll = hash2(pos.x ^ seed, pos.y) % 100;
        if fertile && roll < percent {
            map.set_plant(pos, Some(Plant::Tree));
        } else if fertile && roll >= 100 - BERRY_BUSH_PERCENT {
            map.set_plant(pos, Some(Plant::BerryBush));
            map.set_growth(pos, (hash2(pos.y ^ seed, pos.x) % 140) as f32 / 100.0);
        } else if fertile && roll >= 100 - BERRY_BUSH_PERCENT - HERB_PERCENT {
            map.set_plant(pos, Some(Plant::Herb));
            map.set_growth(pos, (hash2(pos.y ^ seed, pos.x) % 140) as f32 / 100.0);
        }
    }
    map
}

fn mountain_feet(map: &WorldMap) -> Vec<UVec2> {
    map.positions()
        .filter(|p| {
            let p = p.as_ivec2();
            map.move_cost(p).is_some()
                && !map.get(p).unwrap().is_water()
                && [IVec2::X, IVec2::NEG_X, IVec2::Y, IVec2::NEG_Y]
                    .iter()
                    .any(|&d| map.get(p + d).is_some_and(Terrain::is_rock))
        })
        .collect()
}

pub fn pick_rivers(map: &WorldMap, seed: u32) -> Vec<Vec<UVec2>> {
    let feet = mountain_feet(map);
    if feet.is_empty() {
        return Vec::new();
    }
    let flow = FlowField::new(map);
    let mut rng = SimRng::new(u64::from(seed) ^ 0x5249_5645);
    let mut rivers: Vec<Vec<UVec2>> = (0..RIVER_CANDIDATES)
        .map(|_| feet[rng.range_i32(0..feet.len() as i32) as usize])
        .map(|source| flow.trace(source))
        .collect();
    rivers.sort_by_key(|river| Reverse(river.len()));
    let mut picked: Vec<Vec<UVec2>> = Vec::new();
    for river in rivers {
        if picked.len() == RIVER_COUNT {
            break;
        }
        let spaced = picked
            .iter()
            .all(|other| other[0].as_vec2().distance(river[0].as_vec2()) >= RIVER_SOURCE_SPACING);
        if spaced {
            picked.push(river);
        }
    }
    picked
}

pub struct FlowField {
    size: UVec2,
    downstream: Vec<Option<UVec2>>,
}

impl FlowField {
    pub fn new(map: &WorldMap) -> Self {
        let n = map.size.element_product() as usize;
        let index = |p: UVec2| (p.y * map.size.x + p.x) as usize;
        let mut downstream = vec![None; n];
        let mut done = vec![false; n];
        let mut open = BinaryHeap::new();
        for pos in map.positions() {
            let p = pos.as_ivec2();
            let terrain = map.get(p).unwrap();
            let at_edge =
                pos.x == 0 || pos.y == 0 || pos.x == map.size.x - 1 || pos.y == map.size.y - 1;
            if terrain.is_water() || (at_edge && !terrain.is_rock()) {
                done[index(pos)] = true;
                open.push(Reverse((map.height(p).unwrap().to_bits(), pos.x, pos.y)));
            }
        }
        while let Some(Reverse((level, x, y))) = open.pop() {
            let current = UVec2::new(x, y);
            for step in [IVec2::X, IVec2::NEG_X, IVec2::Y, IVec2::NEG_Y] {
                let next = current.as_ivec2() + step;
                if !map.get(next).is_some_and(|t| !t.is_rock()) || done[index(next.as_uvec2())] {
                    continue;
                }
                let next = next.as_uvec2();
                done[index(next)] = true;
                downstream[index(next)] = Some(current);
                let filled = map
                    .height(next.as_ivec2())
                    .unwrap()
                    .max(f32::from_bits(level));
                open.push(Reverse((filled.to_bits(), next.x, next.y)));
            }
        }
        Self {
            size: map.size,
            downstream,
        }
    }

    fn downstream(&self, pos: UVec2) -> Option<UVec2> {
        self.downstream[(pos.y * self.size.x + pos.x) as usize]
    }

    pub fn trace(&self, source: UVec2) -> Vec<UVec2> {
        let mut path = vec![source];
        while let Some(next) = self.downstream(*path.last().unwrap()) {
            path.push(next);
        }
        path
    }
}

fn carve_river(map: &mut WorldMap, path: &[UVec2]) {
    for (i, &pos) in path.iter().enumerate() {
        let mut tiles = vec![pos.as_ivec2()];
        if i >= path.len() / 3 {
            tiles.push(pos.as_ivec2() + IVec2::X);
        }
        if i >= path.len() * 2 / 3 {
            tiles.push(pos.as_ivec2() + IVec2::Y);
        }
        for tile in tiles {
            if map
                .get(tile)
                .is_some_and(|t| !t.is_rock() && t != Terrain::DeepWater)
            {
                map.set(tile.as_uvec2(), Terrain::ShallowWater);
            }
        }
    }
}

pub fn quantile(values: &[f32], share: f32) -> f32 {
    let mut sorted = values.to_vec();
    sorted.sort_by(f32::total_cmp);
    let i = ((sorted.len() as f32 * share) as usize).min(sorted.len() - 1);
    sorted[i]
}

pub fn largest_region(map: &WorldMap) -> Vec<UVec2> {
    let mut seen = vec![false; map.size.element_product() as usize];
    let index = |p: UVec2| (p.y * map.size.x + p.x) as usize;
    let mut best = Vec::new();
    for start in map.positions() {
        if seen[index(start)] || map.move_cost(start.as_ivec2()).is_none() {
            continue;
        }
        let region = flood(map, start, |p| {
            let fresh = !seen[index(p)];
            seen[index(p)] = true;
            fresh
        });
        if region.len() > best.len() {
            best = region;
        }
    }
    best
}

fn flood(map: &WorldMap, start: UVec2, mut visit: impl FnMut(UVec2) -> bool) -> Vec<UVec2> {
    let mut region = Vec::new();
    let mut queue = VecDeque::from([start]);
    visit(start);
    while let Some(pos) = queue.pop_front() {
        region.push(pos);
        for step in [IVec2::X, IVec2::NEG_X, IVec2::Y, IVec2::NEG_Y] {
            let next = pos.as_ivec2() + step;
            if map.move_cost(next).is_some() && visit(next.as_uvec2()) {
                queue.push_back(next.as_uvec2());
            }
        }
    }
    region
}

pub fn spawn_tiles(map: &WorldMap, count: usize) -> Vec<UVec2> {
    let center = map.size.as_vec2() / 2.0;
    let Some(&origin) = largest_region(map).iter().min_by(|a, b| {
        a.as_vec2()
            .distance_squared(center)
            .total_cmp(&b.as_vec2().distance_squared(center))
    }) else {
        return Vec::new();
    };
    let mut seen = std::collections::HashSet::new();
    flood(map, origin, |p| seen.insert(p))
        .into_iter()
        .step_by(2)
        .take(count)
        .collect()
}

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

    const SIZE: UVec2 = UVec2::new(96, 64);

    fn share(map: &WorldMap, terrain: Terrain) -> f32 {
        let n = map
            .positions()
            .filter(|p| map.get(p.as_ivec2()) == Some(terrain))
            .count();
        n as f32 / map.size.element_product() as f32
    }

    #[test]
    fn same_seed_generates_same_world() {
        let (a, b, c) = (generate(SIZE, 7), generate(SIZE, 7), generate(SIZE, 8));
        let tiles = |m: &WorldMap| {
            m.positions()
                .map(|p| {
                    let p = p.as_ivec2();
                    (m.get(p), m.plant(p), m.height(p).map(f32::to_bits))
                })
                .collect::<Vec<_>>()
        };
        assert_eq!(tiles(&a), tiles(&b));
        assert_ne!(tiles(&a), tiles(&c));
    }

    // Elevation thresholds come from quantiles rather than fixed values, so every seed gets the
    // same terrain mix instead of occasionally being all ocean or all mountain.
    #[test]
    fn terrain_shares_are_stable_across_seeds() {
        for seed in 0..5 {
            let map = generate(SIZE, seed);
            assert!((share(&map, Terrain::DeepWater) - DEEP_WATER_SHARE).abs() < 0.02);
            let rock: f32 = [Terrain::Stone, Terrain::IronOre, Terrain::GoldOre]
                .iter()
                .map(|&t| share(&map, t))
                .sum();
            assert!((rock - STONE_SHARE).abs() < 0.02);
            let generated = |t: &Terrain| {
                !matches!(
                    t,
                    Terrain::Gravel
                        | Terrain::Path
                        | Terrain::Paved
                        | Terrain::WoodFloor
                        | Terrain::Bridge
                )
            };
            for terrain in Terrain::ALL.into_iter().filter(generated) {
                assert!(share(&map, terrain) > 0.0, "seed {seed} lacks {terrain:?}");
            }
        }
    }

    #[test]
    fn gold_is_rarer_than_iron() {
        let map = generate(SIZE, 3);
        assert!(share(&map, Terrain::GoldOre) < share(&map, Terrain::IronOre));
    }

    #[test]
    fn mountains_are_the_highest_ground() {
        let map = generate(SIZE, 4);
        let mean_height = |pred: fn(Terrain) -> bool| {
            let hs: Vec<f32> = map
                .positions()
                .map(|p| p.as_ivec2())
                .filter(|&p| pred(map.get(p).unwrap()))
                .map(|p| map.height(p).unwrap())
                .collect();
            hs.iter().sum::<f32>() / hs.len() as f32
        };
        assert!(mean_height(Terrain::is_rock) > mean_height(|t| t == Terrain::Grass));
        assert!(mean_height(|t| t == Terrain::Grass) > mean_height(|t| t == Terrain::DeepWater));
        assert!(
            map.positions()
                .all(|p| (0.0..=MAX_HEIGHT).contains(&map.height(p.as_ivec2()).unwrap()))
        );
    }

    #[test]
    fn worlds_have_berry_bushes_some_of_them_ripe() {
        let map = generate(SIZE, 2);
        let bushes: Vec<_> = map
            .positions()
            .map(|p| p.as_ivec2())
            .filter(|&p| map.plant(p) == Some(Plant::BerryBush))
            .collect();
        assert!(bushes.len() > 20);
        assert!(bushes.iter().any(|&p| map.ripe_plant(p).is_some()));
        assert!(bushes.iter().any(|&p| map.ripe_plant(p).is_none()));
    }

    #[test]
    fn trees_only_grow_on_grass_and_soil() {
        let map = generate(SIZE, 2);
        let trees: Vec<_> = map
            .positions()
            .map(|p| p.as_ivec2())
            .filter(|&p| map.plant(p) == Some(Plant::Tree))
            .collect();
        assert!(!trees.is_empty());
        assert!(
            trees
                .iter()
                .all(|&p| matches!(map.get(p), Some(Terrain::Grass | Terrain::Soil)))
        );
    }

    fn slope_map() -> WorldMap {
        let mut map = WorldMap::filled(UVec2::new(10, 5), Terrain::Grass);
        for pos in map.positions() {
            map.set_height(pos, pos.x as f32);
            if pos.x == 0 {
                map.set(pos, Terrain::DeepWater);
            }
        }
        map
    }

    #[test]
    fn rivers_flow_downhill_into_water() {
        let map = slope_map();
        let path = FlowField::new(&map).trace(UVec2::new(8, 2));
        assert_eq!(
            map.get(path.last().unwrap().as_ivec2()),
            Some(Terrain::DeepWater)
        );
        assert!(path.windows(2).all(|w| w[1].x < w[0].x));
    }

    // Rivers are carved into the terrain after rock is placed; flowing through mountains would
    // leave walkable water channels cutting through what should be solid stone.
    #[test]
    fn rivers_route_around_rock() {
        let mut map = slope_map();
        map.set(UVec2::new(5, 2), Terrain::Stone);
        let path = FlowField::new(&map).trace(UVec2::new(8, 2));
        assert!(
            path.iter()
                .all(|p| !map.get(p.as_ivec2()).unwrap().is_rock())
        );
        assert_eq!(
            map.get(path.last().unwrap().as_ivec2()),
            Some(Terrain::DeepWater)
        );
    }

    // Noise terrain is full of small dips. Flow directions come from a priority flood that
    // spills out of dips at their lowest rim, so rivers cross them instead of ending in (or
    // meandering around inside) every depression.
    #[test]
    fn rivers_cross_depressions_to_reach_water() {
        let mut map = slope_map();
        map.set_height(UVec2::new(5, 2), 0.5);
        let path = FlowField::new(&map).trace(UVec2::new(8, 2));
        assert_eq!(
            map.get(path.last().unwrap().as_ivec2()),
            Some(Terrain::DeepWater)
        );
        assert!(path.len() <= 9);
    }

    #[test]
    fn generated_worlds_have_rivers() {
        let map = generate(WORLD_SIZE, 1);
        assert!(share(&map, Terrain::ShallowWater) > SHALLOW_WATER_SHARE + 0.005);
    }

    #[test]
    fn longest_candidate_rivers_are_picked() {
        let map = generate(WORLD_SIZE, 1);
        let rivers = pick_rivers(&map, 1);
        assert!(!rivers.is_empty() && rivers.len() <= RIVER_COUNT);
        assert!(rivers.windows(2).all(|w| w[0].len() >= w[1].len()));
    }

    #[test]
    fn quantile_picks_value_at_share() {
        let values: Vec<f32> = (0..10).rev().map(|v| v as f32).collect();
        assert_eq!(quantile(&values, 0.0), 0.0);
        assert_eq!(quantile(&values, 0.3), 3.0);
        assert_eq!(quantile(&values, 1.0), 9.0);
    }

    #[test]
    fn largest_region_picks_biggest_walkable_area() {
        let mut map = WorldMap::filled(UVec2::new(6, 3), Terrain::Grass);
        for y in 0..3 {
            map.set(UVec2::new(2, y), Terrain::Stone);
        }
        let region = largest_region(&map);
        assert_eq!(region.len(), 9);
        assert!(region.iter().all(|p| p.x > 2));
    }

    // Colonists must start somewhere they can all reach each other and most of the map,
    // not on a small island cut off by deep water or stone.
    #[test]
    fn spawn_tiles_are_distinct_walkable_and_mutually_reachable() {
        for seed in 0..5 {
            let map = generate(SIZE, seed);
            let tiles = spawn_tiles(&map, 3);
            assert_eq!(tiles.len(), 3);
            let region = largest_region(&map);
            for t in &tiles {
                assert!(region.contains(t));
                assert!(find_path(&map, tiles[0], *t).is_some());
            }
            let mut unique = tiles.clone();
            unique.dedup();
            assert_eq!(unique.len(), 3);
        }
    }
}