monolith

use monolith_planner::Rect;

use slotmap::{SlotMap, new_key_type};

use std::collections::HashMap;

use std::hash::Hash;

new_key_type! { pub struct NodeId; }

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]

#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]

#[cfg_attr(feature = "serde", serde(rename_all = "kebab-case"))]

pub enum Axis {

    Horizontal,

    Vertical,

}

impl Axis {

    pub fn other(self) -> Self {

        match self {

            Self::Horizontal => Self::Vertical,

            Self::Vertical => Self::Horizontal,

        }

    }

    fn index(self) -> usize {

        match self {

            Self::Horizontal => 0,

            Self::Vertical => 1,

        }

    }

}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]

#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]

#[cfg_attr(feature = "serde", serde(rename_all = "kebab-case"))]

pub enum Dir {

    Left,

    Right,

    Up,

    Down,

}

impl Dir {

    pub fn axis(self) -> Axis {

        match self {

            Self::Left | Self::Right => Axis::Horizontal,

            Self::Up | Self::Down => Axis::Vertical,

        }

    }

    pub fn forward(self) -> bool {

        matches!(self, Self::Right | Self::Down)

    }

}

const MIN_SHARE: f64 = 0.05;

/// A window's width and height relative to its neighbours'.

type Size = [f64; 2];

#[derive(Clone, Debug)]

enum Kind<W> {

    Leaf(W, Size),

    Container {

        axis: Axis,

        children: Vec<NodeId>,

        expanded: Option<NodeId>,

    },

}

#[derive(Clone, Debug)]

struct Node<W> {

    parent: Option<NodeId>,

    kind: Kind<W>,

}

/// One output's tiling tree. Every tiled window of every tag lives here, the

/// caller's visibility filter decides which ones take space.

///

/// Children run on the axis opposite their parent. Containers hold at least

/// two children, hidden ones included, and a container left with one child

/// is replaced by it.

#[derive(Clone, Debug)]

pub struct Tree<W> {

    nodes: SlotMap<NodeId, Node<W>>,

    root: Option<NodeId>,

    leaves: HashMap<W, NodeId>,

    stack: Axis,

}

impl<W: Copy + Eq + Hash> Default for Tree<W> {

    fn default() -> Self {

        Self {

            nodes: SlotMap::with_key(),

            root: None,

            leaves: HashMap::new(),

            stack: Axis::Vertical,

        }

    }

}

impl<W: Copy + Eq + Hash> Tree<W> {

    /// The axis tag rows stack along.

    pub fn stack_axis(&self) -> Axis {

        self.stack

    }

    pub fn contains(&self, w: W) -> bool {

        self.leaves.contains_key(&w)

    }

    pub fn is_empty(&self) -> bool {

        self.root.is_none()

    }

    pub fn windows(&self) -> impl Iterator<Item = W> + '_ {

        self.leaves.keys().copied()

    }

    /// Inserts `w` as the next sibling of `anchor`, or as a new row if

    /// `anchor` is not in the tree.

    pub fn insert_after(&mut self, w: W, anchor: W, vis: impl Fn(W) -> bool) {

        let Some(&a) = self.leaves.get(&anchor) else {

            return self.insert_row(w, vis);

        };

        match self.position(a) {

            Some((p, i)) if Some(p) == self.root && self.axis(p) == Some(self.stack) => {

                let n = self.new_leaf(w, self.size(a));

                self.unlink(p, i);

                let c = self.new_container(self.stack.other());

                self.attach(c, 0, a);

                self.attach(c, 1, n);

                self.attach(p, i, c);

                self.reveal(n, true);

            }

            Some((p, i)) => {

                let n = self.new_leaf(w, self.size(p));

                self.attach(p, i + 1, n);

                self.rescale_hidden(n, true, &vis);

                self.reveal(n, true);

            }

            None => {

                let n = self.new_leaf(w, self.size(a));

                let c = self.new_container(self.stack.other());

                self.attach(c, 0, a);

                self.attach(c, 1, n);

                self.root = Some(c);

            }

        }

    }

    /// Inserts `w` after `anchor` by splitting `anchor` along `axis`. They

    /// share the space `anchor` had, and its siblings keep theirs.

    pub fn insert_split(&mut self, w: W, anchor: W, axis: Axis, vis: impl Fn(W) -> bool) {

        let Some(&a) = self.leaves.get(&anchor) else {

            return self.insert_row(w, vis);

        };

        let size = self.size(a);

        let n = self.new_leaf(w, size);

        match self.position(a) {

            Some((p, i)) if self.axis(p) == Some(axis) => {

                self.halve(a, axis);

                self.halve(n, axis);

                self.attach(p, i + 1, n);

            }

            Some((p, i)) => {

                self.unlink(p, i);

                let c = self.new_container(axis);

                self.attach(c, 0, a);

                self.attach(c, 1, n);

                self.attach(p, i, c);

            }

            None => {

                let c = self.new_container(axis);

                self.attach(c, 0, a);

                self.attach(c, 1, n);

                self.root = Some(c);

            }

        }

        self.reveal(n, true);

    }

    fn halve(&mut self, leaf: NodeId, axis: Axis) {

        if let Kind::Leaf(_, size) = &mut self.nodes[leaf].kind {

            size[axis.index()] /= 2.0;

        }

    }

    /// Scales the hidden siblings of a visible `n` arriving in or leaving its

    /// container by as much as the visible ones grow or shrink together, so

    /// the share they come back to does not drift as windows come and go.

    fn rescale_hidden(&mut self, n: NodeId, arriving: bool, vis: &impl Fn(W) -> bool) {

        let (Some((p, _)), true) = (self.position(n), self.visible(n, vis)) else {

            return;

        };

        let axis = self.axis(p).expect("parents are containers");

        let (shown, hidden): (Vec<NodeId>, Vec<NodeId>) = self

            .children(p)

            .iter()

            .filter(|&&c| c != n)

            .partition(|&&c| self.visible(c, vis));

        let rest: f64 = shown.iter().map(|&c| self.extent(c, axis, vis)).sum();

        if rest <= 0.0 {

            return;

        }

        let with = rest + self.extent(n, axis, vis);

        let by = if arriving { with / rest } else { rest / with };

        for c in hidden {

            self.scale(c, axis, by);

        }

    }

    /// Appends `w` as a row of its own along the stacking axis.

    pub fn insert_row(&mut self, w: W, vis: impl Fn(W) -> bool) {

        let size = self.root.map_or([1.0; 2], |r| self.size(r));

        let n = self.new_leaf(w, size);

        match self.root {

            None => self.root = Some(n),

            Some(r) if self.axis(r) == Some(self.stack) => {

                let end = self.children(r).len();

                self.attach(r, end, n);

                self.rescale_hidden(n, true, &vis);

            }

            Some(r) => {

                let c = self.new_container(self.stack);

                self.attach(c, 0, r);

                self.attach(c, 1, n);

                self.root = Some(c);

            }

        }

        self.reveal(n, true);

    }

    pub fn remove(&mut self, w: W, vis: impl Fn(W) -> bool) -> bool {

        let Some(&n) = self.leaves.get(&w) else {

            return false;

        };

        self.rescale_hidden(n, false, &vis);

        self.leaves.remove(&w);

        self.detach(n);

        self.nodes.remove(n);

        true

    }

    pub fn retain(&mut self, mut keep: impl FnMut(W) -> bool) {

        let gone: Vec<W> = self.windows().filter(|&w| !keep(w)).collect();

        for w in gone {

            self.remove(w, |_| false);

        }

    }

    /// Moves a window one step in `dir`. Leaves swap slots with the next

    /// visible sibling, enter a sibling container, or step out beside their

    /// container. Hidden windows never move, and an extended window moves

    /// its group as a rigid block.

    pub fn move_window(&mut self, w: W, dir: Dir, vis: impl Fn(W) -> bool) -> bool {

        let Some(&n) = self.leaves.get(&w).filter(|_| vis(w)) else {

            return false;

        };

        let block = self.group(w, &vis).unwrap_or(n);

        self.until_seen(&vis, |t| {

            let moved = t.move_node(block, dir, &vis, block == n);

            t.reveal(n, moved)

        })

    }

    /// Moves the group or container holding `w` as a block, merging it into

    /// the next lane in `dir` like `merge_window` when `merge` is set. A

    /// window sitting directly in the root is its own block.

    pub fn move_container(&mut self, w: W, dir: Dir, merge: bool, vis: impl Fn(W) -> bool) -> bool {

        let Some(&n) = self.leaves.get(&w).filter(|_| vis(w)) else {

            return false;

        };

        let group = self.group(w, &vis);

        self.until_seen(&vis, |t| {

            let block = group.unwrap_or_else(|| {

                t.nodes[n]

                    .parent

                    .filter(|&p| Some(p) != t.root)

                    .unwrap_or(n)

            });

            if merge {

                t.merge_node(n, block, dir, &vis)

            } else {

                let moved = t.move_node(block, dir, &vis, false);

                t.reveal(n, moved)

            }

        })

    }

    /// Moves a window to the end of the next visible lane in `dir` among its

    /// siblings. A lone window there becomes a stack with it. A window with no

    /// sibling lane that way leaves its stack as `move_window` would, so it

    /// takes one press to step out and another to merge. An extended window

    /// takes its group along.

    pub fn merge_window(&mut self, w: W, dir: Dir, vis: impl Fn(W) -> bool) -> bool {

        let Some(&n) = self.leaves.get(&w).filter(|_| vis(w)) else {

            return false;

        };

        let block = self.group(w, &vis).unwrap_or(n);

        self.until_seen(&vis, |t| t.merge_node(n, block, dir, &vis))

    }

    fn merge_node(&mut self, leaf: NodeId, n: NodeId, dir: Dir, vis: &impl Fn(W) -> bool) -> bool {

        if let Some((p, i)) = self.position(n)

            && self.axis(p) == Some(dir.axis())

            && let Some(j) = self.next_shown(p, i, dir, vis)

        {

            self.join_lane(n, p, j);

            return self.reveal(leaf, true);

        }

        let moved = self.move_node(n, dir, vis, true);

        self.reveal(leaf, moved)

    }

    /// The outermost container `w` is extended in, the block it moves as.

    fn group(&self, w: W, vis: &impl Fn(W) -> bool) -> Option<NodeId> {

        self.shown_path(w, vis)

            .into_iter()

            .take_while(|&(p, c)| self.expanded(p) == Some(c))

            .last()

            .map(|(p, _)| p)

    }

    /// Whether `w` is extended over windows it shares a container with.

    pub fn grouped(&self, w: W, vis: impl Fn(W) -> bool) -> bool {

        self.group(w, &vis).is_some()

    }

    /// The shown windows in the block `w` moves as, `w` alone when it is not

    /// extended.

    pub fn group_windows(&self, w: W, vis: impl Fn(W) -> bool) -> Vec<W> {

        let Some(g) = self.group(w, &vis) else {

            return vec![w];

        };

        let mut out = Vec::new();

        self.shown_leaves(g, &vis, &mut out);

        out

    }

    /// Whether a visible window extended over `w` takes all its room.

    pub fn eclipsed(&self, w: W, vis: impl Fn(W) -> bool) -> bool {

        let Some(&leaf) = self.leaves.get(&w) else {

            return false;

        };

        let mut cur = leaf;

        while let Some(p) = self.nodes[cur].parent {

            if self

                .expanded(p)

                .is_some_and(|e| e != cur && self.visible(e, &vis))

            {

                return true;

            }

            cur = p;

        }

        false

    }

    /// Repeats `step` until the visible layout changes, or undoes it all when

    /// it never does. Steps around hidden windows otherwise ate keypresses

    /// and rearranged windows the user could not see.

    fn until_seen(

        &mut self,

        vis: &impl Fn(W) -> bool,

        mut step: impl FnMut(&mut Self) -> bool,

    ) -> bool {

        let seen = |t: &Self| -> Vec<(W, Rect)> {

            t.arrange(Rect::new(0.0, 0.0, 1000.0, 1000.0), 0.0, vis)

                .into_iter()

                .filter(|p| vis(p.0))

                .collect()

        };

        let same = |a: &[(W, Rect)], b: &[(W, Rect)]| {

            a.len() == b.len()

                && a.iter().all(|(v, r)| {

                    b.iter().any(|(w, s)| {

                        v == w

                            && [r.x - s.x, r.y - s.y, r.w - s.w, r.h - s.h]

                                .iter()

                                .all(|d| d.abs() < 1e-6)

                    })

                })

        };

        let before = self.clone();

        let start = seen(self);

        for _ in 0..=before.nodes.len() {

            if !step(self) {

                break;

            }

            if !same(&start, &seen(self)) {

                return true;

            }

        }

        *self = before;

        false

    }

    /// Undoes any expansion along `leaf`'s ancestry that favours another

    /// child, once it has `moved` or arrived. An expanded container gives its

    /// other children no room, so a window moved into one vanished while

    /// staying focused.

    fn reveal(&mut self, leaf: NodeId, moved: bool) -> bool {

        let mut cur = leaf;

        while moved && let Some(p) = self.nodes[cur].parent {

            if let Kind::Container { expanded, .. } = &mut self.nodes[p].kind

                && expanded.is_some_and(|e| e != cur)

            {

                *expanded = None;

            }

            cur = p;

        }

        moved

    }

    /// Whether every container above `w` is expanded towards it.

    pub fn maximized(&self, w: W, vis: impl Fn(W) -> bool) -> bool {

        let path = self.shown_path(w, &vis);

        !path.is_empty() && path.iter().all(|&(p, c)| self.expanded(p) == Some(c))

    }

    /// `w`'s ancestry as (container, child) pairs, leaving out containers

    /// where the child has no visible sibling. Those look like nothing on

    /// screen, so expanding there would do nothing the user could see.

    fn shown_path(&self, w: W, vis: &impl Fn(W) -> bool) -> Vec<(NodeId, NodeId)> {

        let Some(&leaf) = self.leaves.get(&w) else {

            return Vec::new();

        };

        let mut path = Vec::new();

        let mut cur = leaf;

        while let Some(p) = self.nodes[cur].parent {

            let shared = self

                .children(p)

                .iter()

                .any(|&c| c != cur && self.visible(c, vis));

            if shared {

                path.push((p, cur));

            }

            cur = p;

        }

        path

    }

    /// Transposes the whole tree.

    pub fn flip(&mut self) {

        self.stack = self.stack.other();

        for node in self.nodes.values_mut() {

            match &mut node.kind {

                Kind::Container { axis, .. } => *axis = axis.other(),

                Kind::Leaf(_, size) => size.reverse(),

            }

        }