//! Pure algorithms over `spacesh_proto::LayoutNode`. No I/O.
use spacesh_proto::layout::{LayoutNode, Orient};
use spacesh_proto::ids::SurfaceId;

/// Minimum ratio a panel may shrink to (5%).
const MIN_RATIO: f32 = 0.05;

/// Collect all surface ids in the tree, left-to-right.
pub fn leaves(node: &LayoutNode) -> Vec<SurfaceId> {
    let mut out = Vec::new();
    collect(node, &mut out);
    out
}
fn collect(node: &LayoutNode, out: &mut Vec<SurfaceId>) {
    match node {
        LayoutNode::Leaf { surface_id } => out.push(surface_id.clone()),
        LayoutNode::Split { children, .. } => children.iter().for_each(|c| collect(c, out)),
    }
}

/// Split the leaf `target` by inserting `new_id` as a sibling on `dir`.
/// Returns true if the target was found and split.
pub fn split_leaf(root: &mut LayoutNode, target: &SurfaceId, dir: Orient, after: bool, new_id: SurfaceId) -> bool {
    // If root itself is the target leaf, replace it with a split.
    if let LayoutNode::Leaf { surface_id } = root {
        if surface_id == target {
            let existing = root.clone();
            let new_leaf = LayoutNode::leaf(new_id);
            let children = if after { vec![existing, new_leaf] } else { vec![new_leaf, existing] };
            *root = LayoutNode::Split { orient: dir, ratios: even(children.len()), children };
            return true;
        }
        return false;
    }
    if let LayoutNode::Split { orient, ratios, children } = root {
        // If a direct child is the target leaf AND this split matches `dir`, insert as sibling.
        if *orient == dir {
            if let Some(i) = children.iter().position(|c| is_leaf(c, target)) {
                children.insert(i + if after { 1 } else { 0 }, LayoutNode::leaf(new_id));
                *ratios = even(children.len());
                return true;
            }
        }
        // Otherwise recurse.
        for c in children.iter_mut() {
            if split_leaf(c, target, dir, after, new_id.clone()) {
                return true;
            }
        }
    }
    false
}

/// Remove the leaf `target`. Collapses empty/now-single-child splits and promotes
/// single children. Returns the new root (None if the tree became empty).
pub fn remove_leaf(root: LayoutNode, target: &SurfaceId) -> Option<LayoutNode> {
    match root {
        LayoutNode::Leaf { surface_id } => {
            if &surface_id == target { None } else { Some(LayoutNode::Leaf { surface_id }) }
        }
        LayoutNode::Split { orient, children, .. } => {
            let kept: Vec<LayoutNode> = children
                .into_iter()
                .filter_map(|c| remove_leaf(c, target))
                .collect();
            match kept.len() {
                0 => None,
                1 => Some(kept.into_iter().next().unwrap()), // promote single child
                n => Some(LayoutNode::Split { orient, ratios: even(n), children: kept }),
            }
        }
    }
}

/// Set ratios on the split node addressed by `path` (child indices from root).
/// Normalizes to sum 1.0 and clamps each to >= MIN_RATIO. Returns false if the
/// path is invalid or the length does not match the node's child count.
pub fn set_ratios(root: &mut LayoutNode, path: &[u32], ratios: &[f32]) -> bool {
    let Some(node) = node_at_mut(root, path) else { return false };
    if let LayoutNode::Split { ratios: r, children, .. } = node {
        if ratios.len() != children.len() {
            return false;
        }
        *r = normalize_clamp(ratios);
        true
    } else {
        false
    }
}

/// Move leaf `src` to sit on `edge` of leaf `target`. Returns the new root.
/// No-op (returns the original) if src == target or either is missing.
pub fn move_leaf(root: LayoutNode, src: &SurfaceId, target: &SurfaceId, edge: spacesh_proto::message::Edge) -> LayoutNode {
    use spacesh_proto::message::Edge;
    if src == target || !contains(&root, src) || !contains(&root, target) {
        return root;
    }
    let Some(removed) = remove_leaf(root, src) else { return LayoutNode::leaf(src.clone()) };
    let (orient, after) = match edge {
        Edge::Left => (Orient::H, false),
        Edge::Right => (Orient::H, true),
        Edge::Top => (Orient::V, false),
        Edge::Bottom => (Orient::V, true),
    };
    let mut root = removed;
    split_leaf(&mut root, target, orient, after, src.clone());
    root
}

// ---- helpers ----

fn is_leaf(node: &LayoutNode, id: &SurfaceId) -> bool {
    matches!(node, LayoutNode::Leaf { surface_id } if surface_id == id)
}
fn contains(node: &LayoutNode, id: &SurfaceId) -> bool {
    leaves(node).iter().any(|s| s == id)
}
fn even(n: usize) -> Vec<f32> {
    vec![1.0 / n as f32; n]
}
fn normalize_clamp(ratios: &[f32]) -> Vec<f32> {
    // Two-pass: clamp all to MIN_RATIO, then normalize. If normalization would
    // bring any value back below MIN_RATIO, pin those and redistribute the rest.
    let n = ratios.len();
    if n == 0 { return vec![]; }
    let mut r: Vec<f32> = ratios.iter().map(|v| v.max(MIN_RATIO)).collect();
    // Iteratively pin items that would end up below MIN_RATIO after normalization.
    let mut pinned = vec![false; n];
    for _ in 0..n {
        let pin_sum: f32 = r.iter().zip(&pinned).filter(|(_, p)| **p).map(|(v, _)| *v).sum();
        let free_sum: f32 = r.iter().zip(&pinned).filter(|(_, p)| !**p).map(|(v, _)| *v).sum();
        let remaining = 1.0 - pin_sum;
        let mut changed = false;
        for i in 0..n {
            if pinned[i] { continue; }
            let normalized = if free_sum > 0.0 { r[i] / free_sum * remaining } else { remaining / n as f32 };
            if normalized < MIN_RATIO {
                r[i] = MIN_RATIO;
                pinned[i] = true;
                changed = true;
            }
        }
        if !changed { break; }
    }
    // Final normalization of unpinned values.
    let pin_sum: f32 = r.iter().zip(&pinned).filter(|(_, p)| **p).map(|(v, _)| *v).sum();
    let free_sum: f32 = r.iter().zip(&pinned).filter(|(_, p)| !**p).map(|(v, _)| *v).sum();
    let remaining = (1.0 - pin_sum).max(0.0);
    let mut result = vec![0.0f32; n];
    for i in 0..n {
        if pinned[i] {
            result[i] = r[i];
        } else {
            result[i] = if free_sum > 0.0 { r[i] / free_sum * remaining } else { remaining / n as f32 };
        }
    }
    result
}
fn node_at_mut<'a>(root: &'a mut LayoutNode, path: &[u32]) -> Option<&'a mut LayoutNode> {
    let mut cur = root;
    for &idx in path {
        match cur {
            LayoutNode::Split { children, .. } => {
                cur = children.get_mut(idx as usize)?;
            }
            LayoutNode::Leaf { .. } => return None,
        }
    }
    Some(cur)
}

#[cfg(test)]
mod tests {
    use super::*;
    fn sid(s: &str) -> SurfaceId { SurfaceId(s.into()) }

    #[test]
    fn split_root_leaf_creates_split() {
        let mut root = LayoutNode::leaf(sid("s_1"));
        assert!(split_leaf(&mut root, &sid("s_1"), Orient::H, true, sid("s_2")));
        assert_eq!(leaves(&root), vec![sid("s_1"), sid("s_2")]);
    }

    #[test]
    fn split_same_orient_appends_as_sibling() {
        let mut root = LayoutNode::Split {
            orient: Orient::H, ratios: vec![0.5, 0.5],
            children: vec![LayoutNode::leaf(sid("s_1")), LayoutNode::leaf(sid("s_2"))],
        };
        split_leaf(&mut root, &sid("s_2"), Orient::H, true, sid("s_3"));
        // 3 children in one row, even ratios.
        match &root {
            LayoutNode::Split { children, ratios, .. } => {
                assert_eq!(children.len(), 3);
                assert!((ratios.iter().sum::<f32>() - 1.0).abs() < 1e-5);
            }
            _ => panic!(),
        }
        assert_eq!(leaves(&root), vec![sid("s_1"), sid("s_2"), sid("s_3")]);
    }

    #[test]
    fn remove_promotes_single_child() {
        let root = LayoutNode::Split {
            orient: Orient::H, ratios: vec![0.5, 0.5],
            children: vec![LayoutNode::leaf(sid("s_1")), LayoutNode::leaf(sid("s_2"))],
        };
        let after = remove_leaf(root, &sid("s_2")).unwrap();
        assert_eq!(after, LayoutNode::leaf(sid("s_1"))); // split collapsed to the surviving leaf
    }

    #[test]
    fn remove_last_leaf_returns_none() {
        let root = LayoutNode::leaf(sid("s_1"));
        assert!(remove_leaf(root, &sid("s_1")).is_none());
    }

    #[test]
    fn set_ratios_normalizes_and_clamps() {
        let mut root = LayoutNode::Split {
            orient: Orient::H, ratios: vec![0.5, 0.5],
            children: vec![LayoutNode::leaf(sid("s_1")), LayoutNode::leaf(sid("s_2"))],
        };
        assert!(set_ratios(&mut root, &[], &[0.0, 1.0]));
        if let LayoutNode::Split { ratios, .. } = &root {
            assert!(ratios[0] >= MIN_RATIO);
            assert!((ratios.iter().sum::<f32>() - 1.0).abs() < 1e-5);
        }
    }

    #[test]
    fn set_ratios_wrong_len_rejected() {
        let mut root = LayoutNode::Split {
            orient: Orient::H, ratios: vec![0.5, 0.5],
            children: vec![LayoutNode::leaf(sid("s_1")), LayoutNode::leaf(sid("s_2"))],
        };
        assert!(!set_ratios(&mut root, &[], &[1.0]));
    }

    #[test]
    fn move_leaf_to_right_of_target() {
        let root = LayoutNode::Split {
            orient: Orient::V, ratios: vec![0.5, 0.5],
            children: vec![LayoutNode::leaf(sid("s_1")), LayoutNode::leaf(sid("s_2"))],
        };
        let after = move_leaf(root, &sid("s_1"), &sid("s_2"), spacesh_proto::message::Edge::Right);
        assert_eq!(leaves(&after), vec![sid("s_2"), sid("s_1")]);
    }

    #[test]
    fn move_onto_self_is_noop() {
        let root = LayoutNode::leaf(sid("s_1"));
        let after = move_leaf(root.clone(), &sid("s_1"), &sid("s_1"), spacesh_proto::message::Edge::Right);
        assert_eq!(after, root);
    }
}