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

pub struct Painter<'a, 'c> {
    ctx: &'a mut PainterCtx<'c>,
    region: UiRegion,
    mask: MaskIdx,
    textures: Vec<TextureHandle>,
    primitives: Vec<PrimitiveHandle>,
    children: Vec<Id>,
    sized_children: HashMap<Id, Size>,
    /// 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,
}

pub 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 screen_size: Vec2,
    pub modules: &'a mut Modules,
    pub px_dependent: &'a mut HashSet<Id>,
    draw_started: HashSet<Id>,
}

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: Option<Id>,
    pub mask: MaskIdx,
    pub layer: usize,
    pub desired_size: Size,
}

#[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> {
    pub fn new(data: &'a mut PainterData) -> Self {
        Self {
            widgets: &data.widgets,
            active: &mut data.active,
            layers: &mut data.layers,
            textures: &mut data.textures,
            text: &mut data.text,
            screen_size: data.output_size,
            modules: &mut data.modules,
            px_dependent: &mut data.px_dependent,
            masks: &mut data.masks,
            draw_started: HashSet::default(),
        }
    }

    pub fn redraw(&mut self, id: Id) {
        if self.draw_started.contains(&id) {
            return;
        }
        let Some(active) = self.active.get(&id) else {
            return;
        };

        if let Some(rid) = active.resize {
            let desired = SizeCtx {
                checked: &mut Default::default(),
                text: self.text,
                textures: self.textures,
                widgets: self.widgets,
                size: UiVec2::FULL_SIZE,
                screen_size: self.screen_size,
                px_dependent: &mut Default::default(),
                id,
            }
            .size_inner(id, active.region.size());
            if active.desired_size != desired {
                self.redraw(rid);
                if self.draw_started.contains(&id) {
                    return;
                }
            }
        }

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

        self.draw_inner(
            active.layer,
            id,
            active.region,
            active.parent,
            active.mask,
            Some(active.children),
        );
        self.active.get_mut(&id).unwrap().resize = active.resize;
    }

    pub fn draw(&mut self, id: Id) {
        self.draw_started.clear();
        self.layers.clear();
        self.draw_inner(0, id, UiRegion::full(), None, MaskIdx::NONE, None);
    }

    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 = None;
        if let Some(active) = self.active.get_mut(&id) {
            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;
            }
            let active = self.remove(id).unwrap();
            old_children = active.children;
            resize = active.resize;
        }

        self.draw_started.insert(id);

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

        // draw widgets
        painter.ctx.widgets.get_dyn_dynamic(id).draw(&mut painter);
        let desired_size = SizeCtx {
            text: painter.ctx.text,
            textures: painter.ctx.textures,
            widgets: painter.ctx.widgets,
            checked: &mut Default::default(),
            screen_size: painter.ctx.screen_size,
            px_dependent: &mut Default::default(),
            id: painter.id,
            size: UiVec2::FULL_SIZE,
        }
        .size_raw(id);

        let sized_children = painter.sized_children;

        // add to active
        let instance = WidgetInstance {
            id,
            region,
            parent,
            textures: painter.textures,
            primitives: painter.primitives,
            children: painter.children,
            resize,
            mask: painter.mask,
            desired_size,
            layer,
        };
        for cid in sized_children.keys() {
            if let Some(w) = self.active.get_mut(cid)
                && w.resize.is_none()
            {
                w.resize = Some(id)
            }
        }
        for c in &old_children {
            if !instance.children.contains(c) {
                self.remove_rec(*c);
            }
        }

        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();
        // children will not be changed, so this technically should not be needed
        // probably need unsafe
        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);
        }
        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);
            }
        }
        self.px_dependent.remove(&id);
        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<'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 region_size<W>(&mut self, id: &WidgetId<W>) -> UiVec2 {
        self.size_ctx().size(id).to_uivec2(self.region.size().rel)
    }

    pub fn size_ctx(&mut self) -> SizeCtx<'_> {
        SizeCtx {
            text: self.ctx.text,
            textures: self.ctx.textures,
            widgets: self.ctx.widgets,
            checked: &mut self.sized_children,
            screen_size: self.ctx.screen_size,
            px_dependent: self.ctx.px_dependent,
            id: self.id,
            size: self.region.size(),
        }
    }

    pub fn px_size(&mut self) -> Vec2 {
        self.ctx.px_dependent.insert(self.id);
        self.region.size().to_abs(self.ctx.screen_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 struct SizeCtx<'a> {
    pub text: &'a mut TextData,
    pub textures: &'a mut Textures,
    widgets: &'a Widgets,
    px_dependent: &'a mut HashSet<Id>,
    checked: &'a mut HashMap<Id, Size>,
    /// TODO: should this be pub? rn used for sized
    pub size: UiVec2,
    screen_size: Vec2,
    id: Id,
}

impl SizeCtx<'_> {
    fn size_inner(&mut self, id: Id, size: UiVec2) -> Size {
        let self_size = self.size;
        self.size = size;
        let size = self.widgets.get_dyn_dynamic(id).desired_size(self);
        self.size = self_size;
        self.checked.insert(id, size);
        size
    }
    pub fn size<W>(&mut self, id: &WidgetId<W>) -> Size {
        if let Some(&size) = self.checked.get(&id.id) {
            return size;
        }
        self.size_inner(id.id, self.size)
    }
    fn size_raw(&mut self, id: Id) -> Size {
        self.size_inner(id, self.size)
    }
    pub fn px_size(&mut self) -> Vec2 {
        self.px_dependent.insert(self.id);
        self.size.to_abs(self.screen_size)
    }
    pub fn draw_text(&mut self, buffer: &mut TextBuffer, attrs: &TextAttrs) -> TextTexture {
        self.text.draw(buffer, attrs, self.textures)
    }
    pub fn label(&self) -> &String {
        self.widgets.label(&self.id)
    }
}