use std::collections::HashMap;

use bevy::prelude::*;
use serde::{Deserialize, Serialize};

use crate::{
    colonist::{Colonist, Path},
    family::Partner,
    jobs::CurrentJob,
    mood::Mood,
    needs::is_sleeping,
    population::ColonyLog,
    rng::SimRng,
    sim::{SimSystems, SimTick, TICKS_PER_HOUR},
    world::TilePos,
    worldgen::WorldGenerated,
};

pub struct SocialPlugin;

impl Plugin for SocialPlugin {
    fn build(&self, app: &mut App) {
        app.init_resource::<Relations>()
            .add_systems(Update, clear_relations.run_if(on_message::<WorldGenerated>))
            .add_systems(
                FixedUpdate,
                (mourn, socialize).chain().in_set(SimSystems::Act),
            );
    }
}

pub const SOCIAL_RANGE: i32 = 4;
pub const FRIEND: f32 = 0.5;
pub const RIVAL: f32 = -0.5;
pub const FRIEND_LOST: f32 = 0.3;
pub const RIVAL_LOST: f32 = -0.3;
const CHAT_GAIN: f32 = 0.1;
const ARGUMENT_LOSS: f32 = 0.12;
const CHAT_CHANCE: f32 = 0.35;
const ARGUMENT_CHANCE: f32 = 0.12;
const GRUMPY_MOOD: f32 = 0.3;
const FADE_PER_HOUR: f32 = 0.998;
pub const MOURNING_TICKS: u64 = 2 * 24 * TICKS_PER_HOUR;

#[derive(Component, Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub struct Personality {
    pub warmth: f32,
}

impl Default for Personality {
    fn default() -> Self {
        Self { warmth: 0.5 }
    }
}

#[derive(Component, Debug, Clone, Copy, PartialEq)]
pub struct Mourning {
    pub until: u64,
}

#[derive(Resource, Default, Debug, Clone)]
pub struct Relations(
    pub HashMap<(Entity, Entity), f32>,
    pub HashMap<(Entity, Entity), Bond>,
);

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Bond {
    Friends,
    Rivals,
}

pub fn bond_after(bond: Option<Bond>, opinion: f32) -> Option<Bond> {
    match bond {
        Some(Bond::Friends) if opinion >= FRIEND_LOST => Some(Bond::Friends),
        Some(Bond::Rivals) if opinion <= RIVAL_LOST => Some(Bond::Rivals),
        _ if opinion >= FRIEND => Some(Bond::Friends),
        _ if opinion <= RIVAL => Some(Bond::Rivals),
        _ => None,
    }
}

fn key(a: Entity, b: Entity) -> (Entity, Entity) {
    (a.min(b), a.max(b))
}

impl Relations {
    pub fn opinion(&self, a: Entity, b: Entity) -> f32 {
        self.0.get(&key(a, b)).copied().unwrap_or(0.0)
    }

    pub fn change(&mut self, a: Entity, b: Entity, delta: f32) -> (f32, f32) {
        let value = self.0.entry(key(a, b)).or_insert(0.0);
        let before = *value;
        *value = (*value + delta).clamp(-1.0, 1.0);
        (before, *value)
    }

    pub fn set(&mut self, a: Entity, b: Entity, value: f32) {
        self.0.insert(key(a, b), value.clamp(-1.0, 1.0));
    }

    pub fn of(&self, who: Entity) -> impl Iterator<Item = (Entity, f32)> + '_ {
        self.0.iter().filter_map(move |(&(a, b), &v)| {
            (a == who)
                .then_some((b, v))
                .or((b == who).then_some((a, v)))
        })
    }

    pub fn bond(&self, a: Entity, b: Entity) -> Option<Bond> {
        self.1.get(&key(a, b)).copied()
    }

    pub fn bonded(&self, who: Entity, bond: Bond) -> impl Iterator<Item = Entity> + '_ {
        self.1.iter().filter_map(move |(&(a, b), &kind)| {
            (kind == bond && (a == who || b == who)).then_some(if a == who { b } else { a })
        })
    }

    pub fn friends(&self, who: Entity) -> usize {
        self.bonded(who, Bond::Friends).count()
    }

    pub fn rivals(&self, who: Entity) -> usize {
        self.bonded(who, Bond::Rivals).count()
    }

    pub fn update_bond(&mut self, a: Entity, b: Entity) -> Option<Bond> {
        let k = key(a, b);
        let before = self.1.get(&k).copied();
        let after = bond_after(before, self.opinion(a, b));
        match after {
            Some(bond) => self.1.insert(k, bond),
            None => self.1.remove(&k),
        };
        (after != before).then_some(after).flatten()
    }

    pub fn forget(&mut self, who: Entity) {
        self.0.retain(|&(a, b), _| a != who && b != who);
        self.1.retain(|&(a, b), _| a != who && b != who);
    }
}

fn clear_relations(mut relations: ResMut<Relations>) {
    relations.0.clear();
}

pub fn chat_chance(warmth: f32) -> f32 {
    CHAT_CHANCE * (0.5 + warmth)
}

pub fn argument_chance(warmth: f32, grumpy: bool) -> f32 {
    ARGUMENT_CHANCE * (1.5 - warmth) * if grumpy { 2.0 } else { 1.0 }
}

#[allow(clippy::type_complexity)]
fn socialize(
    tick: Res<SimTick>,
    mut rng: ResMut<SimRng>,
    mut relations: ResMut<Relations>,
    mut log: ResMut<ColonyLog>,
    colonists: Query<(
        Entity,
        &Colonist,
        &TilePos,
        &CurrentJob,
        &Path,
        Option<&Personality>,
        Option<&Mood>,
    )>,
) {
    if !tick.0.is_multiple_of(TICKS_PER_HOUR) {
        return;
    }
    for value in relations.0.values_mut() {
        *value *= FADE_PER_HOUR;
    }
    let pairs: Vec<(Entity, Entity)> = relations.1.keys().copied().collect();
    for (a, b) in pairs {
        relations.update_bond(a, b);
    }
    let awake: Vec<(Entity, &str, UVec2, f32, bool)> = colonists
        .iter()
        .filter(|(_, _, tile, job, path, ..)| !is_sleeping(job, tile, path))
        .map(|(e, c, tile, _, _, personality, mood)| {
            let warmth = personality.copied().unwrap_or_default().warmth;
            let grumpy = mood.is_some_and(|m| m.0 < GRUMPY_MOOD);
            (e, c.name.as_str(), tile.0, warmth, grumpy)
        })
        .collect();
    for (i, a) in awake.iter().enumerate() {
        for b in &awake[i + 1..] {
            let near = (a.2.as_ivec2() - b.2.as_ivec2()).abs().max_element() <= SOCIAL_RANGE;
            if !near {
                continue;
            }
            let warmth = (a.3 + b.3) / 2.0;
            let roll = rng.range_f32(0.0..1.0);
            let chat = chat_chance(warmth);
            let delta = if roll < chat {
                CHAT_GAIN
            } else if roll < chat + argument_chance(warmth, a.4 || b.4) {
                -ARGUMENT_LOSS
            } else {
                continue;
            };
            relations.change(a.0, b.0, delta);
            let event = match relations.update_bond(a.0, b.0) {
                Some(Bond::Friends) => "became friends",
                Some(Bond::Rivals) => "became rivals",
                None => continue,
            };
            log.push(*tick, format!("{} and {} {event}", a.1, b.1));
        }
    }
}

fn mourn(
    mut commands: Commands,
    tick: Res<SimTick>,
    mut relations: ResMut<Relations>,
    mut removed: RemovedComponents<Colonist>,
    colonists: Query<(Entity, Option<&Partner>), With<Colonist>>,
    mourning: Query<(Entity, &Mourning)>,
) {
    for (entity, grief) in &mourning {
        if tick.0 >= grief.until {
            commands.entity(entity).remove::<Mourning>();
        }
    }
    for gone in removed.read() {
        for (entity, partner) in &colonists {
            let close = relations.bond(entity, gone) == Some(Bond::Friends)
                || partner.is_some_and(|p| p.0 == gone);
            if close {
                commands.entity(entity).insert(Mourning {
                    until: tick.0 + MOURNING_TICKS,
                });
            }
        }
        relations.forget(gone);
    }
}

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

    fn e(i: u32) -> Entity {
        Entity::from_raw_u32(i).unwrap()
    }

    #[test]
    fn opinions_are_shared_and_cross_into_friendship_or_rivalry() {
        let mut relations = Relations::default();
        relations.change(e(1), e(2), 0.3);
        assert_eq!(relations.opinion(e(2), e(1)), 0.3);
        relations.set(e(1), e(3), FRIEND);
        relations.set(e(1), e(4), RIVAL);
        assert_eq!(relations.update_bond(e(1), e(3)), Some(Bond::Friends));
        assert_eq!(relations.update_bond(e(1), e(4)), Some(Bond::Rivals));
        assert_eq!((relations.friends(e(1)), relations.rivals(e(1))), (1, 1));
        relations.forget(e(1));
        assert!(relations.0.is_empty());
    }

    // Warm people chat more and argue less; unhappy people pick fights.
    #[test]
    fn warmth_and_mood_steer_conversations() {
        assert!(chat_chance(0.9) > chat_chance(0.1));
        assert!(argument_chance(0.1, false) > argument_chance(0.9, false));
        assert!(argument_chance(0.5, true) > argument_chance(0.5, false));
    }

    // Colonists who spend their days side by side grow close (or come to loathe each other);
    // the colony log notes when a friendship forms.
    #[test]
    fn colonists_near_each_other_form_opinions() {
        let mut app = App::new();
        app.init_resource::<SimTick>()
            .insert_resource(SimRng::new(5))
            .init_resource::<Relations>()
            .init_resource::<ColonyLog>()
            .add_systems(Update, socialize);
        let spawn = |app: &mut App, name: &str, x: u32| {
            app.world_mut()
                .spawn((
                    Colonist { name: name.into() },
                    TilePos(UVec2::new(x, 0)),
                    Personality { warmth: 1.0 },
                ))
                .id()
        };
        let ada = spawn(&mut app, "Ada", 0);
        let bram = spawn(&mut app, "Bram", 2);
        let far = spawn(&mut app, "Cleo", 30);
        for hour in 0..100 {
            app.world_mut().resource_mut::<SimTick>().0 = hour * TICKS_PER_HOUR;
            app.update();
        }
        let relations = app.world().resource::<Relations>();
        assert!(relations.opinion(ada, bram) >= FRIEND);
        assert_eq!(relations.opinion(ada, far), 0.0);
        let log = app.world().resource::<ColonyLog>();
        assert!(
            log.0
                .iter()
                .any(|l| l.ends_with("Ada and Bram became friends"))
        );
    }

    // Losing a friend or partner hurts more than distant grief: they mourn for two days.
    #[test]
    fn friends_mourn_the_dead() {
        let mut app = App::new();
        app.init_resource::<SimTick>()
            .init_resource::<Relations>()
            .add_systems(Update, mourn);
        let spawn =
            |app: &mut App, name: &str| app.world_mut().spawn(Colonist { name: name.into() }).id();
        let (ada, bram, cleo) = (
            spawn(&mut app, "Ada"),
            spawn(&mut app, "Bram"),
            spawn(&mut app, "Cleo"),
        );
        let mut relations = app.world_mut().resource_mut::<Relations>();
        relations.set(ada, bram, 0.8);
        relations.update_bond(ada, bram);
        app.update();
        app.world_mut().despawn(bram);
        app.update();
        assert!(app.world().get::<Mourning>(ada).is_some());
        assert!(app.world().get::<Mourning>(cleo).is_none());
        assert!(app.world().resource::<Relations>().1.is_empty());
    }

    // Friendships and rivalries need a clear swing to end, so an opinion hovering at the line
    // doesn't flip them back and forth (and flood the log).
    #[test]
    fn bonds_have_hysteresis() {
        assert_eq!(bond_after(None, 0.49), None);
        assert_eq!(bond_after(None, 0.5), Some(Bond::Friends));
        assert_eq!(bond_after(Some(Bond::Friends), 0.35), Some(Bond::Friends));
        assert_eq!(bond_after(Some(Bond::Friends), 0.2), None);
        assert_eq!(bond_after(Some(Bond::Rivals), -0.35), Some(Bond::Rivals));
    }
}