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

            }

        }

    }

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

            return false;

        };

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

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

            return false;

        };

        let block = self.nodes[n]

            .parent

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

            .unwrap_or(n);

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

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

            return false;

        };

        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(n, true);

        }

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

        self.reveal(n, moved)

    }

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

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

                        .sum::<f64>()

                        / children.len() as f64

                } else {

                    shown.iter().sum::<f64>() / shown.len() as f64

                }

            }

        }

    }

    /// The size a window placed in `n`'s stead would take, ignoring

    /// visibility since new windows land on any tag.

    fn size(&self, n: NodeId) -> Size {

        let all = |_: W| true;

        [

            self.extent(n, Axis::Horizontal, &all),

            self.extent(n, Axis::Vertical, &all),

        ]

    }

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

            let area = if self.visible(r, &vis) {

                area

            } else {

                Rect::new(area.x + area.w / 2.0, area.y, 0.0, area.h)

            };

            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, vis)).collect();

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

        if let Some(e) = expanded.and_then(|e| children.iter().position(|&c| c == 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, sub, gap, vis, out);

            }

            return;

        }

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

        let extents: Vec<f64> = children

            .iter()

            .zip(&shown)

            .map(|(&c, &s)| if s { self.extent(c, axis, vis) } else { 0.0 })

            .collect();

        let total: f64 = extents.iter().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() {