use std::collections::HashMap;

use bevy::prelude::*;

use crate::{
    building::Blueprints,
    items::ItemKind,
    sim::{SimSystems, SimTick, TICKS_PER_HOUR},
};

pub struct EconomyPlugin;

impl Plugin for EconomyPlugin {
    fn build(&self, app: &mut App) {
        app.init_resource::<ColonyGoals>()
            .add_systems(FixedUpdate, learn_needs.in_set(SimSystems::Decide));
    }
}

const HOURS_OF_MEMORY: f32 = 24.0;
const MAX_GOAL: f32 = 300.0;

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum Stock {
    Wood,
    Stone,
    Iron,
    Gold,
    RawFood,
    Meals,
}

impl Stock {
    pub const ALL: [Stock; 6] = [
        Stock::Wood,
        Stock::Stone,
        Stock::Iron,
        Stock::Gold,
        Stock::RawFood,
        Stock::Meals,
    ];

    pub fn of(kind: ItemKind) -> Stock {
        match kind {
            ItemKind::Wood => Stock::Wood,
            ItemKind::Stone => Stock::Stone,
            ItemKind::Iron => Stock::Iron,
            ItemKind::Gold => Stock::Gold,
            ItemKind::Berries | ItemKind::Meat | ItemKind::Vegetables => Stock::RawFood,
            ItemKind::Meal => Stock::Meals,
        }
    }

    pub fn members(self) -> impl Iterator<Item = ItemKind> {
        ItemKind::ALL
            .into_iter()
            .filter(move |&k| Stock::of(k) == self)
    }

    pub fn is_food(self) -> bool {
        matches!(self, Stock::RawFood | Stock::Meals)
    }

    pub fn label(self) -> &'static str {
        match self {
            Stock::Wood => "Wood",
            Stock::Stone => "Stone",
            Stock::Iron => "Iron",
            Stock::Gold => "Gold",
            Stock::RawFood => "Raw food",
            Stock::Meals => "Meals",
        }
    }

    fn buffer_days(self) -> f32 {
        if self.is_food() { 1.5 } else { 2.0 }
    }

    fn minimum(self) -> f32 {
        match self {
            Stock::Wood => 15.0,
            Stock::Stone => 5.0,
            Stock::Iron | Stock::Gold => 0.0,
            Stock::RawFood => 10.0,
            Stock::Meals => 3.0,
        }
    }

    fn initial_use_per_day(self) -> f32 {
        match self {
            Stock::Wood => 20.0,
            Stock::Stone => 10.0,
            Stock::Iron | Stock::Gold => 0.0,
            Stock::RawFood => 30.0,
            Stock::Meals => 3.0,
        }
    }
}

#[derive(Resource, Debug, Clone)]
pub struct ColonyGoals {
    goals: HashMap<Stock, u32>,
    use_per_day: HashMap<Stock, f32>,
    used_this_hour: HashMap<Stock, u32>,
}

impl Default for ColonyGoals {
    fn default() -> Self {
        let mut goals = Self {
            goals: HashMap::new(),
            use_per_day: Stock::ALL
                .iter()
                .map(|&s| (s, s.initial_use_per_day()))
                .collect(),
            used_this_hour: HashMap::new(),
        };
        goals.recompute(&HashMap::new());
        goals
    }
}

impl ColonyGoals {
    #[cfg(test)]
    pub fn fixed(goals: &[(ItemKind, u32)]) -> Self {
        Self {
            goals: goals.iter().map(|&(k, n)| (Stock::of(k), n)).collect(),
            use_per_day: HashMap::new(),
            used_this_hour: HashMap::new(),
        }
    }

    pub fn goal(&self, stock: Stock) -> u32 {
        self.goals.get(&stock).copied().unwrap_or(0)
    }

    pub fn use_per_day(&self, stock: Stock) -> f32 {
        self.use_per_day.get(&stock).copied().unwrap_or(0.0)
    }

    pub fn record_use(&mut self, kind: ItemKind, count: u32) {
        *self.used_this_hour.entry(Stock::of(kind)).or_insert(0) += count;
    }

    pub fn wanted(
        &self,
        kind: ItemKind,
        stocks: &HashMap<ItemKind, u32>,
        on_the_way: &HashMap<ItemKind, u32>,
    ) -> bool {
        let stock = Stock::of(kind);
        let have: u32 = stock
            .members()
            .map(|k| {
                stocks.get(&k).copied().unwrap_or(0) + on_the_way.get(&k).copied().unwrap_or(0)
            })
            .sum();
        have < self.goal(stock)
    }

    pub fn learn_hour(&mut self, planned: &HashMap<Stock, u32>) {
        for stock in Stock::ALL {
            let used = self.used_this_hour.remove(&stock).unwrap_or(0) as f32;
            let rate = self.use_per_day.entry(stock).or_insert(0.0);
            *rate += (used * 24.0 - *rate) / HOURS_OF_MEMORY;
        }
        self.recompute(planned);
    }

    fn recompute(&mut self, planned: &HashMap<Stock, u32>) {
        for stock in Stock::ALL {
            let learned = (self.use_per_day(stock) * stock.buffer_days())
                .clamp(stock.minimum(), MAX_GOAL)
                .ceil() as u32;
            let goal = learned + planned.get(&stock).copied().unwrap_or(0);
            self.goals.insert(stock, goal);
        }
    }
}

pub fn planned_needs(blueprints: &Blueprints) -> HashMap<Stock, u32> {
    let mut needs = HashMap::new();
    for kind in blueprints.0.values() {
        let (item, count) = kind.cost();
        *needs.entry(Stock::of(item)).or_insert(0) += count;
    }
    needs
}

fn learn_needs(tick: Res<SimTick>, blueprints: Res<Blueprints>, mut goals: ResMut<ColonyGoals>) {
    if tick.0 == 0 || !tick.0.is_multiple_of(TICKS_PER_HOUR) {
        return;
    }
    goals.learn_hour(&planned_needs(&blueprints));
}

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

    fn learn_days(goals: &mut ColonyGoals, kind: ItemKind, per_hour: u32, days: u32) {
        for _ in 0..days * 24 {
            goals.record_use(kind, per_hour);
            goals.learn_hour(&HashMap::new());
        }
    }

    #[test]
    fn targets_follow_what_the_colony_actually_uses() {
        let mut goals = ColonyGoals::default();
        learn_days(&mut goals, ItemKind::Meal, 1, 5);
        let meals = goals.goal(Stock::Meals);
        assert!((32..=40).contains(&meals), "{meals}");
        learn_days(&mut goals, ItemKind::Meal, 0, 5);
        assert_eq!(goals.goal(Stock::Meals), 3);
    }

    #[test]
    fn unused_materials_are_not_stockpiled() {
        let mut goals = ColonyGoals::default();
        learn_days(&mut goals, ItemKind::Wood, 0, 10);
        assert_eq!(goals.goal(Stock::Wood), 15);
        assert_eq!(goals.goal(Stock::Iron), 0);
    }

    // Raw foods stand in for each other (cooking accepts any), so the colony targets them as one
    // group and is free to forage, hunt or farm whichever is at hand.
    #[test]
    fn raw_foods_share_one_target() {
        let goals = ColonyGoals::fixed(&[(ItemKind::Berries, 20)]);
        let stocks = HashMap::from([(ItemKind::Meat, 15), (ItemKind::Vegetables, 4)]);
        assert!(goals.wanted(ItemKind::Berries, &stocks, &HashMap::new()));
        let on_the_way = HashMap::from([(ItemKind::Berries, 1)]);
        assert!(!goals.wanted(ItemKind::Meat, &stocks, &on_the_way));
    }

    #[test]
    fn planned_construction_raises_material_targets() {
        let mut blueprints = Blueprints::default();
        for x in 0..10 {
            blueprints.0.insert(
                UVec2::new(x, 0),
                BuildKind::Structure(crate::structures::StructureKind::Wall),
            );
        }
        let mut goals = ColonyGoals::default();
        let before = goals.goal(Stock::Stone);
        goals.learn_hour(&planned_needs(&blueprints));
        assert!(goals.goal(Stock::Stone) >= before + 19);
    }
}