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) {

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

            return self.insert_row(w);

        };

        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.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) {

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

            return self.insert_row(w);

        };

        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;

        }

    }

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

    pub fn insert_row(&mut self, w: W) {

        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);

            }

            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) -> bool {

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

            return false;

        };

        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);

        }

    }

    /// 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(),

            }

        }

    }

    /// Toggles expansion for `count` levels of `w`'s visible ancestry,

    /// starting `skip` levels up. Level 0 expands the window inside the

    /// nearest container where it has a visible sibling.

    pub fn toggle_expand(

        &mut self,

        w: W,

        skip: usize,

        count: usize,

        vis: impl Fn(W) -> bool,

    ) -> bool {

        let path: Vec<_> = self

            .shown_path(w, &vis)

            .into_iter()

            .skip(skip)

            .take(count)

            .collect();

        if path.is_empty() {

            return false;

        }

        let all = path.iter().all(|&(p, c)| self.expanded(p) == Some(c));

        for (p, c) in path {

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

                *expanded = (!all).then_some(c);

            }

        }

        true

    }

    /// Grows `w` by `delta`, a fraction of the visible length of the nearest

    /// container along `axis` where it has a visible sibling. Hidden siblings

    /// keep their sizes, and with them their share.

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

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