use core::num::NonZeroU16;

use taffy::geometry::{Point as TaffyPoint, Rect as TaffyRect, Size as TaffySize};
use taffy::prelude::{
    AlignItems, AvailableSpace, Dimension, Display, FlexDirection, GridPlacement,
    GridTemplateComponent, JustifyContent, Layout as TaffyLayout, Line as TaffyLine, NodeId, Style,
    TaffyAuto, TaffyTree, auto, fr, length,
};
use taffy::style::{Direction as TaffyDirection, Overflow};

use super::axis::{Axis, CrossAxisAlign, LayoutDirection, MainAxisJustify};
use super::dock::{DockNode, DockState, PanelId, SplitFraction};
use super::geometry::{EdgeInsets, LayoutPos, LayoutPx, LayoutRect, LayoutSize};
use super::primitives::{DockPanel, GridChild, Layout};
use super::retained::{RetainedLayout, ScrollOffset};
use super::scroll::{ScrollAxes, clamp_scroll};
use super::track::{FlexWeight, GridLine, GridLineRef, GridTrack, TrackName, TrackSize};
use crate::theme::Spacing;
use crate::widget_id::WidgetId;

#[derive(Debug, thiserror::Error)]
pub enum LayoutError {
    #[error("dock host references panel {0} that is not provided")]
    MissingDockPanel(PanelId),
    #[error("dock tab node has no panels")]
    EmptyDockTabs,
    #[error("grid span references unknown track name {0}")]
    UnknownTrackName(String),
    #[error("grid declares track name {0} more than once")]
    DuplicateTrackName(String),
    #[error("grid column span resolves to non-increasing lines")]
    GridColumnSpanNotIncreasing,
    #[error("grid row span resolves to non-increasing lines")]
    GridRowSpanNotIncreasing,
    #[error("dock host does not yet render floating surfaces")]
    UnsupportedFloatingSurfaces,
    #[error("taffy: {0}")]
    Taffy(#[from] taffy::TaffyError),
}

#[derive(Clone, Debug, PartialEq)]
pub enum NodeKind {
    Pass,
    Leaf(WidgetId),
    ScrollRegion {
        id: WidgetId,
        axes: ScrollAxes,
        offset: ScrollOffset,
    },
    Splitter {
        id: WidgetId,
        axis: Axis,
        fraction: SplitFraction,
    },
    DockHost {
        id: WidgetId,
    },
    DockSplit {
        axis: Axis,
        fraction: SplitFraction,
    },
    DockTabStrip {
        active: PanelId,
        tabs: Vec<PanelId>,
    },
    DockPanel {
        id: PanelId,
    },
}

#[derive(Clone, Debug)]
pub struct SolvedNode {
    pub kind: NodeKind,
    pub rect: LayoutRect,
    pub content_size: LayoutSize,
    pub children: Vec<SolvedIndex>,
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub struct SolvedIndex(usize);

#[derive(Clone, Debug)]
pub struct SolvedLayout {
    pub nodes: Vec<SolvedNode>,
    pub root: SolvedIndex,
}

impl SolvedLayout {
    #[must_use]
    pub fn root_node(&self) -> &SolvedNode {
        &self.nodes[self.root.0]
    }

    #[must_use]
    pub fn node(&self, index: SolvedIndex) -> &SolvedNode {
        &self.nodes[index.0]
    }
}

pub fn measure(
    layout: &Layout,
    available: LayoutSize,
    retained: &RetainedLayout,
    direction: LayoutDirection,
) -> Result<SolvedLayout, LayoutError> {
    let mut engine = Engine::new(direction);
    let root_taffy = engine.lower(layout, retained)?;
    let mut root_style = engine.style_of(root_taffy)?;
    root_style.size = TaffySize {
        width: length(available.width.value()),
        height: length(available.height.value()),
    };
    engine.tree.set_style(root_taffy, root_style)?;
    engine.tree.compute_layout(
        root_taffy,
        TaffySize {
            width: AvailableSpace::Definite(available.width.value()),
            height: AvailableSpace::Definite(available.height.value()),
        },
    )?;
    let root = engine.collect(root_taffy, LayoutPos::ORIGIN)?;
    Ok(SolvedLayout {
        nodes: engine.solved,
        root,
    })
}

struct Engine {
    tree: TaffyTree<NodeKind>,
    solved: Vec<SolvedNode>,
    direction: LayoutDirection,
}

impl Engine {
    fn new(direction: LayoutDirection) -> Self {
        Self {
            tree: TaffyTree::new(),
            solved: Vec::new(),
            direction,
        }
    }

    fn style_of(&self, node: NodeId) -> Result<Style, LayoutError> {
        Ok(self.tree.style(node)?.clone())
    }

    fn apply_split_grow(
        &mut self,
        a: NodeId,
        b: NodeId,
        fraction: SplitFraction,
    ) -> Result<(), LayoutError> {
        let pa = fraction.value();
        let pb = 1.0 - pa;
        let mut sa = self.style_of(a)?;
        let mut sb = self.style_of(b)?;
        sa.flex_grow = pa;
        sa.flex_shrink = 0.0;
        sa.flex_basis = length(0.0);
        sb.flex_grow = pb;
        sb.flex_shrink = 0.0;
        sb.flex_basis = length(0.0);
        self.tree.set_style(a, sa)?;
        self.tree.set_style(b, sb)?;
        Ok(())
    }
}

const fn flex_direction_for(axis: Axis) -> FlexDirection {
    match axis {
        Axis::Horizontal => FlexDirection::Row,
        Axis::Vertical => FlexDirection::Column,
    }
}

const fn taffy_direction(direction: LayoutDirection) -> TaffyDirection {
    match direction {
        LayoutDirection::Ltr => TaffyDirection::Ltr,
        LayoutDirection::Rtl => TaffyDirection::Rtl,
    }
}

fn fill_along_main(style: Style) -> Style {
    Style {
        flex_grow: 1.0,
        flex_shrink: 1.0,
        flex_basis: length(0.0),
        ..style
    }
}

fn flex_stretch_style(axis: Axis, direction: LayoutDirection) -> Style {
    fill_along_main(Style {
        display: Display::Flex,
        flex_direction: flex_direction_for(axis),
        align_items: Some(AlignItems::Stretch),
        direction: taffy_direction(direction),
        ..Style::DEFAULT
    })
}

impl Engine {
    fn lower(&mut self, layout: &Layout, retained: &RetainedLayout) -> Result<NodeId, LayoutError> {
        match layout {
            Layout::Row {
                gap,
                justify,
                cross,
                children,
            } => self.flex(Axis::Horizontal, *gap, *justify, *cross, children, retained),
            Layout::Column {
                gap,
                justify,
                cross,
                children,
            } => self.flex(Axis::Vertical, *gap, *justify, *cross, children, retained),
            Layout::Stack { children } => self.stack(children, retained),
            Layout::Grid {
                columns,
                rows,
                column_gap,
                row_gap,
                children,
            } => self.grid(columns, rows, *column_gap, *row_gap, children, retained),
            Layout::Inset { padding, child } => self.inset(*padding, child, retained),
            Layout::Spacer { weight } => self.spacer(*weight),
            Layout::Gap { size } => self.gap(*size),
            Layout::ScrollRegion { id, axes, child } => {
                self.scroll_region(*id, *axes, child, retained)
            }
            Layout::Splitter {
                id,
                axis,
                default_fraction,
                a,
                b,
            } => self.splitter(*id, *axis, *default_fraction, a, b, retained),
            Layout::DockHost {
                id,
                state,
                panels,
                tab_strip_height,
            } => self.dock_host(*id, state.as_ref(), panels, *tab_strip_height, retained),
            Layout::Leaf { id } => {
                self.leaf_node(NodeKind::Leaf(*id), fill_along_main(Style::DEFAULT))
            }
        }
    }

    fn leaf_node(&mut self, kind: NodeKind, style: Style) -> Result<NodeId, LayoutError> {
        Ok(self.tree.new_leaf_with_context(style, kind)?)
    }

    fn parent(
        &mut self,
        kind: NodeKind,
        style: Style,
        children: &[NodeId],
    ) -> Result<NodeId, LayoutError> {
        let id = self.tree.new_with_children(style, children)?;
        self.tree.set_node_context(id, Some(kind))?;
        Ok(id)
    }

    fn flex(
        &mut self,
        axis: Axis,
        gap: Spacing,
        justify: MainAxisJustify,
        cross: CrossAxisAlign,
        children: &[Layout],
        retained: &RetainedLayout,
    ) -> Result<NodeId, LayoutError> {
        let kids = children
            .iter()
            .map(|c| self.lower(c, retained))
            .collect::<Result<Vec<_>, _>>()?;
        let style = fill_along_main(Style {
            display: Display::Flex,
            flex_direction: flex_direction_for(axis),
            gap: TaffySize {
                width: length(gap.value_px()),
                height: length(gap.value_px()),
            },
            justify_content: Some(match justify {
                MainAxisJustify::Start => JustifyContent::FlexStart,
                MainAxisJustify::Center => JustifyContent::Center,
                MainAxisJustify::End => JustifyContent::FlexEnd,
                MainAxisJustify::SpaceBetween => JustifyContent::SpaceBetween,
                MainAxisJustify::SpaceAround => JustifyContent::SpaceAround,
            }),
            align_items: Some(match cross {
                CrossAxisAlign::Start => AlignItems::FlexStart,
                CrossAxisAlign::Center => AlignItems::Center,
                CrossAxisAlign::End => AlignItems::FlexEnd,
                CrossAxisAlign::Stretch => AlignItems::Stretch,
            }),
            direction: taffy_direction(self.direction),
            ..Style::DEFAULT
        });
        self.parent(NodeKind::Pass, style, &kids)
    }

    fn stack(
        &mut self,
        children: &[Layout],
        retained: &RetainedLayout,
    ) -> Result<NodeId, LayoutError> {
        let kids = children
            .iter()
            .map(|c| self.lower(c, retained))
            .collect::<Result<Vec<_>, _>>()?;
        let cell = TaffyLine {
            start: GridPlacement::Line(1.into()),
            end: GridPlacement::Line(2.into()),
        };
        let style = fill_along_main(Style {
            display: Display::Grid,
            grid_template_columns: vec![fr::<f32, GridTemplateComponent<String>>(1.0)],
            grid_template_rows: vec![fr::<f32, GridTemplateComponent<String>>(1.0)],
            justify_items: Some(AlignItems::Stretch),
            align_items: Some(AlignItems::Stretch),
            direction: taffy_direction(self.direction),
            ..Style::DEFAULT
        });
        kids.iter()
            .try_fold((), |(), kid| -> Result<(), LayoutError> {
                let mut s = self.style_of(*kid)?;
                s.grid_row = cell.clone();
                s.grid_column = cell.clone();
                self.tree.set_style(*kid, s)?;
                Ok(())
            })?;
        self.parent(NodeKind::Pass, style, &kids)
    }

    fn grid(
        &mut self,
        columns: &[GridTrack],
        rows: &[GridTrack],
        column_gap: Spacing,
        row_gap: Spacing,
        children: &[GridChild],
        retained: &RetainedLayout,
    ) -> Result<NodeId, LayoutError> {
        let kids = children
            .iter()
            .map(|gc| {
                let id = self.lower(&gc.child, retained)?;
                let mut style = self.style_of(id)?;
                let col_start = resolve_line(&gc.span.column_start, columns)?;
                let col_end = resolve_line(&gc.span.column_end, columns)?;
                if col_end.get() <= col_start.get() {
                    return Err(LayoutError::GridColumnSpanNotIncreasing);
                }
                let row_start = resolve_line(&gc.span.row_start, rows)?;
                let row_end = resolve_line(&gc.span.row_end, rows)?;
                if row_end.get() <= row_start.get() {
                    return Err(LayoutError::GridRowSpanNotIncreasing);
                }
                style.grid_column = TaffyLine {
                    start: GridPlacement::Line(col_start.taffy_index().into()),
                    end: GridPlacement::Line(col_end.taffy_index().into()),
                };
                style.grid_row = TaffyLine {
                    start: GridPlacement::Line(row_start.taffy_index().into()),
                    end: GridPlacement::Line(row_end.taffy_index().into()),
                };
                self.tree.set_style(id, style)?;
                Ok::<_, LayoutError>(id)
            })
            .collect::<Result<Vec<_>, _>>()?;
        let style = fill_along_main(Style {
            display: Display::Grid,
            grid_template_columns: columns.iter().map(track_to_taffy).collect(),
            grid_template_rows: rows.iter().map(track_to_taffy).collect(),
            gap: TaffySize {
                width: length(column_gap.value_px()),
                height: length(row_gap.value_px()),
            },
            direction: taffy_direction(self.direction),
            ..Style::DEFAULT
        });
        self.parent(NodeKind::Pass, style, &kids)
    }

    fn inset(
        &mut self,
        padding: EdgeInsets,
        child: &Layout,
        retained: &RetainedLayout,
    ) -> Result<NodeId, LayoutError> {
        let kid = self.lower(child, retained)?;
        let (left, right) = padding.physical(self.direction);
        let style = Style {
            padding: TaffyRect {
                left: length(left.value_px()),
                right: length(right.value_px()),
                top: length(padding.top.value_px()),
                bottom: length(padding.bottom.value_px()),
            },
            ..flex_stretch_style(Axis::Vertical, self.direction)
        };
        self.parent(NodeKind::Pass, style, &[kid])
    }

    fn spacer(&mut self, weight: FlexWeight) -> Result<NodeId, LayoutError> {
        let style = Style {
            flex_grow: weight.as_f32(),
            flex_shrink: 0.0,
            flex_basis: length(0.0),
            ..Style::DEFAULT
        };
        self.leaf_node(NodeKind::Pass, style)
    }

    fn gap(&mut self, size: Spacing) -> Result<NodeId, LayoutError> {
        let style = Style {
            flex_grow: 0.0,
            flex_shrink: 0.0,
            flex_basis: length(size.value_px()),
            ..Style::DEFAULT
        };
        self.leaf_node(NodeKind::Pass, style)
    }

    fn scroll_region(
        &mut self,
        id: WidgetId,
        axes: ScrollAxes,
        child: &Layout,
        retained: &RetainedLayout,
    ) -> Result<NodeId, LayoutError> {
        let kid = self.lower(child, retained)?;
        let style = Style {
            overflow: TaffyPoint {
                x: scroll_overflow(axes.allows_horizontal()),
                y: scroll_overflow(axes.allows_vertical()),
            },
            ..flex_stretch_style(Axis::Vertical, self.direction)
        };
        let offset = retained.scroll_for(id);
        self.parent(NodeKind::ScrollRegion { id, axes, offset }, style, &[kid])
    }

    fn splitter(
        &mut self,
        id: WidgetId,
        axis: Axis,
        default_fraction: SplitFraction,
        a: &Layout,
        b: &Layout,
        retained: &RetainedLayout,
    ) -> Result<NodeId, LayoutError> {
        let fraction = retained.split_for(id, default_fraction);
        let a_id = self.lower(a, retained)?;
        let b_id = self.lower(b, retained)?;
        self.apply_split_grow(a_id, b_id, fraction)?;
        self.parent(
            NodeKind::Splitter { id, axis, fraction },
            flex_stretch_style(axis, self.direction),
            &[a_id, b_id],
        )
    }

    fn dock_host(
        &mut self,
        id: WidgetId,
        state: &DockState,
        panels: &[DockPanel],
        tab_strip_height: Spacing,
        retained: &RetainedLayout,
    ) -> Result<NodeId, LayoutError> {
        if !state.floating.is_empty() {
            return Err(LayoutError::UnsupportedFloatingSurfaces);
        }
        let main = self.dock_node(&state.main, panels, tab_strip_height, retained)?;
        self.parent(
            NodeKind::DockHost { id },
            flex_stretch_style(Axis::Vertical, self.direction),
            &[main],
        )
    }

    fn dock_node(
        &mut self,
        node: &DockNode,
        panels: &[DockPanel],
        tab_strip_height: Spacing,
        retained: &RetainedLayout,
    ) -> Result<NodeId, LayoutError> {
        match node {
            DockNode::Split {
                axis,
                fraction,
                a,
                b,
            } => {
                let a_id = self.dock_node(a, panels, tab_strip_height, retained)?;
                let b_id = self.dock_node(b, panels, tab_strip_height, retained)?;
                self.apply_split_grow(a_id, b_id, *fraction)?;
                self.parent(
                    NodeKind::DockSplit {
                        axis: *axis,
                        fraction: *fraction,
                    },
                    flex_stretch_style(*axis, self.direction),
                    &[a_id, b_id],
                )
            }
            DockNode::Tabs { tabs, active } => {
                let active_id = tabs
                    .get(usize::from(active.get()))
                    .copied()
                    .or_else(|| tabs.first().copied())
                    .ok_or(LayoutError::EmptyDockTabs)?;
                let panel = panels
                    .iter()
                    .find(|p| p.id == active_id)
                    .ok_or(LayoutError::MissingDockPanel(active_id))?;
                let body = self.lower(&panel.child, retained)?;
                let body_panel = self.parent(
                    NodeKind::DockPanel { id: active_id },
                    flex_stretch_style(Axis::Vertical, self.direction),
                    &[body],
                )?;
                if tabs.len() <= 1 {
                    return Ok(body_panel);
                }
                let strip = self.leaf_node(
                    NodeKind::DockTabStrip {
                        active: active_id,
                        tabs: tabs.clone(),
                    },
                    Style {
                        size: TaffySize {
                            width: Dimension::AUTO,
                            height: length(tab_strip_height.value_px()),
                        },
                        ..Style::DEFAULT
                    },
                )?;
                self.parent(
                    NodeKind::Pass,
                    flex_stretch_style(Axis::Vertical, self.direction),
                    &[strip, body_panel],
                )
            }
        }
    }

    fn collect(
        &mut self,
        node: NodeId,
        parent_origin: LayoutPos,
    ) -> Result<SolvedIndex, LayoutError> {
        let layout: TaffyLayout = *self.tree.layout(node)?;
        let origin = LayoutPos::new(
            LayoutPx::saturating(parent_origin.x.value() + layout.location.x),
            LayoutPx::saturating(parent_origin.y.value() + layout.location.y),
        );
        let size = LayoutSize::new(
            LayoutPx::saturating_nonneg(layout.size.width),
            LayoutPx::saturating_nonneg(layout.size.height),
        );
        let rect = LayoutRect::new(origin, size);
        let content_size = LayoutSize::new(
            LayoutPx::saturating_nonneg(layout.content_size.width.max(layout.size.width)),
            LayoutPx::saturating_nonneg(layout.content_size.height.max(layout.size.height)),
        );
        let kind = match self
            .tree
            .get_node_context(node)
            .cloned()
            .unwrap_or(NodeKind::Pass)
        {
            NodeKind::ScrollRegion { id, axes, offset } => NodeKind::ScrollRegion {
                id,
                axes,
                offset: clamp_scroll(offset, size, content_size, axes),
            },
            other => other,
        };
        let child_ids: Vec<NodeId> = self.tree.children(node)?;
        let children = child_ids
            .into_iter()
            .map(|c| self.collect(c, origin))
            .collect::<Result<Vec<_>, _>>()?;
        let solved = SolvedNode {
            kind,
            rect,
            content_size,
            children,
        };
        let idx = SolvedIndex(self.solved.len());
        self.solved.push(solved);
        Ok(idx)
    }
}

fn track_to_taffy(track: &GridTrack) -> GridTemplateComponent<String> {
    match track.size {
        TrackSize::Fixed(spacing) => length(spacing.value_px()),
        TrackSize::Auto => auto(),
        TrackSize::Flex(weight) => fr(weight.as_f32()),
    }
}

const fn scroll_overflow(allowed: bool) -> Overflow {
    if allowed {
        Overflow::Scroll
    } else {
        Overflow::Hidden
    }
}

fn resolve_line(line: &GridLineRef, tracks: &[GridTrack]) -> Result<GridLine, LayoutError> {
    match line {
        GridLineRef::Line(l) => Ok(*l),
        GridLineRef::Name(name) => track_line(name, tracks),
    }
}

fn track_line(name: &TrackName, tracks: &[GridTrack]) -> Result<GridLine, LayoutError> {
    let mut iter = tracks
        .iter()
        .enumerate()
        .filter(|(_, t)| t.name.as_ref() == Some(name));
    let (position, _) = iter
        .next()
        .ok_or_else(|| LayoutError::UnknownTrackName(name.as_str().to_owned()))?;
    if iter.next().is_some() {
        return Err(LayoutError::DuplicateTrackName(name.as_str().to_owned()));
    }
    u16::try_from(position)
        .ok()
        .and_then(|n| n.checked_add(1))
        .and_then(NonZeroU16::new)
        .map(GridLine::new)
        .ok_or_else(|| LayoutError::UnknownTrackName(name.as_str().to_owned()))
}
