use std::collections::HashMap;

use nalgebra_glm::Vec2;
use taffy::prelude::{FromFlex, FromLength, FromPercent, TaffyAuto};
use taffy::{
    AvailableSpace, Dimension, Display, FlexDirection, FlexWrap, JustifyContent, LengthPercentage,
    NodeId, Size, Style, TaffyTree, TrackSizingFunction,
};

use crate::ecs::text::resources::FontEngine;
use crate::ecs::ui::components::{AutoSizeMode, UiLayoutNode};
use crate::ecs::ui::layout_types::{FlowAlignment, FlowDirection, FlowLayout, GridLayout};
use crate::ecs::ui::units::UiValue;

#[derive(Clone, Copy, Debug, Default)]
pub struct TaffyContext {
    pub measured: Size<f32>,
}

#[derive(Clone, Debug)]
pub struct WrapTextMeasure {
    pub text: String,
    pub font_size: f32,
}

pub struct UiTaffy {
    pub tree: TaffyTree<TaffyContext>,
    pub entity_to_node: HashMap<freecs::Entity, NodeId>,
    pub node_to_entity: HashMap<NodeId, freecs::Entity>,
    pub measure_cache: HashMap<freecs::Entity, Size<f32>>,
    pub wrap_text: HashMap<freecs::Entity, WrapTextMeasure>,
}

impl Default for UiTaffy {
    fn default() -> Self {
        Self {
            tree: TaffyTree::new(),
            entity_to_node: HashMap::new(),
            node_to_entity: HashMap::new(),
            measure_cache: HashMap::new(),
            wrap_text: HashMap::new(),
        }
    }
}

impl std::fmt::Debug for UiTaffy {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("UiTaffy")
            .field("entity_count", &self.entity_to_node.len())
            .finish()
    }
}

impl UiTaffy {
    pub fn clear(&mut self) {
        self.tree = TaffyTree::new();
        self.entity_to_node.clear();
        self.node_to_entity.clear();
        self.measure_cache.clear();
        self.wrap_text.clear();
    }

    pub fn release(&mut self, entity: freecs::Entity) {
        if let Some(node_id) = self.entity_to_node.remove(&entity) {
            self.node_to_entity.remove(&node_id);
            let _ = self.tree.remove(node_id);
        }
        self.measure_cache.remove(&entity);
        self.wrap_text.remove(&entity);
    }

    pub fn ensure_leaf(&mut self, entity: freecs::Entity, style: Style) -> NodeId {
        if let Some(&node_id) = self.entity_to_node.get(&entity) {
            let _ = self.tree.set_style(node_id, style);
            return node_id;
        }
        let node_id = self
            .tree
            .new_leaf_with_context(style, TaffyContext::default())
            .expect("taffy new_leaf");
        self.entity_to_node.insert(entity, node_id);
        self.node_to_entity.insert(node_id, entity);
        node_id
    }

    pub fn set_children(&mut self, parent: freecs::Entity, children: &[NodeId]) {
        if let Some(&parent_node) = self.entity_to_node.get(&parent) {
            let _ = self.tree.set_children(parent_node, children);
        }
    }

    pub fn set_measured(&mut self, entity: freecs::Entity, size: Size<f32>) {
        self.measure_cache.insert(entity, size);
    }

    pub fn set_wrap_text(&mut self, entity: freecs::Entity, info: WrapTextMeasure) {
        self.wrap_text.insert(entity, info);
    }

    pub fn clear_wrap_text(&mut self, entity: freecs::Entity) {
        self.wrap_text.remove(&entity);
    }

    pub fn compute_layout(
        &mut self,
        root: freecs::Entity,
        available: Size<AvailableSpace>,
        font_engine: &mut FontEngine,
    ) {
        let Some(&root_node) = self.entity_to_node.get(&root) else {
            return;
        };
        let Self {
            tree,
            node_to_entity,
            measure_cache,
            wrap_text,
            ..
        } = self;
        let _ = tree.compute_layout_with_measure(
            root_node,
            available,
            |known, available_in, node_id, _: Option<&mut TaffyContext>, _style| {
                if let (Some(w), Some(h)) = (known.width, known.height) {
                    return Size {
                        width: w,
                        height: h,
                    };
                }
                let entity = match node_to_entity.get(&node_id) {
                    Some(entity) => *entity,
                    None => return Size::ZERO,
                };
                if let Some(wrap) = wrap_text.get(&entity) {
                    let wrap_width = known.width.or(match available_in.width {
                        AvailableSpace::Definite(value) => Some(value),
                        _ => None,
                    });
                    if let Some(width) = wrap_width
                        && width > 0.0
                    {
                        let wrapped = crate::ecs::ui::text_wrapping::wrap_text(
                            font_engine,
                            &wrap.text,
                            wrap.font_size,
                            width,
                        );
                        let lines = wrapped.lines().count().max(1) as f32;
                        let line_height = wrap.font_size * 1.2;
                        return Size {
                            width,
                            height: line_height * lines,
                        };
                    }
                }
                let cached = measure_cache.get(&entity).copied().unwrap_or(Size::ZERO);
                Size {
                    width: known.width.unwrap_or(cached.width),
                    height: known.height.unwrap_or(cached.height),
                }
            },
        );
    }

    pub fn rect_for(&self, entity: freecs::Entity) -> Option<(Vec2, Vec2)> {
        let node_id = *self.entity_to_node.get(&entity)?;
        let layout = self.tree.layout(node_id).ok()?;
        let position = Vec2::new(layout.location.x, layout.location.y);
        let size = Vec2::new(layout.size.width, layout.size.height);
        Some((position, size))
    }
}

fn mix_dim(absolute: f32, relative_pct: f32, parent: f32) -> Dimension {
    if relative_pct.abs() < f32::EPSILON {
        return Dimension::from_length(absolute);
    }
    if absolute.abs() < f32::EPSILON {
        return Dimension::from_percent(relative_pct * 0.01);
    }
    Dimension::from_length(absolute + relative_pct * 0.01 * parent)
}

fn vec_dimension(
    value: &UiValue<Vec2>,
    parent: Vec2,
    dpi_scale: f32,
    auto: AutoSizeMode,
) -> Size<Dimension> {
    let absolute = value.absolute.unwrap_or(Vec2::new(0.0, 0.0)) * dpi_scale;
    let relative_pct = value.relative.unwrap_or(Vec2::new(0.0, 0.0));

    let width = match auto {
        AutoSizeMode::Width | AutoSizeMode::Both => Dimension::AUTO,
        _ => mix_dim(absolute.x, relative_pct.x, parent.x),
    };
    let height = match auto {
        AutoSizeMode::Height | AutoSizeMode::Both => Dimension::AUTO,
        _ => mix_dim(absolute.y, relative_pct.y, parent.y),
    };
    Size { width, height }
}

fn apply_flow_container(style: &mut Style, flow: &FlowLayout, dpi_scale: f32) {
    let padding_units = LengthPercentage::from_length(flow.padding * dpi_scale);
    let gap = LengthPercentage::from_length(flow.spacing * dpi_scale);

    style.display = Display::Flex;
    style.flex_direction = match flow.direction {
        FlowDirection::Vertical => FlexDirection::Column,
        FlowDirection::Horizontal => FlexDirection::Row,
    };
    style.flex_wrap = if flow.wrap {
        FlexWrap::Wrap
    } else {
        FlexWrap::NoWrap
    };
    style.gap = Size {
        width: gap,
        height: gap,
    };
    style.padding = taffy::Rect {
        left: padding_units,
        right: padding_units,
        top: padding_units,
        bottom: padding_units,
    };
    style.justify_content = Some(match flow.alignment {
        FlowAlignment::Start => JustifyContent::FlexStart,
        FlowAlignment::Center => JustifyContent::Center,
        FlowAlignment::End => JustifyContent::FlexEnd,
    });
    style.align_items = Some(match flow.cross_alignment {
        FlowAlignment::Start => taffy::AlignItems::FlexStart,
        FlowAlignment::Center => taffy::AlignItems::Center,
        FlowAlignment::End => taffy::AlignItems::FlexEnd,
    });
}

fn apply_grid_container(style: &mut Style, grid: &GridLayout, parent_width: f32, dpi_scale: f32) {
    let padding_units = LengthPercentage::from_length(grid.padding * dpi_scale);
    let column_gap = LengthPercentage::from_length(grid.column_spacing * dpi_scale);
    let row_gap = LengthPercentage::from_length(grid.row_spacing * dpi_scale);
    let scaled_padding = grid.padding * dpi_scale;
    let scaled_col_spacing = grid.column_spacing * dpi_scale;

    let columns = if let Some(min_width) = grid.min_column_width {
        let scaled_min = min_width * dpi_scale;
        let inner_width = (parent_width - scaled_padding * 2.0).max(0.0);
        if scaled_min > 0.0 {
            ((inner_width + scaled_col_spacing) / (scaled_min + scaled_col_spacing))
                .floor()
                .max(1.0) as u16
        } else {
            grid.columns.max(1) as u16
        }
    } else {
        grid.columns.max(1) as u16
    };

    style.display = Display::Grid;
    style.grid_template_columns = (0..columns)
        .map(|_| {
            TrackSizingFunction::from(taffy::MinMax {
                min: taffy::MinTrackSizingFunction::from_length(0.0),
                max: taffy::MaxTrackSizingFunction::from_flex(1.0),
            })
        })
        .collect();
    style.grid_auto_rows = vec![taffy::NonRepeatedTrackSizingFunction::from_length(
        grid.row_height * dpi_scale,
    )];
    style.gap = Size {
        width: column_gap,
        height: row_gap,
    };
    style.padding = taffy::Rect {
        left: padding_units,
        right: padding_units,
        top: padding_units,
        bottom: padding_units,
    };
}

pub fn flow_to_style(
    node: &UiLayoutNode,
    flow: &FlowLayout,
    parent_size: Vec2,
    dpi_scale: f32,
) -> Style {
    let mut style = Style {
        size: Size {
            width: Dimension::from_length(parent_size.x),
            height: Dimension::from_length(parent_size.y),
        },
        ..Default::default()
    };
    apply_flow_container(&mut style, flow, dpi_scale);
    apply_size_constraints(&mut style, node, dpi_scale);
    style
}

pub fn grid_to_style(
    node: &UiLayoutNode,
    grid: &GridLayout,
    parent_size: Vec2,
    dpi_scale: f32,
) -> Style {
    let mut style = Style {
        size: Size {
            width: Dimension::from_length(parent_size.x),
            height: Dimension::from_length(parent_size.y),
        },
        ..Default::default()
    };
    apply_grid_container(&mut style, grid, parent_size.x, dpi_scale);
    apply_size_constraints(&mut style, node, dpi_scale);
    style
}

pub fn child_to_style(node: &UiLayoutNode, parent_size: Vec2, dpi_scale: f32) -> Style {
    let is_container = node.flow_layout.is_some() || node.grid_layout.is_some();

    let mut style = if let Some(child_size) = &node.flow_child_size {
        let mut size = vec_dimension(child_size, parent_size, dpi_scale, node.auto_size);
        if is_container {
            if let Dimension::Length(value) = size.width
                && value.abs() < f32::EPSILON
            {
                size.width = Dimension::AUTO;
            }
            if let Dimension::Length(value) = size.height
                && value.abs() < f32::EPSILON
            {
                size.height = Dimension::AUTO;
            }
        }
        Style {
            size,
            ..Default::default()
        }
    } else if matches!(node.auto_size, AutoSizeMode::Both) || is_container {
        Style {
            size: Size {
                width: Dimension::AUTO,
                height: Dimension::AUTO,
            },
            ..Default::default()
        }
    } else {
        Style::default()
    };

    if let Some(grow) = node.flex_grow {
        style.flex_grow = grow;
    }
    style.flex_shrink = node.flex_shrink.unwrap_or(0.0);

    apply_size_constraints(&mut style, node, dpi_scale);

    if let Some(flow) = &node.flow_layout {
        apply_flow_container(&mut style, flow, dpi_scale);
    } else if let Some(grid) = &node.grid_layout {
        apply_grid_container(&mut style, grid, parent_size.x, dpi_scale);
    }

    if !node.visible {
        style.display = Display::None;
    }

    if node.base_layout.is_some() {
        style.position = taffy::Position::Absolute;
    }

    style
}

fn apply_size_constraints(style: &mut Style, node: &UiLayoutNode, dpi_scale: f32) {
    if let Some(min) = node.min_size {
        style.min_size = Size {
            width: Dimension::from_length(min.x * dpi_scale),
            height: Dimension::from_length(min.y * dpi_scale),
        };
    }
    if let Some(max) = node.max_size {
        style.max_size = Size {
            width: Dimension::from_length(max.x * dpi_scale),
            height: Dimension::from_length(max.y * dpi_scale),
        };
    }
}

pub fn available_size_for(parent: Vec2) -> Size<AvailableSpace> {
    Size {
        width: AvailableSpace::Definite(parent.x),
        height: AvailableSpace::Definite(parent.y),
    }
}
