use std::{collections::HashMap, hash::Hash};

use super::layout::{Direction, NodeLayout, Orientation, Rect};

pub type NodeId = u32;

pub enum NodeContent<T> {
    Container(ContainerNode),
    Window(WindowNode<T>),
}

pub struct Node<T> {
    parent: Option<NodeId>,
    layout: NodeLayout,

    content: NodeContent<T>,
}

pub struct ContainerNode {
    children: Vec<NodeId>,
    orientation: Orientation,
}

pub struct WindowNode<T> {
    pub wid: T,
}

pub struct Workspace<T> {
    nodes: HashMap<NodeId, Node<T>>,
    wid_to_nid: HashMap<T, NodeId>,
    last_node_id: NodeId,

    root: NodeId,
    focus: Option<NodeId>,

    margin: u32,
    spacing: u32,

    // reservations_top: Vec<Reservation<T>>,
    reservation_top: u32,

    width: u32,
    height: u32,
}

impl<T> Node<T> {
    pub fn as_window(&self) -> Option<&WindowNode<T>> {
        match &self.content {
            NodeContent::Window(window) => Some(window),
            _ => None,
        }
    }
    pub fn as_container(&self) -> Option<&ContainerNode> {
        match &self.content {
            NodeContent::Container(container) => Some(container),
            _ => None,
        }
    }
    pub fn as_container_mut(&mut self) -> Option<&mut ContainerNode> {
        match &mut self.content {
            NodeContent::Container(container) => Some(container),
            _ => None,
        }
    }
}

impl<T: Eq + Hash + Copy> Workspace<T> {
    pub fn new(width: u32, height: u32) -> Self {
        let root = Node {
            parent: None,
            layout: NodeLayout::empty(),
            content: NodeContent::Container(ContainerNode {
                children: vec![],
                orientation: Orientation::Horizontal,
            }),
        };
        let nodes = HashMap::from_iter([(0, root)]);

        let mut this = Self {
            nodes,
            wid_to_nid: HashMap::new(),
            last_node_id: 0,

            margin: 4,
            spacing: 4,

            root: 0,
            focus: None,

            width,
            height,
            reservation_top: 0,
        };

        this.update_layout();

        this
    }

    pub fn has_windows(&self) -> bool {
        self.all_windows().next().is_some()
    }

    pub fn all_windows(&self) -> impl Iterator<Item = (T, &NodeLayout)> + '_ {
        self.nodes.iter().filter_map(|(_, node)| {
            let window = node.as_window()?;
            Some((window.wid, &node.layout))
        })
    }

    pub fn dirty_windows(&self) -> impl Iterator<Item = (T, Rect)> + '_ {
        self.all_windows().filter_map(|(wid, layout)| {
            if !layout.clear_dirty() {
                return None;
            }
            let rect = layout.get()?;
            Some((wid, rect))
        })
    }

    pub fn window_towards_wrap(&self, direction: Direction) -> Option<T> {
        let current_nid = self.focus?;
        let current_node = self.nodes.get(&current_nid)?;
        let parent_nid = current_node.parent?;
        let parent_node = self.nodes.get(&parent_nid)?;
        let parents_parent_nid = parent_node.parent;
        let parent_container = parent_node.as_container()?;

        let (orientation, delta) = direction.split();

        let nid = if orientation == parent_container.orientation {
            log::info!("Within parent {delta}, {:?}", parent_container.orientation);
            let position_in_parent = parent_container
                .children
                .iter()
                .position(|&n| n == current_nid)?;

            if delta > 0 {
                if position_in_parent < parent_container.children.len() - 1 {
                    parent_container.children[position_in_parent + 1]
                } else {
                    parent_container.children[0]
                }
            } else {
                if position_in_parent > 0 {
                    parent_container.children[position_in_parent - 1]
                } else {
                    parent_container.children[parent_container.children.len() - 1]
                }
            }
        } else if let Some(parents_parent_nid) = parents_parent_nid {
            let parents_parent_node = self.nodes.get(&parents_parent_nid)?;
            let parents_parent_container = parents_parent_node.as_container()?;
            assert_eq!(parents_parent_container.orientation, orientation);
            let position_in_parent = parents_parent_container
                .children
                .iter()
                .position(|&n| n == parent_nid)?;
            log::info!(
                "Within parent's parent {delta}, {:?}",
                parents_parent_container.orientation
            );
            log::info!("position_in_parent = {position_in_parent}");

            if delta > 0 {
                if position_in_parent < parent_container.children.len() - 1 {
                    parents_parent_container.children[position_in_parent + 1]
                } else {
                    parents_parent_container.children[0]
                }
            } else {
                if position_in_parent > 0 {
                    parents_parent_container.children[position_in_parent - 1]
                } else {
                    parents_parent_container.children[parent_container.children.len() - 1]
                }
            }
        } else {
            return None;
        };

        self.first_window_in(nid)
    }

    pub fn first_window_in(&self, nid: NodeId) -> Option<T> {
        let node = self.nodes.get(&nid)?;

        match &node.content {
            NodeContent::Window(window) => Some(window.wid),
            NodeContent::Container(container) => {
                let first = *container.children.first()?;
                self.first_window_in(first)
            }
        }
    }

    pub fn move_window(&mut self, direction: Direction) -> bool
    where
        T: core::fmt::Debug,
    {
        if let Some(container_nid) = self.move_window_inner(direction) {
            self.invalidate_layout(container_nid);
            self.fixup_containers(self.root);
            self.dump(self.root, 0);
            true
        } else {
            false
        }
    }

    fn dump(&self, nid: NodeId, depth: usize)
    where
        T: core::fmt::Debug,
    {
        let Some(node) = self.nodes.get(&nid) else {
            return;
        };

        let rect = node.layout.get();

        match &node.content {
            NodeContent::Window(window) => {
                log::info!(
                    "{:depth$}Window #{:?} {rect:?}",
                    "",
                    window.wid,
                    depth = depth * 2
                );
            }
            NodeContent::Container(container) => {
                log::info!("{:depth$}Container {rect:?} {{", "", depth = depth * 2);
                for &child_nid in container.children.iter() {
                    self.dump(child_nid, depth + 1);
                }
                log::info!("{:depth$}}}", "", depth = depth * 2);
            }
        }
    }

    // pub fn walk_windows<F: FnMut(T)>(&self, mut mapper: F) {
    //     self.walk_windows_inner(&mut mapper, self.root)
    // }

    // fn walk_windows_inner<F: FnMut(T)>(&self, mapper: &mut F, nid: NodeId) {
    //     // let Some(node) = self.nodes.get(&nid) else {
    //     //     return;
    //     // };

    //     // match &node.content {
    //     //     NodeContent::Container(container) => {
    //     //         for &child_nid in container.children.iter() {
    //     //             self.walk_windows_inner(mapper, child_nid);
    //     //         }
    //     //     }
    //     //     NodeContent::Window(window) => {
    //     //         mapper(window.wid);
    //     //     }
    //     // }
    // }

    fn invalidate_layout(&mut self, nid: NodeId) {
        // for reservation in self.reservations_top.iter() {
        //     reservation.layout.clear();
        // }
        self.invalidate_layout_inner(nid);
        self.update_layout();
    }

    fn invalidate_layout_inner(&self, nid: NodeId) {
        let Some(node) = self.nodes.get(&nid) else {
            return;
        };

        node.layout.clear();

        match &node.content {
            NodeContent::Container(container) => {
                for &child_nid in container.children.iter() {
                    self.invalidate_layout_inner(child_nid);
                }
            }
            _ => (),
        }
    }

    fn fixup_containers(&mut self, nid: NodeId) {
        let Some(node) = self.nodes.get_mut(&nid) else {
            return;
        };

        if let NodeContent::Container(container) = &mut node.content {
            if container.children.is_empty() {
                let Some(parent_nid) = node.parent else {
                    return;
                };
                // Remove the empty container
                if let Some(parent_container) = self
                    .nodes
                    .get_mut(&parent_nid)
                    .and_then(Node::as_container_mut)
                {
                    parent_container.children.retain(|&n| n != nid);
                    self.nodes.remove(&nid);
                    self.invalidate_layout(parent_nid);
                }
            } else if container.children.len() == 1 {
                // Remove the empty container and reparent its only child
                if let Some(parent_nid) = node.parent {
                    let child_nid = container.children.remove(0);
                    let [Some(child), Some(parent)] =
                        self.nodes.get_many_mut([&child_nid, &parent_nid])
                    else {
                        return;
                    };

                    child.parent = Some(parent_nid);
                    if let Some(parent_container) = parent.as_container_mut() {
                        parent_container.children.retain(|&n| n != nid);
                        parent_container.children.push(child_nid);
                    }

                    self.invalidate_layout(parent_nid);
                }
            } else {
                // TODO borrowing stuff
                let children = container.children.clone();
                for child_nid in children {
                    self.fixup_containers(child_nid);
                }
            }
        }
    }

    fn move_window_inner(&mut self, direction: Direction) -> Option<NodeId> {
        let current_nid = self.focus?;
        let current_node = self.nodes.get(&current_nid)?;
        let parent_nid = current_node.parent?;
        let parent_node = self.nodes.get_mut(&parent_nid)?;
        let parents_parent_nid = parent_node.parent;
        let parent_container = parent_node.as_container_mut()?;

        let (orientation, delta) = direction.split();

        if parent_container.orientation == orientation {
            let position_in_parent = parent_container
                .children
                .iter()
                .position(|&n| n == current_nid)?;

            // TODO check if this item is also a container

            if delta < 0 && position_in_parent > 0 {
                parent_container
                    .children
                    .swap(position_in_parent - 1, position_in_parent);
                Some(parent_nid)
            } else if delta > 0 && position_in_parent < parent_container.children.len() - 1 {
                parent_container
                    .children
                    .swap(position_in_parent, position_in_parent + 1);
                Some(parent_nid)
            } else {
                None
            }
        } else if let Some(parents_parent_nid) = parents_parent_nid {
            // Go to parent's parent
            let parents_parent_node = self.nodes.get_mut(&parents_parent_nid)?;
            let parents_parent_container = parents_parent_node.as_container_mut()?;
            assert_eq!(parents_parent_container.orientation, orientation);
            todo!()
        } else {
            self.last_node_id += 1;
            let new_root_nid = self.last_node_id;

            // Remove child from parent
            parent_container.children.retain(|&n| n != current_nid);

            self.nodes.get_mut(&current_nid)?.parent = Some(new_root_nid);
            self.nodes.get_mut(&self.root)?.parent = Some(new_root_nid);

            let children = if delta > 0 {
                vec![self.root, current_nid]
            } else {
                vec![current_nid, self.root]
            };

            // Create a new root
            let new_root = Node {
                parent: None,
                layout: NodeLayout::empty(),
                content: NodeContent::Container(ContainerNode {
                    children,
                    orientation,
                }),
            };
            self.root = new_root_nid;
            self.nodes.insert(new_root_nid, new_root);

            Some(self.root)
        }
    }

    pub fn focus_window(&mut self, wid: T) -> (Option<T>, Option<NodeId>) {
        let old = self.focus.take().and_then(|nid| self.wid(nid));
        let new = self.wid_to_nid.get(&wid).copied();
        self.focus = new;
        (old, new)
    }

    pub fn focused_window(&self) -> Option<T> {
        let nid = self.focus?;
        Some(self.nodes.get(&nid)?.as_window()?.wid)
    }

    pub(super) fn window(&self, wid: T) -> Option<&Node<T>> {
        let nid = *self.wid_to_nid.get(&wid)?;
        self.nodes.get(&nid)
    }

    // pub fn reservation(&self, wid: T) -> Option<&Reservation<T>> {
    //     self.reservations_top.iter().find(|r| r.wid == wid)
    // }

    pub(super) fn window_frame(&self, wid: T) -> Option<Rect> {
        self.window(wid)?.layout.get()
    }

    pub(super) fn resize(&mut self, w: u32, h: u32, reservation_top: u32) {
        self.width = w;
        self.height = h;
        self.reservation_top = reservation_top;
        self.update_layout();
    }

    // fn update_reservation_layout(&mut self) -> u32 {
    //     res_y
    // }

    pub(super) fn update_layout(&mut self) {
        self.update_layout_for(
            self.root,
            Rect {
                x: self.margin,
                y: self.margin + self.reservation_top,
                w: self.width - self.margin * 2,
                h: self.height - self.margin * 2 - self.reservation_top,
            },
        );
    }

    pub(super) fn create_window(&mut self, wid: T) -> bool {
        self.create_window_in(self.root, wid)
    }

    pub fn remove_window(&mut self, wid: T) -> bool {
        let Some(nid) = self.wid_to_nid.remove(&wid) else {
            return false;
        };
        if self.remove_inner(nid) {
            if self.focus == Some(nid) {
                let wid = self.first_window_in(self.root);
                if let Some(wid) = wid {
                    self.focus_window(wid);
                } else {
                    self.focus = None;
                }
            }

            self.fixup_containers(self.root);
            self.update_layout();
            true
        } else {
            false
        }
    }

    fn create_window_in(&mut self, container_nid: NodeId, wid: T) -> bool {
        self.last_node_id += 1;
        let nid = self.last_node_id;
        self.nodes.insert(
            nid,
            Node {
                parent: Some(container_nid),
                layout: NodeLayout::empty(),
                content: NodeContent::Window(WindowNode { wid }),
            },
        );
        self.wid_to_nid.insert(wid, nid);
        if !self.add_child_to(container_nid, nid) {
            self.wid_to_nid.remove(&wid);
            self.nodes.remove(&nid);
            return false;
        }
        self.update_layout();
        true
    }

    fn wid(&self, nid: NodeId) -> Option<T> {
        self.nodes.get(&nid).and_then(|node| match &node.content {
            NodeContent::Window(window) => Some(window.wid),
            _ => None,
        })
    }

    fn update_layout_for(&self, node_id: NodeId, rect: Rect) {
        let Some(node) = self.nodes.get(&node_id) else {
            log::warn!("update_layout_for: no node {node_id}");
            return;
        };

        node.layout.set(rect);

        match &node.content {
            NodeContent::Container(container) => {
                if container.children.is_empty() {
                    return;
                }

                match container.orientation {
                    Orientation::Vertical => {
                        let hstep = (rect.h - self.spacing * container.children.len() as u32)
                            / container.children.len() as u32;
                        let ystep = hstep + self.spacing;

                        for (i, &child_nid) in container.children.iter().enumerate() {
                            let i = i as u32;
                            let child_rect = Rect {
                                x: rect.x,
                                y: rect.y + i * ystep,
                                w: rect.w,
                                h: hstep,
                            };
                            log::info!("Frame #{child_nid} size: {child_rect:?}");
                            self.update_layout_for(child_nid, child_rect);
                        }
                    }
                    Orientation::Horizontal => {
                        let wstep = (rect.w - self.spacing * container.children.len() as u32)
                            / container.children.len() as u32;
                        let xstep = wstep + self.spacing;

                        for (i, &child_nid) in container.children.iter().enumerate() {
                            let i = i as u32;
                            let child_rect = Rect {
                                x: rect.x + i * xstep,
                                y: rect.y,
                                w: wstep,
                                h: rect.h,
                            };
                            log::info!("Frame #{child_nid} size: {child_rect:?}");
                            self.update_layout_for(child_nid, child_rect);
                        }
                    }
                }
            }
            _ => (),
        }
    }

    fn remove_inner(&mut self, node_id: NodeId) -> bool {
        if node_id == self.root {
            return false;
        }
        let Some(node) = self.nodes.remove(&node_id) else {
            return false;
        };
        let container_nid = unsafe { node.parent.unwrap_unchecked() };
        let Some(container_node) = self.nodes.get_mut(&container_nid) else {
            return false;
        };
        let NodeContent::Container(container) = &mut container_node.content else {
            return false;
        };
        container_node.layout.clear();
        container.children.retain(|&c| c != node_id);
        // Only root container can be empty
        debug_assert!(!(container_node.parent.is_some() && container.children.is_empty()));
        if container.children.len() == 1 && container_node.parent.is_some() {
            // Lift the remaining node to the parent container
            todo!()
        }

        true
    }

    fn add_child_to(&mut self, container_nid: NodeId, child_nid: NodeId) -> bool {
        let [Some(container_node), Some(child_node)] =
            self.nodes.get_many_mut([&container_nid, &child_nid])
        else {
            return false;
        };
        let NodeContent::Container(container) = &mut container_node.content else {
            return false;
        };

        container_node.layout.clear();
        child_node.layout.clear();
        child_node.parent = Some(container_nid);
        container.children.push(child_nid);

        true
    }
}
