use std::collections::{HashMap, HashSet, VecDeque};

use bevy::prelude::*;

use crate::{
    colonist::Colonist,
    items::{Item, ItemKind},
    sim::{SimSystems, SimTick},
    structures::{Structure, StructureKind},
    world::{TILE_SIZE, TilePos, WorldMap},
    worldgen::WorldGenerated,
    zones::{Zones, can_zone},
};

pub struct BuildingPlugin;

impl Plugin for BuildingPlugin {
    fn build(&self, app: &mut App) {
        app.init_resource::<Blueprints>()
            .add_systems(
                Update,
                (
                    clear_blueprints.run_if(on_message::<WorldGenerated>),
                    draw_blueprints,
                ),
            )
            .add_systems(FixedUpdate, plan_shelter.in_set(SimSystems::Decide));
    }
}

pub const MAX_ROOM_TILES: usize = 120;
const PLAN_INTERVAL_TICKS: u64 = 300;
const MAX_ROOM_MARGIN: i32 = 2;
const CRAMPED_SIDE_COST: i32 = 3;
const WALL_BLUEPRINT_COLOR: Color = Color::srgba(0.5, 0.75, 1.0, 0.8);
const DOOR_BLUEPRINT_COLOR: Color = Color::srgba(1.0, 0.8, 0.4, 0.8);
const ROOF_BLUEPRINT_COLOR: Color = Color::srgba(0.8, 0.6, 1.0, 0.35);

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum BuildKind {
    Structure(StructureKind),
    Roof,
}

impl BuildKind {
    pub fn cost(self) -> (ItemKind, u32) {
        match self {
            BuildKind::Structure(kind) => kind.cost(),
            BuildKind::Roof => (ItemKind::Wood, 1),
        }
    }

    pub fn work(self) -> u32 {
        match self {
            BuildKind::Structure(StructureKind::Campfire) => 360,
            BuildKind::Structure(StructureKind::Torch) => 120,
            BuildKind::Structure(StructureKind::Bed) => 600,
            BuildKind::Structure(StructureKind::Wall) => 300,
            BuildKind::Structure(StructureKind::Door) => 300,
            BuildKind::Roof => 120,
        }
    }

    pub fn name(self) -> &'static str {
        match self {
            BuildKind::Structure(kind) => kind.name(),
            BuildKind::Roof => "roof",
        }
    }

    pub fn blocks_tile(self) -> bool {
        self == BuildKind::Structure(StructureKind::Wall)
    }
}

#[derive(Resource, Default, Debug, Clone)]
pub struct Blueprints(pub HashMap<UVec2, BuildKind>);

impl Blueprints {
    pub fn get(&self, tile: UVec2) -> Option<BuildKind> {
        self.0.get(&tile).copied()
    }

    pub fn tiles(&self) -> HashSet<UVec2> {
        self.0.keys().copied().collect()
    }
}

fn clear_blueprints(mut blueprints: ResMut<Blueprints>) {
    blueprints.0.clear();
}

pub fn enclosed_room(map: &WorldMap, doors: &HashSet<UVec2>, start: UVec2) -> Option<Vec<UVec2>> {
    let mut seen = HashSet::from([start]);
    let mut queue = VecDeque::from([start]);
    while let Some(tile) = queue.pop_front() {
        if seen.len() > MAX_ROOM_TILES {
            return None;
        }
        for step in [IVec2::X, IVec2::NEG_X, IVec2::Y, IVec2::NEG_Y] {
            let next = tile.as_ivec2() + step;
            if !map.in_bounds(next) {
                return None;
            }
            let next = next.as_uvec2();
            let barrier = map.is_wall(next.as_ivec2())
                || doors.contains(&next)
                || map.get(next.as_ivec2()).is_some_and(|t| t.is_rock());
            if !barrier && seen.insert(next) {
                queue.push_back(next);
            }
        }
    }
    let mut room: Vec<UVec2> = seen.into_iter().collect();
    room.sort_by_key(|t| (t.y, t.x));
    Some(room)
}

pub struct ShelterContext<'a> {
    pub map: &'a WorldMap,
    pub structures: &'a [(UVec2, StructureKind)],
    pub items: &'a [UVec2],
}

impl ShelterContext<'_> {
    fn doors(&self) -> HashSet<UVec2> {
        self.structures
            .iter()
            .filter(|(_, kind)| *kind == StructureKind::Door)
            .map(|(pos, _)| *pos)
            .collect()
    }

    fn wall_buildable(&self, tile: IVec2) -> bool {
        can_zone(self.map, tile)
            && self.map.plant(tile).is_none()
            && !self.items.contains(&tile.as_uvec2())
            && !self
                .structures
                .iter()
                .any(|(pos, _)| *pos == tile.as_uvec2())
    }

    fn already_closed(&self, tile: IVec2, doors: &HashSet<UVec2>) -> bool {
        self.map.is_wall(tile)
            || doors.contains(&tile.as_uvec2())
            || self.map.get(tile).is_some_and(|t| t.is_rock())
    }

    pub fn plan(&self, beds: &[UVec2], toward: UVec2) -> Option<HashMap<UVec2, BuildKind>> {
        let doors = self.doors();
        if let Some(&first) = beds.first()
            && let Some(room) = enclosed_room(self.map, &doors, first)
        {
            let roofs: HashMap<UVec2, BuildKind> = room
                .into_iter()
                .filter(|t| !self.map.is_roofed(t.as_ivec2()))
                .map(|t| (t, BuildKind::Roof))
                .collect();
            return (!roofs.is_empty()).then_some(roofs);
        }
        let min = beds.iter().copied().reduce(|a, b| a.min(b))?.as_ivec2();
        let max = beds.iter().copied().reduce(|a, b| a.max(b))?.as_ivec2();
        let margins = 0..=MAX_ROOM_MARGIN;
        let shapes = margins.clone().flat_map(|l| {
            let margins = margins.clone();
            margins.clone().flat_map(move |r| {
                let margins = margins.clone();
                margins
                    .clone()
                    .flat_map(move |d| margins.clone().map(move |u| IVec4::new(l, r, d, u)))
            })
        });
        shapes
            .filter_map(|m| {
                let lo = min - IVec2::new(m.x, m.z) - 1;
                let hi = max + IVec2::new(m.y, m.w) + 1;
                let plan = self.walls_around(lo, hi, toward, &doors)?;
                let cramped = m.to_array().iter().filter(|&&side| side == 0).count() as i32;
                let cost = plan.len() as i32 + CRAMPED_SIDE_COST * cramped;
                Some((cost, m.to_array(), plan))
            })
            .min_by_key(|(cost, margins, _)| (*cost, *margins))
            .map(|(_, _, plan)| plan)
    }

    fn walls_around(
        &self,
        lo: IVec2,
        hi: IVec2,
        toward: UVec2,
        doors: &HashSet<UVec2>,
    ) -> Option<HashMap<UVec2, BuildKind>> {
        let ring: Vec<IVec2> = (lo.y..=hi.y)
            .flat_map(|y| (lo.x..=hi.x).map(move |x| IVec2::new(x, y)))
            .filter(|t| t.x == lo.x || t.x == hi.x || t.y == lo.y || t.y == hi.y)
            .collect();
        let mut plan = HashMap::new();
        for &tile in &ring {
            if !self.map.in_bounds(tile) {
                return None;
            }
            if self.already_closed(tile, doors) {
                continue;
            }
            if !self.wall_buildable(tile) {
                return None;
            }
            plan.insert(tile.as_uvec2(), BuildKind::Structure(StructureKind::Wall));
        }
        let has_door = ring.iter().any(|t| doors.contains(&t.as_uvec2()));
        if !has_door {
            let corner = |t: &UVec2| {
                let t = t.as_ivec2();
                (t.x == lo.x || t.x == hi.x) && (t.y == lo.y || t.y == hi.y)
            };
            let door = plan.keys().filter(|t| !corner(t)).copied().min_by(|a, b| {
                a.as_vec2()
                    .distance_squared(toward.as_vec2())
                    .total_cmp(&b.as_vec2().distance_squared(toward.as_vec2()))
                    .then((a.y, a.x).cmp(&(b.y, b.x)))
            })?;
            plan.insert(door, BuildKind::Structure(StructureKind::Door));
        }
        Some(plan)
    }
}

pub fn plan_shelter(
    tick: Res<SimTick>,
    map: Res<WorldMap>,
    mut zones: ResMut<Zones>,
    structures: Query<(&TilePos, &Structure), Without<Colonist>>,
    items: Query<&TilePos, With<Item>>,
    mut blueprints: ResMut<Blueprints>,
) {
    if !tick.0.is_multiple_of(PLAN_INTERVAL_TICKS) || !blueprints.0.is_empty() {
        return;
    }
    let structures: Vec<(UVec2, StructureKind)> =
        structures.iter().map(|(pos, s)| (pos.0, s.kind)).collect();
    let beds: Vec<UVec2> = structures
        .iter()
        .filter(|(pos, kind)| *kind == StructureKind::Bed && !map.is_roofed(pos.as_ivec2()))
        .map(|(pos, _)| *pos)
        .collect();
    if beds.is_empty() {
        return;
    }
    let items: Vec<UVec2> = items.iter().map(|pos| pos.0).collect();
    let ctx = ShelterContext {
        map: &map,
        structures: &structures,
        items: &items,
    };
    let toward = structures
        .iter()
        .find(|(_, kind)| *kind == StructureKind::Campfire)
        .map(|(pos, _)| *pos)
        .unwrap_or_else(|| zones.home(beds[0]));
    if let Some(plan) = ctx.plan(&beds, toward) {
        for (tile, kind) in &plan {
            if *kind != BuildKind::Roof {
                zones.remove_tile(*tile);
            }
        }
        blueprints.0 = plan;
    }
}

fn draw_blueprints(mut gizmos: Gizmos, blueprints: Res<Blueprints>) {
    for (tile, kind) in &blueprints.0 {
        let color = match kind {
            BuildKind::Structure(StructureKind::Door) => DOOR_BLUEPRINT_COLOR,
            BuildKind::Roof => ROOF_BLUEPRINT_COLOR,
            _ => WALL_BLUEPRINT_COLOR,
        };
        gizmos.rect_2d(
            WorldMap::tile_center(*tile),
            Vec2::splat(TILE_SIZE * 0.85),
            color,
        );
    }
}

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

    fn ctx<'a>(map: &'a WorldMap, structures: &'a [(UVec2, StructureKind)]) -> ShelterContext<'a> {
        ShelterContext {
            map,
            structures,
            items: &[],
        }
    }

    fn count(plan: &HashMap<UVec2, BuildKind>, kind: BuildKind) -> usize {
        plan.values().filter(|k| **k == kind).count()
    }

    const WALL: BuildKind = BuildKind::Structure(StructureKind::Wall);
    const DOOR: BuildKind = BuildKind::Structure(StructureKind::Door);

    #[test]
    fn shelter_is_a_wall_ring_with_one_door_around_the_beds() {
        let map = WorldMap::filled(UVec2::splat(20), Terrain::Grass);
        let beds = [UVec2::new(10, 10), UVec2::new(11, 10)];
        let plan = ctx(&map, &[]).plan(&beds, UVec2::new(10, 2)).unwrap();
        assert_eq!(plan.len(), 6 * 5 - 4 * 3);
        assert_eq!(count(&plan, DOOR), 1);
        let door = plan.iter().find(|(_, k)| **k == DOOR).unwrap().0;
        assert_eq!(door.y, 8);
    }

    // Mountains are free walls: a shelter dug against a cliff only needs walls on the open sides.
    #[test]
    fn rock_counts_as_wall() {
        let mut map = WorldMap::filled(UVec2::splat(20), Terrain::Grass);
        for y in 0..20 {
            map.set(UVec2::new(12, y), Terrain::Stone);
        }
        let plan = ctx(&map, &[])
            .plan(&[UVec2::new(10, 10)], UVec2::new(10, 2))
            .unwrap();
        assert!(plan.keys().all(|t| t.x != 12));
        assert_eq!(count(&plan, WALL) + count(&plan, DOOR), 16 - 5);
    }

    #[test]
    fn shelters_route_around_unbuildable_tiles() {
        let mut map = WorldMap::filled(UVec2::splat(20), Terrain::Grass);
        map.set(UVec2::new(10, 8), Terrain::ShallowWater);
        let plan = ctx(&map, &[])
            .plan(&[UVec2::new(10, 10)], UVec2::new(10, 2))
            .unwrap();
        assert!(!plan.contains_key(&UVec2::new(10, 8)));
        assert_eq!(
            plan.values()
                .filter(|k| **k == BuildKind::Structure(StructureKind::Door))
                .count(),
            1
        );
    }

    // Beds often end up right next to the stockpile; walls may take over empty stockpile tiles
    // (which then stop being storage) instead of the shelter never getting planned at all.
    #[test]
    fn walls_may_replace_empty_stockpile_tiles_but_not_stored_items() {
        let map = WorldMap::filled(UVec2::splat(20), Terrain::Grass);
        let mut zones = Zones::default();
        zones.add_stockpile(&map, IVec2::new(12, 5), IVec2::new(15, 15));
        let items = [UVec2::new(13, 10)];
        let ctx = ShelterContext {
            map: &map,
            structures: &[],
            items: &items,
        };
        let plan = ctx.plan(&[UVec2::new(11, 10)], UVec2::new(5, 10)).unwrap();
        assert!(plan.keys().any(|t| zones.is_stockpile(*t)));
        assert!(!plan.contains_key(&UVec2::new(13, 10)));
    }

    #[test]
    fn enclosed_rooms_get_roof_blueprints() {
        let mut map = WorldMap::filled(UVec2::splat(10), Terrain::Grass);
        for i in 3..=6 {
            for t in [
                UVec2::new(i, 3),
                UVec2::new(i, 6),
                UVec2::new(3, i),
                UVec2::new(6, i),
            ] {
                map.set_wall(t, true);
            }
        }
        map.set_wall(UVec2::new(4, 3), false);
        let structures = [(UVec2::new(4, 3), StructureKind::Door)];
        let plan = ctx(&map, &structures)
            .plan(&[UVec2::new(4, 4)], UVec2::ZERO)
            .unwrap();
        assert_eq!(plan.len(), 4);
        assert!(plan.values().all(|k| *k == BuildKind::Roof));
    }

    #[test]
    fn open_areas_are_not_rooms() {
        let map = WorldMap::filled(UVec2::splat(30), Terrain::Grass);
        assert!(enclosed_room(&map, &HashSet::new(), UVec2::new(15, 15)).is_none());
    }
}