use bevy::prelude::*;

use crate::{
    building::Blueprints,
    colonist::Colonist,
    items::Item,
    pathfinding::nearest_tile,
    plants::fertility,
    sim::{SimSystems, SimTick},
    structures::Structure,
    world::{Terrain, TilePos, WorldMap},
    zones::{Zones, can_farm},
};

pub struct FarmingPlugin;

impl Plugin for FarmingPlugin {
    fn build(&self, app: &mut App) {
        app.add_systems(FixedUpdate, plan_farm.in_set(SimSystems::Decide));
    }
}

pub const FARM_SIZE: u32 = 5;
pub const FARM_CLEARANCE: i32 = 2;
const PLAN_INTERVAL_TICKS: u64 = 600;
const FARM_SEARCH_TILES: usize = 3000;
const SOIL_WEIGHT: f32 = 4.0;

pub struct FarmContext<'a> {
    pub map: &'a WorldMap,
    pub zones: &'a Zones,
    pub blueprints: &'a Blueprints,
    pub occupied: &'a [UVec2],
}

impl FarmContext<'_> {
    fn square(corner: UVec2) -> impl Iterator<Item = UVec2> {
        (0..FARM_SIZE).flat_map(move |y| (0..FARM_SIZE).map(move |x| corner + UVec2::new(x, y)))
    }

    fn clear(&self, tile: UVec2) -> bool {
        let near = |other: UVec2| {
            (other.as_ivec2() - tile.as_ivec2()).abs().max_element() <= FARM_CLEARANCE
        };
        can_farm(self.map, tile.as_ivec2())
            && self.map.plant(tile.as_ivec2()).is_none()
            && !self.zones.stockpile_tiles().any(near)
            && !self.blueprints.0.keys().copied().any(near)
            && !self.occupied.iter().copied().any(near)
    }

    fn score(&self, corner: UVec2, home: UVec2) -> Option<f32> {
        let mut soil = 0.0;
        for tile in Self::square(corner) {
            if !self.clear(tile) {
                return None;
            }
            soil += fertility(self.map.get(tile.as_ivec2()).unwrap_or(Terrain::Sand));
        }
        Some(soil * SOIL_WEIGHT - corner.as_vec2().distance(home.as_vec2()))
    }

    pub fn plan(&self, home: UVec2) -> Option<Vec<UVec2>> {
        let mut best: Option<(f32, UVec2)> = None;
        nearest_tile(self.map, home, FARM_SEARCH_TILES, |corner| {
            if let Some(score) = self.score(corner, home)
                && best.is_none_or(|(b, _)| score > b)
            {
                best = Some((score, corner));
            }
            false
        });
        best.map(|(_, corner)| Self::square(corner).collect())
    }
}

fn plan_farm(
    tick: Res<SimTick>,
    map: Res<WorldMap>,
    mut zones: ResMut<Zones>,
    blueprints: Res<Blueprints>,
    structures: Query<&TilePos, (With<Structure>, Without<Colonist>)>,
    items: Query<&TilePos, With<Item>>,
) {
    if !tick.0.is_multiple_of(PLAN_INTERVAL_TICKS) || zones.has_farm() || !zones.has_stockpile() {
        return;
    }
    let occupied: Vec<UVec2> = structures.iter().chain(items.iter()).map(|p| p.0).collect();
    let home = zones.home(map.size / 2);
    let ctx = FarmContext {
        map: &map,
        zones: &zones,
        blueprints: &blueprints,
        occupied: &occupied,
    };
    if let Some(tiles) = ctx.plan(home) {
        zones.add_farm_tiles(&map, tiles);
    }
}

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

    fn plan(map: &WorldMap, zones: &Zones, occupied: &[UVec2]) -> Vec<UVec2> {
        FarmContext {
            map,
            zones,
            blueprints: &Blueprints::default(),
            occupied,
        }
        .plan(UVec2::new(20, 20))
        .unwrap()
    }

    #[test]
    fn farms_go_on_the_most_fertile_ground_nearby() {
        let mut map = WorldMap::filled(UVec2::splat(40), Terrain::Sand);
        for y in 26..34 {
            for x in 10..18 {
                map.set(UVec2::new(x, y), Terrain::Soil);
            }
        }
        let farm = plan(&map, &Zones::default(), &[]);
        assert_eq!(farm.len(), (FARM_SIZE * FARM_SIZE) as usize);
        assert!(
            farm.iter()
                .all(|t| map.get(t.as_ivec2()) == Some(Terrain::Soil))
        );
    }

    // Beds, fires and the stockpile need room around them for walls and walkways, and a farm
    // packed against them would be dug into by every later shelter plan.
    #[test]
    fn farms_keep_their_distance_from_the_colony() {
        let map = WorldMap::filled(UVec2::splat(40), Terrain::Soil);
        let mut zones = Zones::default();
        zones.add_tiles(&map, [UVec2::new(20, 20)]);
        let bed = UVec2::new(24, 20);
        let farm = plan(&map, &zones, &[bed]);
        for tile in farm {
            for other in [UVec2::new(20, 20), bed] {
                assert!((tile.as_ivec2() - other.as_ivec2()).abs().max_element() > FARM_CLEARANCE);
            }
        }
    }
}