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.

    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;

        };

        self.until_seen(&vis, |t| {

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

            t.reveal(n, moved)

        })

    }

    /// Moves the container holding `w` as a rigid block. A window sitting

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

    pub fn move_container(&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;

        };

        self.until_seen(&vis, |t| {

            let block = t.nodes[n]

                .parent

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

                .unwrap_or(n);

            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.

    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;

        };

        self.until_seen(&vis, |t| {

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

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

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

            {

                t.join_lane(n, p, j);

                return t.reveal(n, true);

            }

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

            t.reveal(n, moved)

        })

    }

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

            return false;

        };

        let mut cur = leaf;

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

            if self.axis(p) == Some(axis) {

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

                    .children(p)

                    .iter()

                    .filter(|&&c| self.visible(c, &vis))

                    .map(|&c| (c, self.extent(c, axis, &vis)))

                    .collect();

                if shown.len() >= 2 {

                    let sum: f64 = shown.iter().map(|s| s.1).sum();

                    let old = self.extent(cur, axis, &vis);

                    let max = 1.0 - MIN_SHARE * (shown.len() - 1) as f64;

                    let new = (old / sum + delta).clamp(MIN_SHARE, max) * sum;

                    let scale = (sum - new) / (sum - old);

                    for (c, _) in shown {

                        self.scale(c, axis, if c == cur { new / old } else { scale });

                    }

                    return true;

                }

            }

            cur = p;

        }

        false

    }

    /// Scales every window under `n` along `axis`, hidden ones too, so they

    /// keep their proportions for when they are shown again.

    fn scale(&mut self, n: NodeId, axis: Axis, by: f64) {

        match &mut self.nodes[n].kind {

            Kind::Leaf(_, size) => size[axis.index()] *= by,

            Kind::Container { children, .. } => {

                for c in children.clone() {

                    self.scale(c, axis, by);

                }

            }

        }

    }

    /// How much room `n` asks for along `axis` among its siblings. A

    /// container asks for the mean of what its visible children ask for, so

    /// windows of other tags sharing a lane change nothing on screen, and a

    /// window carries its own size wherever it moves.

    fn extent(&self, n: NodeId, axis: Axis, vis: &impl Fn(W) -> bool) -> f64 {

        match &self.nodes[n].kind {

            Kind::Leaf(_, size) => size[axis.index()],

            Kind::Container { children, .. } => {

                let shown: Vec<f64> = children

                    .iter()

                    .filter(|&&c| self.visible(c, vis))

                    .map(|&c| self.extent(c, axis, vis))

                    .collect();

                if shown.is_empty() {

                    children

                        .iter()