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

pub enum Axis {

    Horizontal,

    Vertical,

}

impl Axis {

    pub fn other(self) -> Self {

        match self {

            Self::Horizontal => Self::Vertical,

            Self::Vertical => Self::Horizontal,

        }

    }

}

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

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;

#[derive(Clone, Debug)]

struct Child {

    node: NodeId,

    weight: f64,

}

#[derive(Clone, Debug)]

enum Kind<W> {

    Leaf(W),

    Container {

        axis: Axis,

        children: Vec<Child>,

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

        };

        let n = self.new_leaf(w);

        match self.position(a) {

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

                let weight = self.unlink(p, i).weight;

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

                self.attach(c, 0, a, 1.0);

                self.attach(c, 1, n, 1.0);

                self.attach(p, i, c, weight);

            }

            Some((p, i)) => {

                let weight = self.mean_weight(p);

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

            }

            None => {

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

                self.attach(c, 0, a, 1.0);

                self.attach(c, 1, n, 1.0);

                self.root = Some(c);

            }

        }

    }

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

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

        let n = self.new_leaf(w);

        match self.root {

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

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

                let weight = self.mean_weight(r);

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

                self.attach(r, end, n, weight);

            }

            Some(r) => {

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

                self.attach(c, 0, r, 1.0);

                self.attach(c, 1, n, 1.0);

                self.root = Some(c);

            }

        }

    }

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

            return false;

        };

        self.move_node(n, dir, &vis, true)

    }

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

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

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

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

            return false;

        };

        let block = self.nodes[n]

            .parent

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

            .unwrap_or(n);

        self.move_node(block, dir, &vis, false)

    }

    /// Transposes the whole tree.

    pub fn flip(&mut self) {

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

        for node in self.nodes.values_mut() {

            if let Kind::Container { axis, .. } = &mut node.kind {

                *axis = axis.other();

            }

        }

    }

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

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

    pub fn toggle_expand(&mut self, w: W, skip: usize, count: usize) -> bool {

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

            return false;

        };

        let mut path = Vec::new();

        let mut cur = leaf;

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

            path.push((p, cur));

            cur = p;

        }

        let path: Vec<_> = path.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 weights, 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, i)) = self.position(cur) {

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

                let shown: Vec<usize> = (0..self.children(p).len())

                    .filter(|&j| self.visible(self.children(p)[j].node, &vis))

                    .collect();

                if shown.len() >= 2 {

                    let children = self.children_mut(p);

                    let sum: f64 = shown.iter().map(|&j| children[j].weight).sum();

                    let old = children[i].weight;

                    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 j in shown {

                        children[j].weight = if j == i {

                            new

                        } else {

                            children[j].weight * scale

                        };

                    }

                    return true;

                }

            }

            cur = p;

        }

        false

    }

    /// Lays the tree out in `area`. Hidden windows come back as zero-size

    /// rects on the boundary where they would reappear.

    pub fn arrange(&self, area: Rect, gap: f64, vis: impl Fn(W) -> bool) -> Vec<(W, Rect)> {

        let mut out = Vec::with_capacity(self.leaves.len());

        if let Some(r) = self.root {

            self.place(r, area, gap, &vis, &mut out);

        }

        out

    }

    fn place(

        &self,

        n: NodeId,

        r: Rect,

        gap: f64,

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

        out: &mut Vec<(W, Rect)>,

    ) {

        let (axis, children, expanded) = match &self.nodes[n].kind {

            Kind::Leaf(w) => return out.push((*w, r)),

            Kind::Container {

                axis,

                children,

                expanded,

            } => (*axis, children, *expanded),

        };

        let shown: Vec<bool> = children.iter().map(|c| self.visible(c.node, vis)).collect();

        let (start, len) = span(r, axis);

        if let Some(e) = expanded.and_then(|e| children.iter().position(|c| c.node == e))

            && shown[e]

        {

            for (i, c) in children.iter().enumerate() {

                let sub = match i.cmp(&e) {

                    std::cmp::Ordering::Equal => r,

                    std::cmp::Ordering::Less => with_span(r, axis, start, 0.0),

                    std::cmp::Ordering::Greater => with_span(r, axis, start + len, 0.0),

                };

                self.place(c.node, sub, gap, vis, out);

            }

            return;

        }

        let count = shown.iter().filter(|s| **s).count();

        let total: f64 = children

            .iter()

            .zip(&shown)

            .filter(|(_, s)| **s)

            .map(|(c, _)| c.weight)

            .sum();

        let gaps = count.saturating_sub(1) as f64;

        let gap = if gaps > 0.0 { gap.min(len / gaps) } else { gap };

        let avail = len - gap * gaps;

        let last = shown.iter().rposition(|s| *s);

        let mut pos = start;

        let mut seen = false;

        for (i, c) in children.iter().enumerate() {

            if shown[i] {

                if seen {

                    pos += gap;

                }

                let l = avail * c.weight / total;

                self.place(c.node, with_span(r, axis, pos, l), gap, vis, out);

                pos += l;

                seen = true;

            } else {

                let between = seen && last.is_some_and(|l| l > i);

                let at = if between { pos + gap / 2.0 } else { pos };

                self.place(c.node, with_span(r, axis, at, 0.0), gap, vis, out);

            }

        }

    }

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

        let mut cur = n;

        while let Some((p, i)) = self.position(cur) {

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

                if cur != n {

                    let at = if dir.forward() { i + 1 } else { i };

                    self.relocate(n, p, at);

                    return true;

                }

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

                    let sib = self.children(p)[j].node;

                    if enter && self.axis(sib).is_some() {

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

                        self.relocate(n, sib, end);

                    } else {

                        let children = self.children_mut(p);

                        children[i].node = sib;

                        children[j].node = n;

                    }

                    return true;

                }

            }

            cur = p;

        }

        if cur == n || self.axis(cur) == Some(dir.axis()) {

            return false;

        }

        let (old, i) = self.position(n).expect("n is below the root");

        self.unlink(old, i);

        self.normalize(old);

        let r = self.root.expect("the rest of the tree survives");

        if self.axis(r) == Some(dir.axis()) {

            let weight = self.mean_weight(r);

            let at = if dir.forward() {

                self.children(r).len()

            } else {

                0

            };

            self.attach(r, at, n, weight);

        } else {

            let c = self.new_container(dir.axis());

            self.attach(c, 0, r, 1.0);

            let at = if dir.forward() {

                self.children(c).len()

            } else {

                0

            };

            self.attach(c, at, n, 1.0);

            self.root = Some(c);

        }

        true

    }

    fn next_shown(&self, p: NodeId, i: usize, dir: Dir, vis: &impl Fn(W) -> bool) -> Option<usize> {

        let children = self.children(p);

        if dir.forward() {

            (i + 1..children.len()).find(|&j| self.visible(children[j].node, vis))

        } else {

            (0..i).rev().find(|&j| self.visible(children[j].node, vis))

        }

    }

    fn relocate(&mut self, n: NodeId, p: NodeId, at: usize) {

        let (old, i) = self.position(n).expect("relocated nodes have a parent");

        let weight = self.mean_weight(p);

        let at = if old == p && i < at { at - 1 } else { at };

        self.unlink(old, i);

        self.attach(p, at, n, weight);

        self.normalize(old);

    }

    /// Places `node` among `p`'s children. A container running along `p`'s

    /// axis is dissolved into `p` so axes keep alternating.

    fn attach(&mut self, p: NodeId, at: usize, node: NodeId, weight: f64) {

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

            let Some(Node {

                kind: Kind::Container { children, .. },

                ..

            }) = self.nodes.remove(node)

            else {

                unreachable!("only containers have an axis");

            };

            let total: f64 = children.iter().map(|c| c.weight).sum();

            for (k, c) in children.into_iter().enumerate() {

                self.nodes[c.node].parent = Some(p);

                let weight = weight * c.weight / total;

                self.children_mut(p).insert(

                    at + k,

                    Child {

                        node: c.node,

                        weight,

                    },

                );

            }

        } else {

            self.nodes[node].parent = Some(p);

            self.children_mut(p).insert(at, Child { node, weight });

        }

    }

    fn unlink(&mut self, p: NodeId, i: usize) -> Child {

        let Kind::Container {

            children, expanded, ..

        } = &mut self.nodes[p].kind

        else {

            unreachable!("parents are containers");

        };

        let child = children.remove(i);

        if *expanded == Some(child.node) {

            *expanded = None;

        }

        self.nodes[child.node].parent = None;

        child

    }

    fn detach(&mut self, n: NodeId) {