use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};

use bevy::prelude::*;

use crate::{
    actions::Action,
    economy::Stock,
    items::ItemKind,
    pathfinding::nearest_tile,
    plants::fertility,
    selection::{Cursor, update_cursor},
    ui::PointerOverUi,
    world::{TILE_SIZE, Terrain, WorldMap},
    worldgen::WorldGenerated,
};

pub struct ZonesPlugin;

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

pub const NEW_STOCKPILE_SIZE: u32 = 4;
const STOCKPILE_SITE_SEARCH_TILES: usize = 4000;
pub const PREVIEW_ADD_COLOR: Color = Color::srgba(0.4, 0.7, 1.0, 0.9);
pub const PREVIEW_REMOVE_COLOR: Color = Color::srgba(1.0, 0.4, 0.4, 0.9);

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum ZoneKind {
    Stockpile,
    Farm,
}

#[derive(
    Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize,
)]
pub enum StoragePriority {
    Low,
    #[default]
    Normal,
    High,
}

impl StoragePriority {
    pub const DESCENDING: [StoragePriority; 3] = [
        StoragePriority::High,
        StoragePriority::Normal,
        StoragePriority::Low,
    ];

    pub fn next(self) -> Self {
        match self {
            StoragePriority::Low => StoragePriority::Normal,
            StoragePriority::Normal => StoragePriority::High,
            StoragePriority::High => StoragePriority::Low,
        }
    }

    pub fn label(self) -> &'static str {
        match self {
            StoragePriority::Low => "Low",
            StoragePriority::Normal => "Normal",
            StoragePriority::High => "High",
        }
    }
}

#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct StorageRule {
    pub refused: Vec<Stock>,
    pub priority: StoragePriority,
}

impl StorageRule {
    pub fn accepts(&self, kind: ItemKind) -> bool {
        !self.refused.contains(&Stock::of(kind))
    }

    pub fn toggle(&mut self, stock: Stock) {
        if self.refused.contains(&stock) {
            self.refused.retain(|s| *s != stock);
        } else {
            self.refused.push(stock);
        }
    }
}

static DEFAULT_RULE: StorageRule = StorageRule {
    refused: Vec::new(),
    priority: StoragePriority::Normal,
};

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

#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
pub struct ZonesSave {
    stockpile: Vec<UVec2>,
    farm: Vec<UVec2>,
    rules: Vec<(UVec2, StorageRule)>,
}

impl Zones {
    pub fn save(&self) -> ZonesSave {
        ZonesSave {
            stockpile: self.stockpile.iter().copied().collect(),
            farm: self.farm.iter().copied().collect(),
            rules: self.rules.iter().map(|(t, r)| (*t, r.clone())).collect(),
        }
    }

    pub fn load(save: ZonesSave) -> Self {
        Self {
            stockpile: save.stockpile.into_iter().collect(),
            farm: save.farm.into_iter().collect(),
            rules: save.rules.into_iter().collect(),
        }
    }

    pub fn rule(&self, tile: UVec2) -> &StorageRule {
        self.rules.get(&tile).unwrap_or(&DEFAULT_RULE)
    }

    pub fn accepts(&self, tile: UVec2, kind: ItemKind) -> bool {
        self.is_stockpile(tile) && self.rule(tile).accepts(kind)
    }

    pub fn stockpile_region(&self, start: UVec2) -> Vec<UVec2> {
        if !self.is_stockpile(start) {
            return Vec::new();
        }
        let mut seen = HashSet::from([start]);
        let mut stack = vec![start];
        while let Some(tile) = stack.pop() {
            for step in [IVec2::X, IVec2::NEG_X, IVec2::Y, IVec2::NEG_Y] {
                let next = tile.as_ivec2() + step;
                if next.min_element() >= 0
                    && self.is_stockpile(next.as_uvec2())
                    && seen.insert(next.as_uvec2())
                {
                    stack.push(next.as_uvec2());
                }
            }
        }
        seen.into_iter().collect()
    }

    pub fn update_region_rule(&mut self, start: UVec2, change: impl Fn(&mut StorageRule)) {
        let mut rule = self.rule(start).clone();
        change(&mut rule);
        for tile in self.stockpile_region(start) {
            if rule == DEFAULT_RULE {
                self.rules.remove(&tile);
            } else {
                self.rules.insert(tile, rule.clone());
            }
        }
    }

    pub fn is_farm(&self, tile: UVec2) -> bool {
        self.farm.contains(&tile)
    }

    pub fn farm_tiles(&self) -> impl Iterator<Item = UVec2> + '_ {
        self.farm.iter().copied()
    }

    pub fn has_farm(&self) -> bool {
        !self.farm.is_empty()
    }

    pub fn add_farm_tiles(&mut self, map: &WorldMap, tiles: impl IntoIterator<Item = UVec2>) {
        for tile in tiles {
            if can_farm(map, tile.as_ivec2()) && !self.is_stockpile(tile) {
                self.farm.insert(tile);
            }
        }
    }

    pub fn add_zone(&mut self, kind: ZoneKind, map: &WorldMap, a: IVec2, b: IVec2) {
        let tiles = rect_tiles(a, b).filter_map(|t| t.try_into().ok());
        match kind {
            ZoneKind::Stockpile => self.add_tiles(map, tiles),
            ZoneKind::Farm => self.add_farm_tiles(map, tiles),
        }
    }

    pub fn remove_zone(&mut self, kind: ZoneKind, a: IVec2, b: IVec2) {
        let set = match kind {
            ZoneKind::Stockpile => &mut self.stockpile,
            ZoneKind::Farm => &mut self.farm,
        };
        for tile in rect_tiles(a, b) {
            if tile.x >= 0 && tile.y >= 0 {
                set.remove(&tile.as_uvec2());
                if kind == ZoneKind::Stockpile {
                    self.rules.remove(&tile.as_uvec2());
                }
            }
        }
    }

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

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

    pub fn home(&self, fallback: UVec2) -> UVec2 {
        let centroid = self.stockpile.iter().map(|t| t.as_vec2()).sum::<Vec2>()
            / self.stockpile.len().max(1) as f32;
        self.stockpile
            .iter()
            .copied()
            .min_by(|a, b| {
                a.as_vec2()
                    .distance_squared(centroid)
                    .total_cmp(&b.as_vec2().distance_squared(centroid))
                    .then((a.y, a.x).cmp(&(b.y, b.x)))
            })
            .unwrap_or(fallback)
    }

    pub fn remove_tile(&mut self, tile: UVec2) {
        self.stockpile.remove(&tile);
        self.rules.remove(&tile);
    }

    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()
    }

    #[cfg(test)]
    pub fn add_stockpile(&mut self, map: &WorldMap, a: IVec2, b: IVec2) {
        self.add_zone(ZoneKind::Stockpile, map, a, b);
    }

    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.is_farm(tile) {
                self.stockpile.insert(tile);
            }
        }
    }

    #[cfg(test)]
    pub fn remove_stockpile(&mut self, a: IVec2, b: IVec2) {
        self.remove_zone(ZoneKind::Stockpile, a, b);
    }
}

pub fn plan_expansion(
    map: &WorldMap,
    zones: &Zones,
    excluded: &HashSet<UVec2>,
    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)
                || zones.is_farm(t)
                || planned.contains(&t)
                || excluded.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() || t == Terrain::Bridge)
}

pub fn can_farm(map: &WorldMap, tile: IVec2) -> bool {
    can_zone(map, tile) && map.get(tile).is_some_and(|t| fertility(t) > 0.0)
}

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 mode: Option<ZoneKind>,
    drag_start: Option<UVec2>,
}

impl ZoneTool {
    pub fn active(&self) -> bool {
        self.mode.is_some()
    }

    fn toggle(&mut self, kind: ZoneKind) {
        self.mode = (self.mode != Some(kind)).then_some(kind);
    }
}

fn clear_zones(mut zones: ResMut<Zones>) {
    *zones = Zones::default();
}

pub fn new_stockpile_tiles(corner: UVec2) -> impl Iterator<Item = UVec2> {
    (0..NEW_STOCKPILE_SIZE)
        .flat_map(move |y| (0..NEW_STOCKPILE_SIZE).map(move |x| corner + UVec2::new(x, y)))
}

fn is_open_square(map: &WorldMap, corner: UVec2) -> bool {
    new_stockpile_tiles(corner)
        .all(|t| can_zone(map, t.as_ivec2()) && map.plant(t.as_ivec2()).is_none())
}

pub fn find_stockpile_site(map: &WorldMap, near: UVec2) -> Option<UVec2> {
    nearest_tile(map, near, STOCKPILE_SITE_SEARCH_TILES, |p| {
        is_open_square(map, p)
    })
    .or_else(|| {
        nearest_tile(map, near, STOCKPILE_SITE_SEARCH_TILES, |p| {
            can_zone(map, p.as_ivec2())
        })
    })
}

pub 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.toggle(ZoneKind::Stockpile),
            Action::ToggleFarmTool => tool.toggle(ZoneKind::Farm),
            Action::Cancel
            | Action::CycleBuildTool
            | Action::ToggleDesignateTool
            | Action::ToggleDemolishTool => tool.mode = None,
            _ => continue,
        }
        tool.drag_start = None;
        return;
    }
    let Some(kind) = tool.mode else {
        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_zone(kind, start.as_ivec2(), end.as_ivec2());
        } else {
            zones.add_zone(kind, &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, &HashSet::new(), 8);
        assert_eq!(plan.len(), 8);
        assert!(plan.iter().all(|t| !zones.is_stockpile(*t)));
        let beyond_one_ring = plan_expansion(&map, &zones, &HashSet::new(), 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, &HashSet::new(), 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(), &HashSet::new(), 8).is_empty());
    }

    #[test]
    fn farms_need_fertile_ground_and_never_overlap_stockpiles() {
        let mut map = WorldMap::filled(UVec2::new(4, 1), Terrain::Soil);
        map.set(UVec2::new(1, 0), Terrain::Stone);
        let mut zones = Zones::default();
        zones.add_stockpile(&map, IVec2::new(3, 0), IVec2::new(3, 0));
        zones.add_zone(ZoneKind::Farm, &map, IVec2::ZERO, IVec2::new(3, 0));
        let mut farm: Vec<_> = zones.farm_tiles().collect();
        farm.sort_by_key(|t| t.x);
        assert_eq!(farm, vec![UVec2::new(0, 0), UVec2::new(2, 0)]);
        zones.add_stockpile(&map, IVec2::ZERO, IVec2::ZERO);
        assert!(!zones.is_stockpile(UVec2::ZERO));
    }

    #[test]
    fn storage_rules_apply_to_one_connected_stockpile() {
        let map = WorldMap::filled(UVec2::new(10, 4), Terrain::Grass);
        let mut zones = Zones::default();
        zones.add_stockpile(&map, IVec2::ZERO, IVec2::new(2, 1));
        zones.add_stockpile(&map, IVec2::new(6, 0), IVec2::new(8, 1));
        zones.update_region_rule(UVec2::ZERO, |rule| {
            rule.toggle(Stock::Stone);
            rule.priority = StoragePriority::High;
        });
        assert!(!zones.accepts(UVec2::new(2, 1), ItemKind::Stone));
        assert!(zones.accepts(UVec2::new(2, 1), ItemKind::Wood));
        assert_eq!(zones.rule(UVec2::new(1, 0)).priority, StoragePriority::High);
        assert!(zones.accepts(UVec2::new(7, 0), ItemKind::Stone));
        assert_eq!(
            zones.rule(UVec2::new(7, 0)).priority,
            StoragePriority::Normal
        );
    }

    #[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 {
                mode: Some(ZoneKind::Stockpile),
                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>().mode = None;
        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_without_stockpiles() {
        let map = WorldMap::filled(UVec2::splat(8), Terrain::Grass);
        let mut app = App::new();
        app.init_resource::<Zones>()
            .add_message::<WorldGenerated>()
            .add_systems(Update, clear_zones.run_if(on_message::<WorldGenerated>));
        app.world_mut()
            .resource_mut::<Zones>()
            .add_stockpile(&map, IVec2::ZERO, IVec2::ONE);
        app.world_mut().write_message(WorldGenerated);
        app.update();
        assert!(!app.world().resource::<Zones>().has_stockpile());
    }

    #[test]
    fn new_stockpiles_go_on_the_nearest_fully_open_ground() {
        let mut map = WorldMap::filled(UVec2::splat(12), Terrain::Grass);
        for y in 0..12 {
            map.set_plant(UVec2::new(4, y), Some(crate::world::Plant::Tree));
            map.set(UVec2::new(7, y), Terrain::DeepWater);
        }
        let site = find_stockpile_site(&map, UVec2::new(5, 5)).unwrap();
        assert!(new_stockpile_tiles(site).all(|t| t.x < 4));
    }
}