use std::collections::HashSet;

use bevy::prelude::*;

use crate::{
    actions::Action,
    selection::{Cursor, update_cursor},
    ui::PointerOverUi,
    world::{TILE_SIZE, WorldMap},
    worldgen::{WorldGenerated, spawn_tiles},
};

pub struct ZonesPlugin;

impl Plugin for ZonesPlugin {
    fn build(&self, app: &mut App) {
        app.init_resource::<Zones>()
            .init_resource::<ZoneTool>()
            .add_systems(
                Update,
                place_starting_stockpile.run_if(on_message::<WorldGenerated>),
            )
            .add_systems(
                PostUpdate,
                (zone_tool, draw_zone_preview).chain().after(update_cursor),
            );
    }
}

const STARTING_STOCKPILE_OFFSET: IVec2 = IVec2::new(2, -2);
const STARTING_STOCKPILE_SIZE: IVec2 = IVec2::new(4, 4);
const PREVIEW_ADD_COLOR: Color = Color::srgba(0.4, 0.7, 1.0, 0.9);
const PREVIEW_REMOVE_COLOR: Color = Color::srgba(1.0, 0.4, 0.4, 0.9);

#[derive(Resource, Default, Debug, Clone, PartialEq)]
pub struct Zones {
    stockpile: HashSet<UVec2>,
}

impl Zones {
    pub fn is_stockpile(&self, tile: UVec2) -> bool {
        self.stockpile.contains(&tile)
    }

    #[cfg(test)]
    pub fn stockpile_tiles(&self) -> impl Iterator<Item = UVec2> + '_ {
        self.stockpile.iter().copied()
    }

    pub fn has_stockpile(&self) -> bool {
        !self.stockpile.is_empty()
    }

    pub fn empty_stockpile_tiles(&self, occupied: &HashSet<UVec2>) -> usize {
        self.stockpile.difference(occupied).count()
    }

    pub fn add_stockpile(&mut self, map: &WorldMap, a: IVec2, b: IVec2) {
        self.add_tiles(map, rect_tiles(a, b).filter_map(|t| t.try_into().ok()));
    }

    pub fn add_tiles(&mut self, map: &WorldMap, tiles: impl IntoIterator<Item = UVec2>) {
        for tile in tiles {
            if can_zone(map, tile.as_ivec2()) {
                self.stockpile.insert(tile);
            }
        }
    }

    pub fn remove_stockpile(&mut self, a: IVec2, b: IVec2) {
        for tile in rect_tiles(a, b) {
            if tile.x >= 0 && tile.y >= 0 {
                self.stockpile.remove(&tile.as_uvec2());
            }
        }
    }
}

pub fn plan_expansion(map: &WorldMap, zones: &Zones, count: usize) -> Vec<UVec2> {
    if zones.stockpile.is_empty() {
        return Vec::new();
    }
    let centroid =
        zones.stockpile.iter().map(|t| t.as_vec2()).sum::<Vec2>() / zones.stockpile.len() as f32;
    let preference = |t: UVec2| {
        (
            map.plant(t.as_ivec2()).is_some(),
            t.as_vec2().distance_squared(centroid),
            t.y,
            t.x,
        )
    };
    let mut planned: Vec<UVec2> = Vec::new();
    while planned.len() < count {
        let taken = |t: UVec2| zones.is_stockpile(t) || planned.contains(&t);
        let best = zones
            .stockpile
            .iter()
            .chain(planned.iter())
            .flat_map(|t| {
                [IVec2::X, IVec2::NEG_X, IVec2::Y, IVec2::NEG_Y].map(|d| t.as_ivec2() + d)
            })
            .filter(|&t| can_zone(map, t) && !taken(t.as_uvec2()))
            .map(|t| t.as_uvec2())
            .min_by(|&a, &b| {
                let (ka, kb) = (preference(a), preference(b));
                ka.0.cmp(&kb.0)
                    .then(ka.1.total_cmp(&kb.1))
                    .then((ka.2, ka.3).cmp(&(kb.2, kb.3)))
            });
        match best {
            Some(tile) => planned.push(tile),
            None => break,
        }
    }
    planned
}

pub fn can_zone(map: &WorldMap, tile: IVec2) -> bool {
    map.move_cost(tile).is_some() && !map.get(tile).is_some_and(|t| t.is_water())
}

fn rect_tiles(a: IVec2, b: IVec2) -> impl Iterator<Item = IVec2> {
    let (min, max) = (a.min(b), a.max(b));
    (min.y..=max.y).flat_map(move |y| (min.x..=max.x).map(move |x| IVec2::new(x, y)))
}

#[derive(Resource, Default, Debug)]
pub struct ZoneTool {
    pub active: bool,
    drag_start: Option<UVec2>,
}

fn place_starting_stockpile(map: Res<WorldMap>, mut zones: ResMut<Zones>) {
    *zones = Zones::default();
    if let Some(&spawn) = spawn_tiles(&map, 1).first() {
        let corner = spawn.as_ivec2() + STARTING_STOCKPILE_OFFSET;
        zones.add_stockpile(&map, corner, corner + STARTING_STOCKPILE_SIZE - 1);
    }
}

fn removing(keys: &ButtonInput<KeyCode>) -> bool {
    keys.any_pressed([KeyCode::ShiftLeft, KeyCode::ShiftRight])
}

fn zone_tool(
    mut actions: MessageReader<Action>,
    keys: Res<ButtonInput<KeyCode>>,
    buttons: Res<ButtonInput<MouseButton>>,
    over_ui: Res<PointerOverUi>,
    cursor: Res<Cursor>,
    map: Res<WorldMap>,
    mut tool: ResMut<ZoneTool>,
    mut zones: ResMut<Zones>,
) {
    for action in actions.read() {
        match action {
            Action::ToggleStockpileTool => tool.active = !tool.active,
            Action::Cancel => tool.active = false,
            _ => continue,
        }
        tool.drag_start = None;
        return;
    }
    if !tool.active {
        return;
    }
    if buttons.just_pressed(MouseButton::Left) && !over_ui.0 {
        tool.drag_start = cursor.tile;
    }
    if !buttons.just_released(MouseButton::Left) {
        return;
    }
    if let (Some(start), Some(end)) = (tool.drag_start.take(), cursor.tile) {
        if removing(&keys) {
            zones.remove_stockpile(start.as_ivec2(), end.as_ivec2());
        } else {
            zones.add_stockpile(&map, start.as_ivec2(), end.as_ivec2());
        }
    }
}

fn draw_zone_preview(
    keys: Res<ButtonInput<KeyCode>>,
    cursor: Res<Cursor>,
    tool: Res<ZoneTool>,
    mut gizmos: Gizmos,
) {
    let (Some(start), Some(end)) = (tool.drag_start, cursor.tile) else {
        return;
    };
    let (min, max) = (start.min(end), start.max(end));
    let center = (WorldMap::tile_center(min) + WorldMap::tile_center(max)) / 2.0;
    let size = (max - min + 1).as_vec2() * TILE_SIZE;
    let color = if removing(&keys) {
        PREVIEW_REMOVE_COLOR
    } else {
        PREVIEW_ADD_COLOR
    };
    gizmos.rect_2d(center, size, color);
}

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

    #[test]
    fn stockpiles_only_cover_walkable_dry_land() {
        let mut map = WorldMap::filled(UVec2::new(4, 1), Terrain::Grass);
        map.set(UVec2::new(1, 0), Terrain::Stone);
        map.set(UVec2::new(2, 0), Terrain::ShallowWater);
        let mut zones = Zones::default();
        zones.add_stockpile(&map, IVec2::new(3, 0), IVec2::new(-1, 0));
        let mut tiles: Vec<_> = zones.stockpile_tiles().collect();
        tiles.sort_by_key(|t| t.x);
        assert_eq!(tiles, vec![UVec2::new(0, 0), UVec2::new(3, 0)]);
    }

    #[test]
    fn expansion_grows_compactly_around_the_stockpile() {
        let map = WorldMap::filled(UVec2::splat(10), Terrain::Grass);
        let mut zones = Zones::default();
        zones.add_stockpile(&map, IVec2::new(4, 4), IVec2::new(5, 5));
        let plan = plan_expansion(&map, &zones, 8);
        assert_eq!(plan.len(), 8);
        assert!(plan.iter().all(|t| !zones.is_stockpile(*t)));
        let beyond_one_ring = plan_expansion(&map, &zones, 20);
        assert_eq!(beyond_one_ring.len(), 20);
        assert!(
            plan.iter()
                .all(|t| (3..=6).contains(&t.x) && (3..=6).contains(&t.y))
        );
    }

    // Trees can be zoned, but stockpiling on open ground keeps forests intact and avoids
    // colonists cutting trees just to make room, so expansion only takes trees as a last resort.
    #[test]
    fn expansion_prefers_open_ground_over_trees_and_skips_unzonable_tiles() {
        let mut map = WorldMap::filled(UVec2::splat(10), Terrain::Grass);
        map.set(UVec2::new(3, 4), Terrain::Stone);
        map.set_plant(UVec2::new(5, 4), Some(crate::world::Plant::Tree));
        let mut zones = Zones::default();
        zones.add_stockpile(&map, IVec2::new(4, 4), IVec2::new(4, 4));
        let plan = plan_expansion(&map, &zones, 2);
        assert_eq!(plan.len(), 2);
        assert!(!plan.contains(&UVec2::new(3, 4)));
        assert!(!plan.contains(&UVec2::new(5, 4)));
    }

    #[test]
    fn no_expansion_without_a_stockpile() {
        let map = WorldMap::filled(UVec2::splat(10), Terrain::Grass);
        assert!(plan_expansion(&map, &Zones::default(), 8).is_empty());
    }

    #[test]
    fn removing_clears_a_rectangle() {
        let map = WorldMap::filled(UVec2::new(4, 4), Terrain::Grass);
        let mut zones = Zones::default();
        zones.add_stockpile(&map, IVec2::ZERO, IVec2::splat(3));
        zones.remove_stockpile(IVec2::new(1, 1), IVec2::new(2, 2));
        assert_eq!(zones.stockpile_tiles().count(), 12);
        assert!(!zones.is_stockpile(UVec2::new(1, 2)));
        assert!(zones.is_stockpile(UVec2::new(0, 2)));
    }

    fn tool_app() -> App {
        let mut app = App::new();
        app.insert_resource(WorldMap::filled(UVec2::splat(8), Terrain::Grass))
            .init_resource::<Zones>()
            .init_resource::<ButtonInput<KeyCode>>()
            .init_resource::<ButtonInput<MouseButton>>()
            .insert_resource(ZoneTool {
                active: true,
                drag_start: None,
            })
            .init_resource::<Cursor>()
            .init_resource::<PointerOverUi>()
            .add_message::<Action>()
            .add_systems(Update, zone_tool);
        app
    }

    fn drag(app: &mut App, from: UVec2, to: UVec2) {
        app.world_mut().resource_mut::<Cursor>().tile = Some(from);
        app.world_mut()
            .resource_mut::<ButtonInput<MouseButton>>()
            .press(MouseButton::Left);
        app.update();
        app.world_mut().resource_mut::<Cursor>().tile = Some(to);
        let mut buttons = app.world_mut().resource_mut::<ButtonInput<MouseButton>>();
        buttons.clear();
        buttons.release(MouseButton::Left);
        app.update();
        app.world_mut()
            .resource_mut::<ButtonInput<MouseButton>>()
            .clear();
    }

    #[test]
    fn dragging_adds_and_shift_dragging_removes() {
        let mut app = tool_app();
        drag(&mut app, UVec2::new(1, 1), UVec2::new(3, 2));
        assert_eq!(app.world().resource::<Zones>().stockpile_tiles().count(), 6);
        app.world_mut()
            .resource_mut::<ButtonInput<KeyCode>>()
            .press(KeyCode::ShiftLeft);
        drag(&mut app, UVec2::new(3, 2), UVec2::new(3, 1));
        assert_eq!(app.world().resource::<Zones>().stockpile_tiles().count(), 4);
    }

    #[test]
    fn dragging_does_nothing_when_the_tool_is_off() {
        let mut app = tool_app();
        app.world_mut().resource_mut::<ZoneTool>().active = false;
        drag(&mut app, UVec2::new(1, 1), UVec2::new(3, 2));
        assert_eq!(app.world().resource::<Zones>().stockpile_tiles().count(), 0);
    }

    #[test]
    fn new_worlds_start_with_a_stockpile() {
        let mut app = App::new();
        app.insert_resource(crate::worldgen::generate(UVec2::new(96, 64), 1))
            .init_resource::<Zones>()
            .add_message::<WorldGenerated>()
            .add_systems(
                Update,
                place_starting_stockpile.run_if(on_message::<WorldGenerated>),
            );
        app.world_mut().write_message(WorldGenerated);
        app.update();
        assert!(app.world().resource::<Zones>().stockpile_tiles().count() > 8);
    }
}