use iris_core::{
    WidgetId,
    util::{HashMap, HashSet},
};
use std::{
    any::{Any, TypeId},
    marker::PhantomData,
};

#[derive(Clone, Copy, PartialEq, Eq, Hash)]
struct Key {
    id: WidgetId,
    ty: TypeId,
    i: usize,
}

#[derive(Default)]
pub struct WidgetState {
    widgets: HashMap<WidgetId, HashSet<(TypeId, usize)>>,
    counts: HashMap<(WidgetId, TypeId), usize>,
    map: HashMap<Key, Box<dyn Any>>,
}

impl WidgetState {
    pub fn new() -> Self {
        Self::default()
    }
    pub fn add<T: 'static>(&mut self, id: WidgetId, data: T) -> WeakState<T> {
        let ty = TypeId::of::<T>();
        let count = self.counts.entry((id, ty)).or_default();
        let i = *count;
        let key = Key { ty, i, id };
        self.map.insert(key, Box::new(data));
        self.widgets.entry(id).or_default().insert((ty, i));
        *count += 1;
        WeakState {
            key,
            _pd: PhantomData,
        }
    }
    pub fn remove(&mut self, id: WidgetId) {
        for &(ty, i) in self.widgets.get(&id).into_iter().flatten() {
            self.map.remove(&Key { id, ty, i });
        }
    }
    pub fn get<T: 'static>(&self, state: WeakState<T>) -> &T {
        self.map.get(&state.key).unwrap().downcast_ref().unwrap()
    }
    pub fn get_mut<T: 'static>(&mut self, state: WeakState<T>) -> &mut T {
        self.map
            .get_mut(&state.key)
            .unwrap()
            .downcast_mut()
            .unwrap()
    }
}

#[derive(Clone, Copy)]
pub struct WeakState<T> {
    key: Key,
    _pd: PhantomData<T>,
}

pub trait HasWidgetState {
    fn widget_state(&self) -> &WidgetState;
    fn widget_state_mut(&mut self) -> &mut WidgetState;
}

impl<'a, T: 'static> FnOnce<(&'a mut WidgetState,)> for WeakState<T> {
    type Output = &'a mut T;

    extern "rust-call" fn call_once(self, (state,): (&'a mut WidgetState,)) -> Self::Output {
        state.get_mut(self)
    }
}