use std::collections::HashMap;

use bevy::{
    app::{App, PreUpdate},
    ecs::{
        entity::Entity,
        query::With,
        resource::Resource,
        system::{Query, ResMut},
    },
    math::{Dir3, Vec3},
    transform::components::Transform,
};

use crate::bunnies::bunny::Bunny;

pub fn add_locator(app: &mut App) -> &mut App {
    app.init_resource::<BunnyLocator>()
        .add_systems(PreUpdate, update_locator)
}

fn update_locator(
    mut locator: ResMut<BunnyLocator>,
    bunnies: Query<(Entity, &Transform), With<Bunny>>,
) {
    bunnies
        .iter()
        .for_each(|(bunny, transform)| locator.set(bunny, transform.translation, transform.back()));
}

#[derive(Resource, Default)]
pub struct BunnyLocator {
    bunnies: HashMap<Entity, (Vec3, Dir3)>,
}

impl BunnyLocator {
    fn set(&mut self, bunny: Entity, position: Vec3, direction: Dir3) {
        self.bunnies.insert(bunny, (position, direction));
    }

    pub fn get_nearby_count(&self, center: Vec3, radius: f32) -> usize {
        self.bunnies
            .iter()
            .filter(|(_, (pos, _))| center.distance(*pos) <= radius)
            .count()
    }

    pub(super) fn get_nearby_except_self(
        &self,
        this: Entity,
        center: Vec3,
        radius: f32,
    ) -> Vec<(Entity, (Vec3, Dir3))> {
        self.bunnies
            .iter()
            .filter(|(e, (pos, _))| **e != this && center.distance(*pos) <= radius)
            .map(|(b, (p, d))| (*b, (*p, *d)))
            .collect::<Vec<_>>()
    }
}