use crate::{
    ActiveData, Axis, EventsLike, Painter, SizeCtx, Ui, UiRegion, UiVec2, WidgetId,
    render::MaskIdx,
    util::{HashSet, forget_ref},
};
use std::ops::{Deref, DerefMut};

/// state maintained between widgets during painting
pub struct DrawState<'a> {
    pub(super) ui: &'a mut Ui,
    pub(super) events: &'a mut dyn EventsLike,
    draw_started: HashSet<WidgetId>,
}

impl<'a> DrawState<'a> {
    pub fn new(ui: &'a mut Ui, events: &'a mut dyn EventsLike) -> Self {
        Self {
            ui,
            events,
            draw_started: Default::default(),
        }
    }

    pub fn redraw_updates(&mut self) {
        while let Some(&id) = self.widgets.needs_redraw.iter().next() {
            self.redraw(id);
        }
        self.ui.free(self.events);
    }

    /// redraws a widget that's currently active (drawn)
    pub fn redraw(&mut self, id: WidgetId) {
        self.widgets.needs_redraw.remove(&id);
        self.draw_started.remove(&id);
        // check if parent depends on the desired size of this, if so then redraw it first
        for axis in [Axis::X, Axis::Y] {
            if let Some(&(outer, old)) = self.cache.size.axis_dyn(axis).get(&id)
                && let Some(current) = self.active.get(&id)
                && let Some(pid) = current.parent
            {
                self.cache.size.axis_dyn(axis).remove(&id);
                let new = self.size_ctx(id, outer).len_axis(id, axis);
                self.cache.size.axis_dyn(axis).insert(id, (outer, new));
                if new != old {
                    self.redraw(pid);
                }
            }
        }

        if self.draw_started.contains(&id) {
            return;
        }

        let Some(active) = self.remove(id, false) else {
            return;
        };

        self.draw_inner(
            active.layer,
            id,
            active.region,
            active.parent,
            active.mask,
            Some(active.children),
        );
    }

    pub(super) fn size_ctx<'b>(&'b mut self, source: WidgetId, outer: UiVec2) -> SizeCtx<'b> {
        SizeCtx {
            source,
            cache: &mut self.ui.cache,
            text: &mut self.ui.text,
            textures: &mut self.ui.textures,
            widgets: &self.ui.widgets,
            outer,
            output_size: self.ui.output_size,
            id: source,
        }
    }

    pub fn redraw_all(&mut self) {
        // free all resources & cache
        for (id, active) in self.ui.active.drain() {
            let data = self.ui.widgets.data(id).unwrap();
            self.events.undraw(data, &active);
        }
        self.ui.cache.clear();
        self.ui.free(self.events);
        self.layers.clear();
        self.widgets.needs_redraw.clear();

        if let Some(id) = &self.ui.root {
            self.draw_inner(0, id.id(), UiRegion::FULL, None, MaskIdx::NONE, None);
        }
    }

    pub(super) fn draw_inner(
        &mut self,
        layer: usize,
        id: WidgetId,
        region: UiRegion,
        parent: Option<WidgetId>,
        mask: MaskIdx,
        old_children: Option<Vec<WidgetId>>,
    ) {
        let mut old_children = old_children.unwrap_or_default();
        if let Some(active) = self.ui.active.get_mut(&id)
            && !self.ui.widgets.needs_redraw.contains(&id)
        {
            // check to see if we can skip drawing first
            if active.region == region {
                return;
            } else if active.region.size() == region.size() {
                // TODO: epsilon?
                let from = active.region;
                self.mov(id, from, region);
                return;
            }
            // if not, then maintain resize and track old children to remove unneeded
            let active = self.remove(id, false).unwrap();
            old_children = active.children;
        }

        // draw widget
        self.draw_started.insert(id);

        let mut painter = Painter {
            state: self,
            region,
            mask,
            layer,
            id,
            textures: Vec::new(),
            primitives: Vec::new(),
            children: Vec::new(),
        };

        let mut widget = painter.state.widgets.get_dyn_dynamic(id);
        widget.draw(&mut painter);
        drop(widget);

        let Painter {
            state: _,
            region,
            mask,
            textures,
            primitives,
            children,
            layer,
            id,
        } = painter;

        // add to active
        let active = ActiveData {
            id,
            region,
            parent,
            textures,
            primitives,
            children,
            mask,
            layer,
        };

        // remove old children that weren't kept
        for c in &old_children {
            if !active.children.contains(c) {
                self.remove_rec(*c);
            }
        }

        // update modules
        let data = self.ui.widgets.data(id).unwrap();
        self.events.draw(data, &active);
        self.active.insert(id, active);
    }

    fn mov(&mut self, id: WidgetId, from: UiRegion, to: UiRegion) {
        let active = self.ui.active.get_mut(&id).unwrap();
        for h in &active.primitives {
            let region = self.ui.layers[h.layer].region_mut(h);
            *region = region.outside(&from).within(&to);
        }
        active.region = active.region.outside(&from).within(&to);
        // SAFETY: children cannot be recursive
        let children = unsafe { forget_ref(&active.children) };
        for child in children {
            self.mov(*child, from, to);
        }
    }

    /// NOTE: instance textures are cleared and self.textures freed
    fn remove(&mut self, id: WidgetId, undraw: bool) -> Option<ActiveData> {
        let mut active = self.active.remove(&id);
        if let Some(active) = &mut active {
            for h in &active.primitives {
                let mask = self.layers.free(h);
                if mask != MaskIdx::NONE {
                    self.masks.remove(mask);
                }
            }
            active.textures.clear();
            self.textures.free();
            if undraw {
                let data = self.ui.widgets.data(id).unwrap();
                self.events.undraw(data, active);
            }
        }
        active
    }

    fn remove_rec(&mut self, id: WidgetId) -> Option<ActiveData> {
        self.cache.remove(id);
        let inst = self.remove(id, true);
        if let Some(inst) = &inst {
            for c in &inst.children {
                self.remove_rec(*c);
            }
        }
        inst
    }
}

impl Deref for DrawState<'_> {
    type Target = Ui;

    fn deref(&self) -> &Self::Target {
        self.ui
    }
}
impl DerefMut for DrawState<'_> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.ui
    }
}