use std::{hash::Hash, rc::Rc};

use crate::{
    layout::{IdFnTag, Ui, UiModule, Widget, WidgetId, WidgetIdFn, WidgetLike},
    util::{HashMap, Id},
};

pub trait UiCtx {
    fn ui(&mut self) -> &mut Ui;
}

impl UiCtx for Ui {
    fn ui(&mut self) -> &mut Ui {
        self
    }
}

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

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

pub trait Eventable<W, Tag> {
    fn on<E: Event, Ctx: 'static>(
        self,
        event: E,
        f: impl EventFn<Ctx, E::Data>,
    ) -> impl WidgetIdFn<W> + Eventable<W, IdFnTag>;

    fn id_on<E: Event, Ctx: 'static>(
        self,
        event: E,
        f: impl Fn(&WidgetId<W>, &mut Ctx, E::Data) + 'static,
    ) -> impl WidgetIdFn<W> + Eventable<W, IdFnTag>
    where
        W: Widget;

    fn edit_on<E: Event>(
        self,
        event: E,
        f: impl Fn(&mut W, E::Data) + 'static,
    ) -> impl WidgetIdFn<W> + Eventable<W, IdFnTag>
    where
        W: Widget;
}

impl<W: WidgetLike<Tag>, Tag> Eventable<W::Widget, Tag> for W {
    fn on<E: Event, Ctx: 'static>(
        self,
        event: E,
        f: impl EventFn<Ctx, E::Data>,
    ) -> impl WidgetIdFn<W::Widget> {
        move |ui| {
            let id = self.add(ui);
            ui.register_event(&id, event, f);
            id
        }
    }

    fn id_on<E: Event, Ctx: 'static>(
        self,
        event: E,
        f: impl Fn(&WidgetId<W::Widget>, &mut Ctx, E::Data) + 'static,
    ) -> impl WidgetIdFn<W::Widget>
    where
        W::Widget: Widget,
    {
        self.with_id(move |ui, id| {
            let id2 = id.clone();
            id.on(event, move |ctx, pos| f(&id2, ctx, pos)).add(ui)
        })
    }

    fn edit_on<E: Event>(
        self,
        event: E,
        f: impl Fn(&mut W::Widget, E::Data) + 'static,
    ) -> impl WidgetIdFn<W::Widget>
    where
        W::Widget: Widget,
    {
        self.id_on(event, move |id, ui: &mut Ui, pos| f(&mut ui[id], pos))
    }
}

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: Id, event: E, f: impl EventFn<Ctx, E::Data>);
    fn run<'a>(
        &self,
        id: &Id,
        event: E,
    ) -> Option<impl Fn(&mut Ctx, E::Data) + use<'a, Self, E, Ctx>>;
}

type EventFnMap<Ctx, Data> = HashMap<Id, 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: &Id) {
        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: Id, 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: &Id, event: E) -> Option<impl Fn(&mut 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: &mut Ctx, data: E::Data| {
                for f in &fs {
                    f(ctx, data.clone())
                }
            })
        } else {
            None
        }
    }
}

impl<E: HashableEvent, Ctx: 'static> DefaultEventModule<E, Ctx> {
    pub fn run_all(&self, ctx: &mut Ctx, event: E, data: E::Data)
    where
        E::Data: Clone,
    {
        if let Some(map) = self.map.get(&event) {
            for fs in map.values() {
                for f in fs {
                    f(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: UiCtx + 'static, W>(
        ctx: &mut Ctx,
        id: &WidgetId<W>,
        event: E,
        data: E::Data,
    ) {
        if let Some(f) = ctx
            .ui()
            .data
            .modules
            .get_mut::<E::Module<Ctx>>()
            .run(&id.id, event)
        {
            f(ctx, data);
        }
    }
}

pub trait EventCtx: UiCtx {
    fn run_event<E: Event + Clone, W>(&mut self, id: &WidgetId<W>, event: E, data: E::Data);
}

impl<Ctx: UiCtx + 'static> EventCtx for Ctx {
    fn run_event<E: Event + Clone, W>(&mut self, id: &WidgetId<W>, event: E, data: E::Data) {
        Ui::run_event(self, id, event.clone(), data.clone());
    }
}