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

use bevy::prelude::*;

use crate::world::{CHEAPEST_MOVE_COST, WorldMap};

const STRAIGHT: u32 = 10;
const DIAGONAL: u32 = 14;
const MIN_TERRAIN_COST: u32 = CHEAPEST_MOVE_COST;

const NEIGHBORS: [IVec2; 8] = [
    IVec2::new(1, 0),
    IVec2::new(-1, 0),
    IVec2::new(0, 1),
    IVec2::new(0, -1),
    IVec2::new(1, 1),
    IVec2::new(1, -1),
    IVec2::new(-1, 1),
    IVec2::new(-1, -1),
];

pub fn find_path(map: &WorldMap, start: UVec2, goal: UVec2) -> Option<Vec<UVec2>> {
    map.move_cost(goal.as_ivec2())?;
    let goal = goal.as_ivec2();
    astar(map, start, |p| p == goal, |p| octile(p, goal))
}

pub fn find_path_adjacent(map: &WorldMap, start: UVec2, target: UVec2) -> Option<Vec<UVec2>> {
    let target = target.as_ivec2();
    astar(
        map,
        start,
        |p| is_adjacent(p, target),
        |p| octile(p, target).saturating_sub(DIAGONAL * MIN_TERRAIN_COST),
    )
}

pub fn is_adjacent(a: IVec2, b: IVec2) -> bool {
    a != b && (a - b).abs().max_element() == 1
}

pub fn can_step(map: &WorldMap, from: IVec2, step: IVec2) -> bool {
    let next = from + step;
    map.move_cost(next).is_some()
        && (step.x == 0
            || step.y == 0
            || (map.move_cost(from + IVec2::new(step.x, 0)).is_some()
                && map.move_cost(from + IVec2::new(0, step.y)).is_some()))
}

pub fn nearest_tile(
    map: &WorldMap,
    start: UVec2,
    max_tiles: usize,
    mut is_target: impl FnMut(UVec2) -> bool,
) -> Option<UVec2> {
    let index = |p: UVec2| (p.y * map.size.x + p.x) as usize;
    let mut seen = vec![false; (map.size.x * map.size.y) as usize];
    let mut seen_count = 1;
    if let Some(cell) = seen.get_mut(index(start)) {
        *cell = true;
    }
    let mut queue = VecDeque::from([start]);
    while let Some(pos) = queue.pop_front() {
        if is_target(pos) {
            return Some(pos);
        }
        if seen_count >= max_tiles {
            continue;
        }
        for step in NEIGHBORS {
            let next = pos.as_ivec2() + step;
            if can_step(map, pos.as_ivec2(), step) && !seen[index(next.as_uvec2())] {
                seen[index(next.as_uvec2())] = true;
                seen_count += 1;
                queue.push_back(next.as_uvec2());
            }
        }
    }
    None
}

#[derive(Default, Debug, Clone)]
pub struct Reach {
    width: u32,
    seen: Vec<bool>,
}

#[derive(Resource, Default, Debug, Clone)]
pub struct ColonyReach(pub Option<Reach>);

impl ColonyReach {
    pub fn allows(&self, tile: UVec2) -> bool {
        self.0.as_ref().is_none_or(|reach| reach.contains(tile))
    }

    pub fn allows_near(&self, tile: UVec2) -> bool {
        self.0.as_ref().is_none_or(|reach| reach.near(tile))
    }
}

impl Reach {
    pub fn flood(map: &WorldMap, start: UVec2) -> Self {
        let width = map.size.x;
        let mut seen = vec![false; map.size.element_product() as usize];
        let index = |p: UVec2| (p.y * width + p.x) as usize;
        if map.move_cost(start.as_ivec2()).is_some() {
            seen[index(start)] = true;
            let mut queue = VecDeque::from([start]);
            while let Some(pos) = queue.pop_front() {
                for step in NEIGHBORS {
                    let next = pos.as_ivec2() + step;
                    if can_step(map, pos.as_ivec2(), step) && !seen[index(next.as_uvec2())] {
                        seen[index(next.as_uvec2())] = true;
                        queue.push_back(next.as_uvec2());
                    }
                }
            }
        }
        Self { width, seen }
    }

    pub fn contains(&self, p: UVec2) -> bool {
        p.x < self.width
            && self
                .seen
                .get((p.y * self.width + p.x) as usize)
                .copied()
                .unwrap_or(false)
    }

    pub fn near(&self, p: UVec2) -> bool {
        (-1..=1).any(|dy| {
            (-1..=1).any(|dx| {
                let q = p.as_ivec2() + IVec2::new(dx, dy);
                q.min_element() >= 0 && self.contains(q.as_uvec2())
            })
        })
    }
}

pub fn search_reach(max_tiles: usize) -> u32 {
    ((max_tiles as f32).sqrt() / 2.0).ceil() as u32
}

pub fn nearest_of(
    map: &WorldMap,
    start: UVec2,
    max_tiles: usize,
    targets: &HashSet<UVec2>,
) -> Option<UVec2> {
    let reach = search_reach(max_tiles) as i32;
    let in_reach = |p: &&UVec2| (p.as_ivec2() - start.as_ivec2()).abs().max_element() <= reach;
    if !targets.iter().any(|p| in_reach(&p)) {
        return None;
    }
    nearest_tile(map, start, max_tiles, |p| targets.contains(&p))
}

fn octile(a: IVec2, b: IVec2) -> u32 {
    let d = (a - b).abs();
    let (min, max) = (d.x.min(d.y) as u32, d.x.max(d.y) as u32);
    (DIAGONAL * min + STRAIGHT * (max - min)) * MIN_TERRAIN_COST
}

#[derive(Default)]
struct Scratch {
    generation: u32,
    stamp: Vec<u32>,
    cost: Vec<u32>,
    from: Vec<IVec2>,
    open: BinaryHeap<Reverse<(u32, i32, i32)>>,
}

thread_local! {
    static SCRATCH: std::cell::RefCell<Scratch> = std::cell::RefCell::new(Scratch::default());
}

fn astar(
    map: &WorldMap,
    start: UVec2,
    is_goal: impl Fn(IVec2) -> bool,
    heuristic: impl Fn(IVec2) -> u32,
) -> Option<Vec<UVec2>> {
    astar_limited(map, start, is_goal, heuristic, usize::MAX)
}

fn astar_limited(
    map: &WorldMap,
    start: UVec2,
    is_goal: impl Fn(IVec2) -> bool,
    heuristic: impl Fn(IVec2) -> u32,
    max_expanded: usize,
) -> Option<Vec<UVec2>> {
    let start = start.as_ivec2();
    if !map.in_bounds(start) {
        return None;
    }
    if is_goal(start) {
        return Some(vec![]);
    }
    SCRATCH.with_borrow_mut(|scratch| {
        let tiles = map.size.element_product() as usize;
        if scratch.stamp.len() != tiles {
            *scratch = Scratch {
                stamp: vec![0; tiles],
                cost: vec![0; tiles],
                from: vec![IVec2::ZERO; tiles],
                ..default()
            };
        }
        scratch.generation = scratch.generation.wrapping_add(1);
        if scratch.generation == 0 {
            scratch.stamp.fill(0);
            scratch.generation = 1;
        }
        let Scratch {
            generation,
            stamp,
            cost,
            from,
            open,
        } = scratch;
        let generation = *generation;
        let width = map.size.x as usize;
        let index = |p: IVec2| p.y as usize * width + p.x as usize;
        open.clear();
        stamp[index(start)] = generation;
        cost[index(start)] = 0;
        from[index(start)] = start;
        open.push(Reverse((heuristic(start), start.x, start.y)));
        let mut expanded = 0;
        while let Some(Reverse((_, x, y))) = open.pop() {
            let current = IVec2::new(x, y);
            if is_goal(current) {
                let mut path = vec![current.as_uvec2()];
                let mut at = current;
                while from[index(at)] != start {
                    at = from[index(at)];
                    path.push(at.as_uvec2());
                }
                path.reverse();
                return Some(path);
            }
            expanded += 1;
            if expanded > max_expanded {
                return None;
            }
            let current_cost = cost[index(current)];
            for step in NEIGHBORS {
                let next = current + step;
                if !can_step(map, current, step) {
                    continue;
                }
                let Some(terrain_cost) = map.step_cost(current, next) else {
                    continue;
                };
                let diagonal = step.x != 0 && step.y != 0;
                let base = if diagonal { DIAGONAL } else { STRAIGHT };
                let new_cost = current_cost + base * terrain_cost;
                let i = index(next);
                if stamp[i] != generation || new_cost < cost[i] {
                    stamp[i] = generation;
                    cost[i] = new_cost;
                    from[i] = current;
                    open.push(Reverse((new_cost + heuristic(next), next.x, next.y)));
                }
            }
        }
        None
    })
}

pub const WANDER_SEARCH: usize = 800;

pub fn find_nearby_path(
    map: &WorldMap,
    start: UVec2,
    goal: UVec2,
    max_expanded: usize,
) -> Option<Vec<UVec2>> {
    map.move_cost(goal.as_ivec2())?;
    let goal = goal.as_ivec2();
    astar_limited(map, start, |p| p == goal, |p| octile(p, goal), max_expanded)
}

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

    // Wandering picks random nearby goals; one walled off by a river or cliff made A* search the
    // whole connected region before failing. Short trips cap the search instead.
    #[test]
    fn nearby_paths_give_up_early_on_unreachable_goals() {
        let mut map = WorldMap::filled(UVec2::splat(100), Terrain::Grass);
        for y in 0..100 {
            map.set(UVec2::new(50, y), Terrain::DeepWater);
        }
        let start = UVec2::new(45, 50);
        assert!(find_nearby_path(&map, start, UVec2::new(40, 52), 200).is_some());
        assert_eq!(find_nearby_path(&map, start, UVec2::new(55, 50), 200), None);
        assert_eq!(find_path(&map, start, UVec2::new(55, 50)), None);
    }

    // Target searches must use the same no-corner-cutting rule as A*. A stack on a tile touching
    // the colony only diagonally between two rocks was "found" by the search, A* then explored
    // the whole map before failing, and the job was picked again at once: ~10 ms, many times a
    // second, which made Max speed crawl on a 50-day colony.
    #[test]
    fn searches_never_pick_targets_that_paths_cannot_reach() {
        let mut map = WorldMap::filled(UVec2::splat(6), Terrain::Grass);
        for y in 3..6 {
            for x in 0..6 {
                map.set(UVec2::new(x, y), Terrain::Stone);
            }
        }
        map.set(UVec2::new(2, 3), Terrain::Grass);
        map.set(UVec2::new(3, 4), Terrain::Gravel);
        let pocket = UVec2::new(3, 4);
        assert!(map.move_cost(IVec2::new(2, 3)).is_some());
        assert_eq!(find_path(&map, UVec2::ZERO, pocket), None);
        let targets = HashSet::from([pocket]);
        assert_eq!(nearest_of(&map, UVec2::ZERO, 100, &targets), None);
    }

    // Idle colonists look for food and cooking input several times per game minute. A stack far
    // beyond what the flood fill can cover made every one of those searches visit the full tile
    // budget and fail, which dominated frame time at high speeds. The reach is a radius: the
    // flood fills a square, so its half-width (not its width, as it once was) bounds it.
    #[test]
    fn target_searches_skip_targets_beyond_their_reach() {
        let map = WorldMap::filled(UVec2::splat(300), Terrain::Grass);
        let start = UVec2::splat(10);
        let max_tiles = 400;
        let near = start + UVec2::new(search_reach(max_tiles) / 4, 0);
        let far = start + UVec2::new(search_reach(max_tiles) + 1, 0);
        assert_eq!(
            nearest_of(&map, start, max_tiles, &HashSet::from([far])),
            None
        );
        assert_eq!(
            nearest_of(&map, start, max_tiles, &HashSet::from([far, near])),
            Some(near)
        );
        let edge = start + UVec2::new(search_reach(max_tiles) - 1, 0);
        assert_eq!(
            nearest_of(&map, start, max_tiles, &HashSet::from([edge])),
            Some(edge)
        );
    }

    fn map_from(rows: &[&str]) -> WorldMap {
        let size = UVec2::new(rows[0].len() as u32, rows.len() as u32);
        let mut map = WorldMap::filled(size, Terrain::Grass);
        for (row, line) in rows.iter().enumerate() {
            let y = size.y - 1 - row as u32;
            for (x, c) in line.chars().enumerate() {
                let terrain = match c {
                    '#' => Terrain::Stone,
                    '~' => Terrain::ShallowWater,
                    _ => Terrain::Grass,
                };
                map.set(UVec2::new(x as u32, y), terrain);
            }
        }
        map
    }

    fn is_connected(start: UVec2, path: &[UVec2]) -> bool {
        std::iter::once(start)
            .chain(path.iter().copied())
            .collect::<Vec<_>>()
            .windows(2)
            .all(|w| (w[0].as_ivec2() - w[1].as_ivec2()).abs().max_element() == 1)
    }

    #[test]
    fn path_excludes_start_and_ends_at_goal() {
        let map = map_from(&["....."]);
        let path = find_path(&map, UVec2::new(0, 0), UVec2::new(4, 0)).unwrap();
        assert_eq!(
            path,
            vec![
                UVec2::new(1, 0),
                UVec2::new(2, 0),
                UVec2::new(3, 0),
                UVec2::new(4, 0)
            ]
        );
    }

    #[test]
    fn path_to_self_is_empty() {
        let map = map_from(&["..."]);
        assert_eq!(find_path(&map, UVec2::ZERO, UVec2::ZERO), Some(vec![]));
    }

    #[test]
    fn path_goes_around_walls() {
        let map = map_from(&[
            ".....", //
            ".###.", ".#.#.", "...#.",
        ]);
        let start = UVec2::new(2, 1);
        let goal = UVec2::new(4, 0);
        let path = find_path(&map, start, goal).unwrap();
        assert!(is_connected(start, &path));
        assert_eq!(path.last(), Some(&goal));
        assert!(path.iter().all(|p| map.move_cost(p.as_ivec2()).is_some()));
    }

    #[test]
    fn unreachable_or_impassable_goal_has_no_path() {
        let map = map_from(&[
            "..#..", //
            "..#..",
        ]);
        assert_eq!(find_path(&map, UVec2::new(0, 0), UVec2::new(4, 0)), None);
        assert_eq!(find_path(&map, UVec2::new(0, 0), UVec2::new(2, 0)), None);
    }

    // Colonists move in straight lines between tile centers, so a diagonal step between two
    // blocked orthogonal neighbours would visually clip through the corners of walls.
    #[test]
    fn diagonal_steps_never_cut_wall_corners() {
        let map = map_from(&[
            "#.", //
            ".#",
        ]);
        assert_eq!(find_path(&map, UVec2::new(0, 0), UVec2::new(1, 1)), None);
        let map = map_from(&[
            "..", //
            "#.",
        ]);
        let path = find_path(&map, UVec2::new(0, 1), UVec2::new(1, 0)).unwrap();
        assert_eq!(path, vec![UVec2::new(1, 1), UVec2::new(1, 0)]);
    }

    #[test]
    fn prefers_detour_over_slow_terrain() {
        let map = map_from(&[
            ".....", //
            "..~..",
        ]);
        let path = find_path(&map, UVec2::new(0, 0), UVec2::new(4, 0)).unwrap();
        assert!(!path.contains(&UVec2::new(2, 0)));
    }

    #[test]
    fn walks_around_a_hill_rather_than_over_it() {
        let mut map = map_from(&[
            ".....", //
            ".....",
        ]);
        map.set_height(UVec2::new(2, 0), 10.0);
        let path = find_path(&map, UVec2::new(0, 0), UVec2::new(4, 0)).unwrap();
        assert!(!path.contains(&UVec2::new(2, 0)));
    }

    #[test]
    fn detours_along_paths_when_they_are_faster() {
        let mut map = map_from(&[
            "..........", //
            "..........",
        ]);
        for x in 0..10 {
            map.set(UVec2::new(x, 1), Terrain::Paved);
        }
        let path = find_path(&map, UVec2::new(0, 0), UVec2::new(9, 0)).unwrap();
        assert!(path.iter().filter(|p| p.y == 1).count() >= 6);
    }

    // A wall can go up on a tile a colonist is standing on (or just about to enter); routes must
    // still start from that tile so nobody gets permanently stuck inside a wall.
    #[test]
    fn colonists_can_walk_out_of_a_blocked_tile() {
        let mut map = map_from(&["....."]);
        map.set_wall(UVec2::new(2, 0), true);
        let path = find_path(&map, UVec2::new(2, 0), UVec2::new(4, 0)).unwrap();
        assert_eq!(path, vec![UVec2::new(3, 0), UVec2::new(4, 0)]);
    }

    #[test]
    fn crosses_slow_terrain_when_it_is_the_only_way() {
        let map = map_from(&["..~.."]);
        let path = find_path(&map, UVec2::new(0, 0), UVec2::new(4, 0)).unwrap();
        assert!(path.contains(&UVec2::new(2, 0)));
    }

    #[test]
    fn path_adjacent_stops_next_to_target() {
        let map = map_from(&["......"]);
        let path = find_path_adjacent(&map, UVec2::new(0, 0), UVec2::new(5, 0)).unwrap();
        assert_eq!(path.last(), Some(&UVec2::new(4, 0)));
    }

    #[test]
    fn path_adjacent_works_for_impassable_targets() {
        let map = map_from(&["...#"]);
        let path = find_path_adjacent(&map, UVec2::new(0, 0), UVec2::new(3, 0)).unwrap();
        assert_eq!(path.last(), Some(&UVec2::new(2, 0)));
    }

    #[test]
    fn path_adjacent_from_neighbor_is_empty() {
        let map = map_from(&["..."]);
        let path = find_path_adjacent(&map, UVec2::new(1, 0), UVec2::new(2, 0));
        assert_eq!(path, Some(vec![]));
    }

    #[test]
    fn path_adjacent_to_enclosed_target_is_none() {
        let walled = map_from(&[
            "#####", //
            "#...#", "#####", ".....",
        ]);
        assert_eq!(
            find_path_adjacent(&walled, UVec2::new(0, 0), UVec2::new(2, 2)),
            None
        );
    }

    #[test]
    fn nearest_tile_finds_closest_match_by_walking_distance() {
        let map = map_from(&[
            ".#...", //
            ".#.#.", "...#.",
        ]);
        let found = nearest_tile(&map, UVec2::new(0, 0), 100, |p| p.x == 4);
        assert_eq!(found, Some(UVec2::new(4, 2)));
        assert_eq!(
            nearest_tile(&map, UVec2::new(0, 0), 100, |p| p.x == 9),
            None
        );
    }

    #[test]
    fn finds_paths_across_a_generated_world() {
        let map = crate::worldgen::generate(crate::worldgen::WORLD_SIZE, 1);
        let region = crate::worldgen::largest_region(&map);
        let (start, goal) = (region[0], region[region.len() - 1]);
        let path = find_path(&map, start, goal).unwrap();
        assert!(is_connected(start, &path));
        assert_eq!(path.last(), Some(&goal));
    }
}