use std::collections::HashMap;

pub struct GenerationalRegistry<T> {
    pub entries: Vec<Option<T>>,
    pub generations: Vec<u32>,
    pub reference_counts: Vec<usize>,
    pub free_indices: Vec<u32>,
    pub name_to_index: HashMap<String, u32>,
    pub index_to_name: Vec<Option<String>>,
}

impl<T> Default for GenerationalRegistry<T> {
    fn default() -> Self {
        Self {
            entries: Vec::new(),
            generations: Vec::new(),
            reference_counts: Vec::new(),
            free_indices: Vec::new(),
            name_to_index: HashMap::new(),
            index_to_name: Vec::new(),
        }
    }
}

impl<T> GenerationalRegistry<T> {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn insert(&mut self, name: String, value: T) -> (u32, u32) {
        if let Some(&existing_index) = self.name_to_index.get(&name) {
            let generation = self.generations[existing_index as usize];
            self.entries[existing_index as usize] = Some(value);
            return (existing_index, generation);
        }

        let index = if let Some(free_index) = self.free_indices.pop() {
            self.generations[free_index as usize] += 1;
            self.entries[free_index as usize] = Some(value);
            self.reference_counts[free_index as usize] = 0;
            self.index_to_name[free_index as usize] = Some(name.clone());
            free_index
        } else {
            let index = self.entries.len() as u32;
            self.entries.push(Some(value));
            self.generations.push(0);
            self.reference_counts.push(0);
            self.index_to_name.push(Some(name.clone()));
            index
        };

        self.name_to_index.insert(name, index);
        (index, self.generations[index as usize])
    }

    pub fn remove(&mut self, index: u32, generation: u32) -> Option<T> {
        if !self.is_valid(index, generation) {
            return None;
        }

        let slot = index as usize;
        let value = self.entries[slot].take();

        if let Some(name) = self.index_to_name[slot].take() {
            self.name_to_index.remove(&name);
        }

        self.reference_counts[slot] = 0;
        self.free_indices.push(index);

        value
    }

    pub fn is_valid(&self, index: u32, generation: u32) -> bool {
        let slot = index as usize;
        slot < self.generations.len()
            && self.generations[slot] == generation
            && self.entries[slot].is_some()
    }

    pub fn lookup_index(&self, name: &str) -> Option<(u32, u32)> {
        self.name_to_index.get(name).map(|&index| {
            let generation = self.generations[index as usize];
            (index, generation)
        })
    }

    pub fn add_reference(&mut self, index: u32) {
        if let Some(count) = self.reference_counts.get_mut(index as usize) {
            *count += 1;
        }
    }

    pub fn remove_reference(&mut self, index: u32) {
        if let Some(count) = self.reference_counts.get_mut(index as usize) {
            *count = count.saturating_sub(1);
        }
    }

    pub fn reference_count(&self, index: u32) -> usize {
        self.reference_counts
            .get(index as usize)
            .copied()
            .unwrap_or(0)
    }

    pub fn remove_unused(&mut self) -> Vec<String> {
        let mut removed = Vec::new();

        for index in 0..self.entries.len() {
            if self.reference_counts[index] == 0 && self.entries[index].is_some() {
                if let Some(name) = self.index_to_name[index].take() {
                    self.name_to_index.remove(&name);
                    removed.push(name);
                }
                self.entries[index] = None;
                self.free_indices.push(index as u32);
            }
        }

        removed
    }

    pub fn len(&self) -> usize {
        self.entries.iter().filter(|entry| entry.is_some()).count()
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub fn clear(&mut self) {
        self.entries.clear();
        self.generations.clear();
        self.reference_counts.clear();
        self.free_indices.clear();
        self.name_to_index.clear();
        self.index_to_name.clear();
    }
}

pub fn registry_entry<T>(
    registry: &GenerationalRegistry<T>,
    index: u32,
    generation: u32,
) -> Option<&T> {
    if registry.is_valid(index, generation) {
        registry.entries[index as usize].as_ref()
    } else {
        None
    }
}

pub fn registry_entry_mut<T>(
    registry: &mut GenerationalRegistry<T>,
    index: u32,
    generation: u32,
) -> Option<&mut T> {
    if registry.is_valid(index, generation) {
        registry.entries[index as usize].as_mut()
    } else {
        None
    }
}

pub fn registry_entry_by_name<'a, T>(
    registry: &'a GenerationalRegistry<T>,
    name: &str,
) -> Option<&'a T> {
    registry
        .name_to_index
        .get(name)
        .and_then(|&index| registry.entries[index as usize].as_ref())
}

pub fn registry_entry_by_name_mut<'a, T>(
    registry: &'a mut GenerationalRegistry<T>,
    name: &str,
) -> Option<&'a mut T> {
    registry
        .name_to_index
        .get(name)
        .copied()
        .and_then(|index| registry.entries[index as usize].as_mut())
}
