use crate::{Ui, UiModule, Widget, WidgetId, WidgetRef, util::HashMap};
use std::{
    hash::Hash,
    ops::{Index, IndexMut},
    rc::Rc,
};

pub trait Event: Sized {
    type Module<Ctx: 'static>: EventModule<Self, Ctx>;
    type Data: Clone;
}

pub struct EventCtx<'a, Ctx, Data> {
    pub ui: &'a mut Ui,
    pub state: &'a mut Ctx,
    pub data: Data,
}

pub struct EventIdCtx<'a, Ctx, Data, W: ?Sized> {
    pub widget: WidgetRef<W>,
    pub ui: &'a mut Ui,
    pub state: &'a mut Ctx,
    pub data: Data,
}

impl<'a, Ctx, Data, W2, W: Widget> Index<WidgetRef<W>> for EventIdCtx<'a, Ctx, Data, W2> {
    type Output = W;

    fn index(&self, index: WidgetRef<W>) -> &Self::Output {
        &self.ui[index]
    }
}

impl<'a, Ctx, Data, W2, W: Widget> IndexMut<WidgetRef<W>> for EventIdCtx<'a, Ctx, Data, W2> {
    fn index_mut(&mut self, index: WidgetRef<W>) -> &mut Self::Output {
        &mut self.ui[index]
    }
}

impl<'a, Ctx, Data, W: Widget> EventIdCtx<'a, Ctx, Data, W> {
    pub fn widget(&mut self) -> &mut W {
        &mut self.ui[self.widget]
    }
}

pub trait EventFn<Ctx, Data>: Fn(EventCtx<Ctx, Data>) + 'static {}
impl<F: Fn(EventCtx<Ctx, Data>) + 'static, Ctx, Data> EventFn<Ctx, Data> for F {}

pub trait WidgetEventFn<Ctx, Data, W: ?Sized>: Fn(EventIdCtx<Ctx, Data, W>) + 'static {}
impl<F: Fn(EventIdCtx<Ctx, Data, W>) + 'static, Ctx, Data, W: ?Sized> WidgetEventFn<Ctx, Data, W>
    for F
{
}

pub trait DefaultEvent: Hash + Eq + 'static {
    type Data: Clone = ();
}

impl<E: DefaultEvent> Event for E {
    type Module<Ctx: 'static> = DefaultEventModule<E, Ctx>;
    type Data = E::Data;
}

pub trait EventModule<E: Event, Ctx>: UiModule + Default {
    fn register(&mut self, id: WidgetId, event: E, f: impl EventFn<Ctx, E::Data>);
    fn run<'a>(
        &self,
        id: WidgetId,
        event: E,
    ) -> Option<impl Fn(EventCtx<Ctx, E::Data>) + use<'a, Self, E, Ctx>>;
}

type EventFnMap<Ctx, Data> = HashMap<WidgetId, Vec<Rc<dyn EventFn<Ctx, Data>>>>;
pub struct DefaultEventModule<E: Event, Ctx> {
    map: HashMap<E, EventFnMap<Ctx, <E as Event>::Data>>,
}

impl<E: Event + 'static, Ctx: 'static> UiModule for DefaultEventModule<E, Ctx> {
    fn on_remove(&mut self, id: WidgetId) {
        for map in self.map.values_mut() {
            map.remove(&id);
        }
    }
}

pub trait HashableEvent: Event + Hash + Eq + 'static {}
impl<E: Event + Hash + Eq + 'static> HashableEvent for E {}

impl<E: HashableEvent, Ctx: 'static> EventModule<E, Ctx> for DefaultEventModule<E, Ctx> {
    fn register(&mut self, id: WidgetId, event: E, f: impl EventFn<Ctx, <E as Event>::Data>) {
        self.map
            .entry(event)
            .or_default()
            .entry(id)
            .or_default()
            .push(Rc::new(f));
    }

    fn run<'a>(
        &self,
        id: WidgetId,
        event: E,
    ) -> Option<impl Fn(EventCtx<Ctx, E::Data>) + use<'a, E, Ctx>> {
        if let Some(map) = self.map.get(&event)
            && let Some(fs) = map.get(&id)
        {
            let fs = fs.clone();
            Some(move |ctx: EventCtx<Ctx, <E as Event>::Data>| {
                for f in &fs {
                    f(EventCtx {
                        ui: ctx.ui,
                        state: ctx.state,
                        data: ctx.data.clone(),
                    })
                }
            })
        } else {
            None
        }
    }
}

impl<E: HashableEvent, Ctx: 'static> DefaultEventModule<E, Ctx> {
    pub fn run_all(&self, event: E, ctx: EventCtx<Ctx, E::Data>)
    where
        E::Data: Clone,
    {
        if let Some(map) = self.map.get(&event) {
            for fs in map.values() {
                for f in fs {
                    f(EventCtx {
                        ui: ctx.ui,
                        state: ctx.state,
                        data: ctx.data.clone(),
                    })
                }
            }
        }
    }
}

impl<E: Event + 'static, Ctx: 'static> Default for DefaultEventModule<E, Ctx> {
    fn default() -> Self {
        Self {
            map: Default::default(),
        }
    }
}

impl Ui {
    pub fn run_event<E: Event, Ctx: 'static, W>(
        &mut self,
        ctx: &mut Ctx,
        id: WidgetRef<W>,
        event: E,
        data: E::Data,
    ) {
        if let Some(f) = self
            .data
            .modules
            .get_mut::<E::Module<Ctx>>()
            .run(id.id(), event)
        {
            f(EventCtx {
                ui: self,
                state: ctx,
                data,
            });
        }
    }
}