use crate::{
    layout::{
        Axis, Layers, Len, Modules, Size, TextAttrs, TextBuffer, TextData, TextTexture,
        TextureHandle, Textures, UiRegion, UiVec2, Vec2, WidgetId, Widgets,
    },
    render::{Mask, MaskIdx, Primitive, PrimitiveHandle, PrimitiveInst},
    util::{HashMap, HashSet, Id, TrackedArena},
};

/// makes your surfaces look pretty
pub struct Painter<'a, 'c> {
    ctx: &'a mut PainterCtx<'c>,
    region: UiRegion,
    mask: MaskIdx,
    textures: Vec<TextureHandle>,
    primitives: Vec<PrimitiveHandle>,
    children: Vec<Id>,
    children_width: HashMap<Id, (UiVec2, Len)>,
    children_height: HashMap<Id, (UiVec2, Len)>,
    /// whether this widget depends on region's final pixel size or not
    /// TODO: decide if point (pt) should be used here instead of px
    pub layer: usize,
    id: Id,
}

/// context for a painter; lets you draw and redraw widgets
struct PainterCtx<'a> {
    pub widgets: &'a Widgets,
    pub active: &'a mut HashMap<Id, WidgetInstance>,
    pub layers: &'a mut Layers,
    pub textures: &'a mut Textures,
    pub masks: &'a mut TrackedArena<Mask, u32>,
    pub text: &'a mut TextData,
    pub output_size: Vec2,
    pub modules: &'a mut Modules,
    pub cache_width: HashMap<Id, (UiVec2, Len)>,
    pub cache_height: HashMap<Id, (UiVec2, Len)>,
    pub needs_redraw: HashSet<Id>,
    draw_started: HashSet<Id>,
}

/// stores information for children about the highest level parent that needed their size
/// so that they can redraw the parent if their size changes
#[derive(Clone, Copy, Debug, Default)]
pub struct ResizeRef {
    x: Option<(Id, (UiVec2, Len))>,
    y: Option<(Id, (UiVec2, Len))>,
}

/// important non rendering data for retained drawing
#[derive(Debug)]
pub struct WidgetInstance {
    pub id: Id,
    pub region: UiRegion,
    pub parent: Option<Id>,
    pub textures: Vec<TextureHandle>,
    pub primitives: Vec<PrimitiveHandle>,
    pub children: Vec<Id>,
    pub resize: ResizeRef,
    pub mask: MaskIdx,
    pub layer: usize,
}

/// data to be stored in Ui to create PainterCtxs easily
#[derive(Default)]
pub struct PainterData {
    pub widgets: Widgets,
    pub active: HashMap<Id, WidgetInstance>,
    pub layers: Layers,
    pub textures: Textures,
    pub text: TextData,
    pub output_size: Vec2,
    pub modules: Modules,
    pub px_dependent: HashSet<Id>,
    pub masks: TrackedArena<Mask, u32>,
}

impl<'a> PainterCtx<'a> {
    /// redraws a widget that's currently active (drawn)
    /// can be called on something already drawn or removed,
    /// will just return if so
    pub fn redraw(&mut self, id: Id) {
        self.needs_redraw.remove(&id);
        if self.draw_started.contains(&id) {
            return;
        }
        let Some(active) = self.active.get(&id) else {
            return;
        };
        let mut resize = active.resize;

        // set resize back after redrawing
        let finish = |s: &mut Self, resize| {
            if let Some(active) = s.active.get_mut(&id) {
                // might need to get_or_insert here instead of just assuming
                active.resize = resize;
            }
        };

        // check if a parent depends on the desired size of this, if so then redraw it first
        // TODO: this is stupid having 2 of these, don't ask me what the consequences are
        let mut ret = false;
        if let Some((rid, (outer, old_desired))) = &mut resize.x {
            let new_desired = SizeCtx {
                source: id,
                cache_width: &mut self.cache_width,
                cache_height: &mut self.cache_height,
                text: self.text,
                textures: self.textures,
                widgets: self.widgets,
                outer: *outer,
                output_size: self.output_size,
                checked_width: &mut Default::default(),
                checked_height: &mut Default::default(),
                id,
            }
            .width_inner(id);
            if new_desired != *old_desired {
                // unsure if I need to walk down the tree here
                self.redraw(*rid);
                *old_desired = new_desired;
                if self.draw_started.contains(&id) {
                    ret = true;
                }
            }
        }

        if let Some((rid, (outer, old_desired))) = &mut resize.y {
            // NOTE: might need hack in Span here (or also do it properly here)
            let new_desired = SizeCtx {
                source: id,
                cache_width: &mut self.cache_width,
                cache_height: &mut self.cache_height,
                text: self.text,
                textures: self.textures,
                widgets: self.widgets,
                outer: *outer,
                output_size: self.output_size,
                checked_width: &mut Default::default(),
                checked_height: &mut Default::default(),
                id,
            }
            .height_inner(id);
            if new_desired != *old_desired {
                self.redraw(*rid);
                *old_desired = new_desired;
                if self.draw_started.contains(&id) {
                    ret = true;
                }
            }
        }

        if ret {
            return finish(self, resize);
        }

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

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

    fn draw_inner(
        &mut self,
        layer: usize,
        id: Id,
        region: UiRegion,
        parent: Option<Id>,
        mask: MaskIdx,
        old_children: Option<Vec<Id>>,
    ) {
        // I have no idea if these checks work lol
        // the idea is u can't redraw stuff u already drew,
        // and if parent is different then there's another copy with a different parent
        // but this has a very weird issue where you can't move widgets unless u remove first
        // so swapping is impossible rn I think?
        // there's definitely better solutions like a counter (>1 = panic) but don't care rn
        // if self.draw_started.contains(&id) {
        //     panic!(
        //         "Cannot draw the same widget ({}) twice (1)",
        //         self.widgets.data(&id).unwrap().label
        //     );
        // }
        let mut old_children = old_children.unwrap_or_default();
        let mut resize = ResizeRef::default();
        if let Some(active) = self.active.get_mut(&id)
            && !self.needs_redraw.contains(&id)
        {
            // check to see if we can skip drawing first
            if active.parent != parent {
                panic!("Cannot draw the same widget twice (2)");
            }
            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).unwrap();
            old_children = active.children;
            resize = active.resize;
        }

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

        let mut painter = Painter {
            region,
            mask,
            layer,
            id,
            textures: Vec::new(),
            primitives: Vec::new(),
            ctx: self,
            children: Vec::new(),
            children_width: Default::default(),
            children_height: Default::default(),
        };

        painter.ctx.widgets.get_dyn_dynamic(id).draw(&mut painter);

        let children_width = painter.children_width;
        let children_height = painter.children_height;

        // add to active
        let instance = WidgetInstance {
            id,
            region,
            parent,
            textures: painter.textures,
            primitives: painter.primitives,
            children: painter.children,
            resize,
            mask: painter.mask,
            layer,
        };
        // set resize for children who's size this widget depends on
        for (cid, outer) in children_width {
            if let Some(w) = self.active.get_mut(&cid)
                && w.resize.x.is_none()
            {
                w.resize.x = Some((id, outer))
            }
        }
        for (cid, outer) in children_height {
            if let Some(w) = self.active.get_mut(&cid)
                && w.resize.y.is_none()
            {
                w.resize.y = Some((id, outer))
            }
        }

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

        // update modules
        for m in self.modules.iter_mut() {
            m.on_draw(&instance);
        }
        self.active.insert(id, instance);
    }

    fn mov(&mut self, id: Id, from: UiRegion, to: UiRegion) {
        let active = self.active.get_mut(&id).unwrap();
        for h in &active.primitives {
            let region = self.layers[h.layer].primitives.region_mut(h);
            *region = region.outside(&from).within(&to);
        }
        active.region = active.region.outside(&from).within(&to);
        for m in self.modules.iter_mut() {
            m.on_move(active);
        }
        // children will not be changed, so this technically should not be needed
        // probably need unsafe
        let children = active.children.clone();
        for child in children {
            self.mov(child, from, to);
        }
    }

    /// NOTE: instance textures are cleared and self.textures freed
    fn remove(&mut self, id: Id) -> Option<WidgetInstance> {
        let mut inst = self.active.remove(&id);
        if let Some(inst) = &mut inst {
            for h in &inst.primitives {
                let mask = self.layers.free(h);
                if mask != MaskIdx::NONE {
                    self.masks.remove(mask);
                }
            }
            inst.textures.clear();
            self.textures.free();
            for m in self.modules.iter_mut() {
                m.on_undraw(inst);
            }
        }
        inst
    }

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

impl PainterData {
    fn ctx(&mut self, needs_redraw: HashSet<Id>) -> PainterCtx<'_> {
        PainterCtx {
            widgets: &self.widgets,
            active: &mut self.active,
            layers: &mut self.layers,
            textures: &mut self.textures,
            text: &mut self.text,
            output_size: self.output_size,
            modules: &mut self.modules,
            masks: &mut self.masks,
            cache_width: Default::default(),
            cache_height: Default::default(),
            draw_started: Default::default(),
            needs_redraw,
        }
    }

    pub fn draw(&mut self, id: Id) {
        let mut ctx = self.ctx(Default::default());
        ctx.draw_started.clear();
        ctx.layers.clear();
        ctx.draw_inner(0, id, UiRegion::FULL, None, MaskIdx::NONE, None);
    }

    pub fn redraw(&mut self, ids: HashSet<Id>) {
        let mut ctx = self.ctx(ids);
        while let Some(&id) = ctx.needs_redraw.iter().next() {
            ctx.redraw(id);
        }
    }
}

impl<'a, 'c> Painter<'a, 'c> {
    fn primitive_at<P: Primitive>(&mut self, primitive: P, region: UiRegion) {
        let h = self.ctx.layers.write(
            self.layer,
            PrimitiveInst {
                id: self.id,
                primitive,
                region,
                mask_idx: self.mask,
            },
        );
        if self.mask != MaskIdx::NONE {
            // TODO: I have no clue if this works at all :joy:
            self.ctx.masks.push_ref(self.mask);
        }
        self.primitives.push(h);
    }

    /// Writes a primitive to be rendered
    pub fn primitive<P: Primitive>(&mut self, primitive: P) {
        self.primitive_at(primitive, self.region)
    }

    pub fn primitive_within<P: Primitive>(&mut self, primitive: P, region: UiRegion) {
        self.primitive_at(primitive, region.within(&self.region));
    }

    pub fn set_mask(&mut self, region: UiRegion) {
        assert!(self.mask == MaskIdx::NONE);
        self.mask = self.ctx.masks.push(Mask { region });
    }

    /// Draws a widget within this widget's region.
    pub fn widget<W>(&mut self, id: &WidgetId<W>) {
        self.widget_at(id, self.region);
    }

    /// Draws a widget somewhere within this one.
    /// Useful for drawing child widgets in select areas.
    pub fn widget_within<W>(&mut self, id: &WidgetId<W>, region: UiRegion) {
        self.widget_at(id, region.within(&self.region));
    }

    fn widget_at<W>(&mut self, id: &WidgetId<W>, region: UiRegion) {
        self.children.push(id.id);
        self.ctx
            .draw_inner(self.layer, id.id, region, Some(self.id), self.mask, None);
    }

    pub fn texture_within(&mut self, handle: &TextureHandle, region: UiRegion) {
        self.textures.push(handle.clone());
        self.primitive_at(handle.primitive(), region.within(&self.region));
    }

    pub fn texture(&mut self, handle: &TextureHandle) {
        self.textures.push(handle.clone());
        self.primitive(handle.primitive());
    }

    pub fn texture_at(&mut self, handle: &TextureHandle, region: UiRegion) {
        self.textures.push(handle.clone());
        self.primitive_at(handle.primitive(), region);
    }

    /// returns (handle, offset from top left)
    pub fn render_text(&mut self, buffer: &mut TextBuffer, attrs: &TextAttrs) -> TextTexture {
        self.ctx.text.draw(buffer, attrs, self.ctx.textures)
    }

    pub fn region(&self) -> UiRegion {
        self.region
    }

    pub fn size<W>(&mut self, id: &WidgetId<W>) -> Size {
        self.size_ctx().size(id)
    }

    pub fn len_axis<W>(&mut self, id: &WidgetId<W>, axis: Axis) -> Len {
        match axis {
            Axis::X => self.size_ctx().width(id),
            Axis::Y => self.size_ctx().height(id),
        }
    }

    pub fn size_ctx(&mut self) -> SizeCtx<'_> {
        SizeCtx {
            text: self.ctx.text,
            textures: self.ctx.textures,
            widgets: self.ctx.widgets,
            output_size: self.ctx.output_size,
            checked_width: &mut self.children_width,
            checked_height: &mut self.children_height,
            cache_width: &mut self.ctx.cache_width,
            cache_height: &mut self.ctx.cache_height,
            source: self.id,
            id: self.id,
            outer: self.region.size(),
        }
    }

    pub fn output_size(&self) -> Vec2 {
        self.ctx.output_size
    }

    pub fn px_size(&mut self) -> Vec2 {
        self.region.size().to_abs(self.ctx.output_size)
    }

    pub fn text_data(&mut self) -> &mut TextData {
        self.ctx.text
    }

    pub fn child_layer(&mut self) {
        self.layer = self.ctx.layers.child(self.layer);
    }

    pub fn next_layer(&mut self) {
        self.layer = self.ctx.layers.next(self.layer);
    }

    pub fn label(&self) -> &str {
        &self.ctx.widgets.data(&self.id).unwrap().label
    }

    pub fn id(&self) -> &Id {
        &self.id
    }
}

pub struct SizeCtx<'a> {
    pub text: &'a mut TextData,
    pub textures: &'a mut Textures,
    source: Id,
    widgets: &'a Widgets,
    cache_width: &'a mut HashMap<Id, (UiVec2, Len)>,
    cache_height: &'a mut HashMap<Id, (UiVec2, Len)>,
    checked_width: &'a mut HashMap<Id, (UiVec2, Len)>,
    checked_height: &'a mut HashMap<Id, (UiVec2, Len)>,
    /// TODO: should this be pub? rn used for sized
    pub outer: UiVec2,
    output_size: Vec2,
    id: Id,
}

impl SizeCtx<'_> {
    pub fn id(&self) -> &Id {
        &self.id
    }

    pub fn source(&self) -> &Id {
        &self.source
    }

    fn width_inner(&mut self, id: Id) -> Len {
        // first check cache
        // TODO: is this needed? broken rn bc does not store children during upper size check,
        // so if something actually using check_* hits cache it fails to add them
        // if let Some(&(outer, len)) = self.cache_width.get(&id)
        //     && outer == self.outer
        // {
        //     self.checked_width.insert(id, (self.outer, len));
        //     return len;
        // }
        // store self vars that need to be maintained
        let self_outer = self.outer;
        let self_id = self.id;
        // get size of input id
        self.id = id;
        let len = self.widgets.get_dyn_dynamic(id).desired_width(self);
        // restore vars & update cache + checked
        self.outer = self_outer;
        self.id = self_id;
        self.cache_width.insert(id, (self.outer, len));
        self.checked_width.insert(id, (self.outer, len));
        len
    }

    // TODO: should be refactored to share code w width_inner
    fn height_inner(&mut self, id: Id) -> Len {
        // if let Some(&(outer, len)) = self.cache_height.get(&id)
        //     && outer == self.outer
        // {
        //     prntln!("FAIL {id:?}");
        //     self.checked_height.insert(id, (self.outer, len));
        //     return len;
        // }
        let self_outer = self.outer;
        let self_id = self.id;
        self.id = id;
        let len = self.widgets.get_dyn_dynamic(id).desired_height(self);
        self.outer = self_outer;
        self.id = self_id;
        self.cache_height.insert(id, (self.outer, len));
        self.checked_height.insert(id, (self.outer, len));
        len
    }

    pub fn width<W>(&mut self, id: &WidgetId<W>) -> Len {
        self.width_inner(id.id)
    }

    pub fn height<W>(&mut self, id: &WidgetId<W>) -> Len {
        self.height_inner(id.id)
    }

    pub fn len_axis<W>(&mut self, id: &WidgetId<W>, axis: Axis) -> Len {
        match axis {
            Axis::X => self.width(id),
            Axis::Y => self.height(id),
        }
    }

    pub fn size<W>(&mut self, id: &WidgetId<W>) -> Size {
        Size {
            x: self.width(id),
            y: self.height(id),
        }
    }

    pub fn px_size(&mut self) -> Vec2 {
        self.outer.to_abs(self.output_size)
    }

    pub fn output_size(&mut self) -> Vec2 {
        self.output_size
    }

    pub fn draw_text(&mut self, buffer: &mut TextBuffer, attrs: &TextAttrs) -> TextTexture {
        self.text.draw(buffer, attrs, self.textures)
    }

    pub fn label(&self, id: &Id) -> &String {
        self.widgets.label(id)
    }
}