use std::ops::Deref;

use crate::util::{Id, IdNum, IdTracker};

pub struct Arena<T, I> {
    data: Vec<T>,
    tracker: IdTracker<I>,
}

impl<T, I: const IdNum> Arena<T, I> {
    pub fn new() -> Self {
        Self {
            data: Vec::new(),
            tracker: IdTracker::default(),
        }
    }

    pub fn push(&mut self, value: T) -> Id<I> {
        let id = self.tracker.next();
        let i = id.idx();
        if i == self.data.len() {
            self.data.push(value);
        } else {
            self.data[i] = value;
        }
        id
    }

    pub fn remove(&mut self, id: Id<I>) -> T
    where
        T: Copy,
    {
        let i = id.idx();
        self.tracker.free(id);
        self.data[i]
    }
}

impl<T, I: const IdNum> Default for Arena<T, I> {
    fn default() -> Self {
        Self::new()
    }
}

pub struct TrackedArena<T, I> {
    inner: Arena<T, I>,
    refs: Vec<u32>,
    pub changed: bool,
}

impl<T, I: const IdNum> TrackedArena<T, I> {
    pub fn new() -> Self {
        Self {
            inner: Arena::default(),
            refs: Vec::new(),
            changed: true,
        }
    }

    pub fn push(&mut self, value: T) -> Id<I> {
        self.changed = true;
        let id = self.inner.push(value);
        let i = id.idx();
        if i == self.refs.len() {
            self.refs.push(0);
        }
        id
    }

    pub fn push_ref(&mut self, i: Id<I>) {
        self.refs[i.idx()] += 1;
    }

    pub fn remove(&mut self, id: Id<I>) -> T
    where
        T: Copy,
    {
        let i = id.idx();
        self.refs[i] -= 1;
        if self.refs[i] == 0 {
            self.changed = true;
            self.inner.remove(id)
        } else {
            self[i]
        }
    }
}

impl<T, I: const IdNum> Default for TrackedArena<T, I> {
    fn default() -> Self {
        Self::new()
    }
}

impl<T, I> Deref for TrackedArena<T, I> {
    type Target = Vec<T>;

    fn deref(&self) -> &Self::Target {
        &self.inner.data
    }
}

impl<T, I> Deref for Arena<T, I> {
    type Target = Vec<T>;

    fn deref(&self) -> &Self::Target {
        &self.data
    }
}