use crate::ecs::world::Entity;
use std::collections::HashMap;

pub struct EntityMapping {
    entity_map: HashMap<u64, Entity>,
    reverse_entity_map: HashMap<Entity, u64>,
    next_entity_id: u64,
}

impl Default for EntityMapping {
    fn default() -> Self {
        Self::new()
    }
}

impl EntityMapping {
    pub fn new() -> Self {
        Self {
            entity_map: HashMap::new(),
            reverse_entity_map: HashMap::new(),
            next_entity_id: 1,
        }
    }

    pub fn register(&mut self, entity: Entity) -> u64 {
        if let Some(&existing_id) = self.reverse_entity_map.get(&entity) {
            return existing_id;
        }
        let plugin_entity_id = self.generate_id();
        self.entity_map.insert(plugin_entity_id, entity);
        self.reverse_entity_map.insert(entity, plugin_entity_id);
        plugin_entity_id
    }

    pub fn get(&self, plugin_entity_id: u64) -> Option<Entity> {
        self.entity_map.get(&plugin_entity_id).copied()
    }

    pub fn unregister(&mut self, plugin_entity_id: u64) {
        if let Some(entity) = self.entity_map.remove(&plugin_entity_id) {
            self.reverse_entity_map.remove(&entity);
        }
    }

    pub fn cleanup_invalid<F>(&mut self, is_valid: F)
    where
        F: Fn(Entity) -> bool,
    {
        let stale_ids: Vec<u64> = self
            .entity_map
            .iter()
            .filter(|(_, entity)| !is_valid(**entity))
            .map(|(id, _)| *id)
            .collect();

        for id in stale_ids {
            self.unregister(id);
        }
    }

    fn generate_id(&mut self) -> u64 {
        let id = self.next_entity_id;
        self.next_entity_id += 1;
        id
    }
}
