//! Iterative measure/arrange/paint over derived visual data. Every occurrence
//! has its own geometry; measurement caches never become semantic object state.
use crate::{
    document::{Frame, MountGeometry, RenderCommand},
    text::{self, TextStyle},
    ui_bridge::{MAX_DIMENSION, UiKind, UiNode},
};
use anyhow::{Context as _, Result, bail};
use std::collections::{HashMap, HashSet};

const MAX_LAYOUT_VISITS: usize = 10_000;

#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum MainAlign {
    #[default]
    Start,
    Center,
    End,
    SpaceBetween,
}
impl MainAlign {
    pub fn name(self) -> &'static str {
        match self {
            Self::Start => "start",
            Self::Center => "center",
            Self::End => "end",
            Self::SpaceBetween => "space-between",
        }
    }
    pub fn parse(value: &str) -> Result<Self, endor_core::VmError> {
        match value {
            "start" => Ok(Self::Start),
            "center" => Ok(Self::Center),
            "end" => Ok(Self::End),
            "space-between" => Ok(Self::SpaceBetween),
            _ => Err(endor_core::VmError::InvalidArgument(
                "mainAlign expects start, center, end or space-between",
            )),
        }
    }
}

/// Main-axis constraints interpreted independently by each parent occurrence.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct FlexPolicy {
    pub basis: Option<f32>,
    pub min: f32,
    pub max: f32,
    pub shrink: f32,
}
impl Default for FlexPolicy {
    fn default() -> Self {
        Self {
            basis: None,
            min: 0.0,
            max: MAX_DIMENSION as f32,
            shrink: 0.0,
        }
    }
}

#[derive(Default, Debug, Clone, Copy, PartialEq)]
pub struct Size {
    pub width: f32,
    pub height: f32,
}

#[derive(Debug, Clone, PartialEq)]
enum Shape {
    Text(String, u32, f32),
    Rect(f32, f32),
    Container(bool, f32, f32, Vec<Size>),
}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
struct Key {
    parent: Option<u64>,
    slot: u64,
    block: String,
}

#[derive(Debug, Clone, Default)]
pub struct LayoutRuntime {
    sizes: HashMap<String, (Shape, Size)>,
    ids: HashMap<Key, u64>,
    next_id: u64,
    pub version: u64,
    pub measured: usize,
    pub arranged: usize,
    pub settle_passes: usize,
    pub slices: usize,
    pub viewport: Option<Size>,
}

struct Item<'a> {
    node: &'a UiNode,
    parent: Option<usize>,
    children: Vec<usize>,
    id: u64,
    constraints: [f32; 4],
}

fn bounded(value: f32) -> Result<f32> {
    if !value.is_finite() || !(0.0..=MAX_DIMENSION as f32).contains(&value) {
        bail!("UI layout exceeds the 1000000 coordinate/size limit");
    }
    Ok(value)
}

pub fn root_constraints(viewport: Option<Size>) -> [f32; 4] {
    let size = viewport.unwrap_or(Size {
        width: MAX_DIMENSION as f32,
        height: MAX_DIMENSION as f32,
    });
    [0.0, 0.0, size.width, size.height]
}
pub fn child_constraints(node: &UiNode, parent: [f32; 4]) -> [f32; 4] {
    if let UiKind::Scroll { width, height, .. } = node.kind {
        return arranged_child_constraints(
            node,
            parent,
            constrained(Size { width, height }, parent),
        );
    }
    match container(node) {
        Some((true, padding, _)) => [
            0.0,
            0.0,
            MAX_DIMENSION as f32,
            (parent[3] - padding * 2.0).max(0.0),
        ],
        Some((false, padding, _)) => [
            0.0,
            0.0,
            (parent[2] - padding * 2.0).max(0.0),
            MAX_DIMENSION as f32,
        ],
        None => parent,
    }
}
/// Scroll cross-axis constraints come from this occurrence's actual viewport,
/// including a source parent's allocation rather than only constructor size.
pub fn arranged_child_constraints(node: &UiNode, available: [f32; 4], bounds: Size) -> [f32; 4] {
    if let UiKind::Scroll { padding, axis, .. } = node.kind {
        return [
            0.0,
            0.0,
            if axis.horizontal() {
                MAX_DIMENSION as f32
            } else {
                (bounds.width - padding * 2.0).max(0.0)
            },
            if axis.vertical() {
                MAX_DIMENSION as f32
            } else {
                (bounds.height - padding * 2.0).max(0.0)
            },
        ];
    }
    child_constraints(node, available)
}

pub fn constrained(size: Size, constraints: [f32; 4]) -> Size {
    Size {
        width: size.width.clamp(constraints[0], constraints[2]),
        height: size.height.clamp(constraints[1], constraints[3]),
    }
}

fn container(node: &UiNode) -> Option<(bool, f32, f32)> {
    match &node.kind {
        UiKind::Row { padding, gap, .. } => Some((true, *padding, *gap)),
        UiKind::Scroll { padding, gap, .. } => Some((false, *padding, *gap)),
        UiKind::Column { padding, gap, .. } | UiKind::Button { padding, gap, .. } => {
            Some((false, *padding, *gap))
        }
        _ => None,
    }
}

pub fn main_axis(node: &UiNode) -> Option<usize> {
    container(node).map(|(row, _, _)| usize::from(!row))
}

fn measure_shape(shape: &Shape) -> Result<Size> {
    let size = match shape {
        Shape::Text(value, size, width) => {
            let style = TextStyle::ui(*size);
            let intrinsic = text::measure(value, style);
            let wrapped = text::measure_wrapped(value, style, *width);
            Size {
                width: intrinsic.width.min(*width),
                height: wrapped.height,
            }
        }
        Shape::Rect(width, height) => Size {
            width: *width,
            height: *height,
        },
        Shape::Container(row, padding, gap, children) => {
            let along = children
                .iter()
                .map(|s| if *row { s.width } else { s.height })
                .sum::<f32>()
                + padding * 2.0
                + gap * children.len().saturating_sub(1) as f32;
            let cross = children
                .iter()
                .map(|s| if *row { s.height } else { s.width })
                .fold(0.0_f32, f32::max)
                + padding * 2.0;
            if *row {
                Size {
                    width: along,
                    height: cross,
                }
            } else {
                Size {
                    width: cross,
                    height: along,
                }
            }
        }
    };
    bounded(size.width)?;
    bounded(size.height)?;
    Ok(size)
}

fn shape(node: &UiNode, children: Vec<Size>, constraints: [f32; 4]) -> Result<Shape> {
    Ok(match &node.kind {
        UiKind::Text { value, size, .. } => Shape::Text(value.clone(), *size, constraints[2]),
        UiKind::Rect { width, height, .. } => Shape::Rect(*width, *height),
        UiKind::Spacer(size) => Shape::Rect(*size, *size),
        UiKind::Scroll { width, height, .. } => Shape::Rect(*width, *height),
        _ => {
            let (row, padding, gap) = container(node).context("unknown native visual protocol")?;
            Shape::Container(row, padding, gap, children)
        }
    })
}

pub fn measure_native(node: &UiNode, children: Vec<Size>, constraints: [f32; 4]) -> Result<Size> {
    measure_shape(&shape(node, children, constraints)?)
}

/// Private derived scratch space for one occurrence during layout. Contains no
/// managed references; source code cannot mutate it through heap attributes.
#[derive(Debug, Default)]
pub struct Context {
    pub node: Option<std::sync::Arc<UiNode>>,
    pub children: Vec<Size>,
    pub preferred: Vec<Size>,
    pub bounds: Option<[f32; 4]>,
}

pub fn paint_native(node: &UiNode, x: f32, y: f32, size: Size) -> Option<RenderCommand> {
    let size_for_text = size;
    match &node.kind {
        UiKind::Text { value, size, color } => Some(RenderCommand::Text {
            x,
            y,
            value: value.clone(),
            size: *size,
            color: color.clone(),
            wrap_bounds: Some([size_for_text.width, size_for_text.height]),
        }),
        UiKind::Rect { radius, color, .. } => Some(RenderCommand::RoundedRect {
            x,
            y,
            width: size.width,
            height: size.height,
            radius: *radius,
            color: color.clone(),
        }),
        UiKind::Button { .. } => Some(RenderCommand::RoundedRect {
            x,
            y,
            width: size.width,
            height: size.height,
            radius: 7.0,
            color: "#E3EAF7".into(),
        }),
        _ => None,
    }
}

/// Allocate main-axis space with bounded growth and basis-weighted shrinking.
/// Preferred sizes come from this occurrence's first measurement in the solve;
/// remeasurement must not turn a prior allocation into a new intrinsic basis.
pub fn flex_sizes(node: &UiNode, bounds: Size, children: &[Size], preferred: &[Size]) -> Vec<Size> {
    let Some((row, padding, gap)) = container(node) else {
        return children.to_vec();
    };
    let nodes = node.children();
    if !nodes.iter().any(|child| child.has_flex()) {
        return children.to_vec();
    }
    let space = ((if row { bounds.width } else { bounds.height }) as f64
        - padding as f64 * 2.0
        - gap as f64 * children.len().saturating_sub(1) as f64)
        .max(0.0);
    let bases: Vec<f64> = nodes
        .iter()
        .zip(preferred)
        .map(|(child, size)| {
            child.flex_policy.basis.unwrap_or_else(|| {
                if child.flex > 0.0 {
                    0.0
                } else if row {
                    size.width
                } else {
                    size.height
                }
            }) as f64
        })
        .collect();
    let mut sizes: Vec<f64> = nodes
        .iter()
        .zip(&bases)
        .map(|(child, basis)| {
            basis.clamp(child.flex_policy.min as f64, child.flex_policy.max as f64)
        })
        .collect();
    let free = space - sizes.iter().sum::<f64>();
    let growing = free > 0.0;
    let weights: Vec<f64> = nodes
        .iter()
        .zip(&bases)
        .map(|(child, basis)| {
            if growing {
                child.flex as f64
            } else {
                child.flex_policy.shrink as f64 * basis
            }
        })
        .collect();
    let limits: Vec<f64> = nodes
        .iter()
        .map(|child| {
            if growing {
                child.flex_policy.max as f64
            } else {
                child.flex_policy.min as f64
            }
        })
        .collect();
    let mut active: Vec<usize> = (0..nodes.len())
        .filter(|&i| {
            weights[i] > 0.0
                && if growing {
                    sizes[i] < limits[i]
                } else {
                    sizes[i] > limits[i]
                }
        })
        .collect();
    let mut remaining = free.abs();
    while remaining > 0.0 && !active.is_empty() {
        let total = active.iter().map(|&i| weights[i]).sum::<f64>();
        let mut spent = 0.0;
        let mut saturated = false;
        for &i in &active {
            let capacity = (limits[i] - sizes[i]).abs();
            let amount = (remaining * weights[i] / total).min(capacity);
            sizes[i] += if growing { amount } else { -amount };
            spent += amount;
            saturated |= amount >= capacity;
        }
        remaining = (remaining - spent).max(0.0);
        // Without saturation this round distributed all available space. Do not
        // spin on floating point residue; a saturated round freezes a child.
        if !saturated {
            break;
        }
        active.retain(|&i| {
            if growing {
                sizes[i] < limits[i]
            } else {
                sizes[i] > limits[i]
            }
        });
    }
    let mut output = children.to_vec();
    for (size, along) in output.iter_mut().zip(sizes) {
        if row {
            size.width = along as f32;
        } else {
            size.height = along as f32;
        }
    }
    output
}

pub fn arrange_native(
    node: &UiNode,
    x: f32,
    y: f32,
    bounds: Size,
    children: &[Size],
) -> Result<Vec<(f32, f32)>> {
    let Some((row, padding, gap)) = container(node) else {
        return Ok(vec![]);
    };
    let occupied = children
        .iter()
        .map(|child| if row { child.width } else { child.height })
        .sum::<f32>()
        + gap * children.len().saturating_sub(1) as f32;
    let available = if row { bounds.width } else { bounds.height };
    let spare = (available - padding * 2.0 - occupied).max(0.0);
    let offset = match node.main_align {
        MainAlign::Center => spare / 2.0,
        MainAlign::End => spare,
        _ => 0.0,
    };
    let distributed_gap = if node.main_align == MainAlign::SpaceBetween && children.len() > 1 {
        spare / (children.len() - 1) as f32
    } else {
        0.0
    };
    let mut cursor = padding + offset;
    let mut points = Vec::with_capacity(children.len());
    for child in children {
        let point = if row {
            (x + cursor, y + padding)
        } else {
            (x + padding, y + cursor)
        };
        bounded(point.0)?;
        bounded(point.1)?;
        bounded(point.0 + child.width)?;
        bounded(point.1 + child.height)?;
        points.push(point);
        cursor += if row { child.width } else { child.height };
        cursor += gap + distributed_gap;
    }
    Ok(points)
}

impl LayoutRuntime {
    pub fn invalidate(&mut self, block: &str) {
        self.sizes.remove(block);
    }
    pub fn reserve_identity_floor(&mut self, floor: u64) {
        self.next_id = self.next_id.max(floor);
    }
    pub fn identity_high_watermark(&self) -> u64 {
        self.next_id
    }

    pub fn clear_occurrences(&mut self) {
        self.ids.clear();
    }
    pub fn frame(&mut self, root: &UiNode) -> Result<Frame> {
        // Stage all cache/identity changes; a rejected frame leaves the last
        // valid occurrence registry intact and never publishes partial geometry.
        let mut next = self.clone();
        let frame = next.build(root, false)?;
        *self = next;
        Ok(frame)
    }

    /// Allocate stable occurrence identities without measuring or publishing geometry.
    pub fn structure(&mut self, root: &UiNode) -> Result<Frame> {
        let mut next = self.clone();
        let frame = next.build(root, true)?;
        *self = next;
        Ok(frame)
    }

    pub fn measure_node(
        &mut self,
        node: &UiNode,
        children: Vec<Size>,
        constraints: [f32; 4],
    ) -> Result<Size> {
        let shape = shape(node, children, constraints)?;
        if let Some((_, size)) = self
            .sizes
            .get(&node.block_id)
            .filter(|(old, _)| *old == shape)
        {
            return Ok(*size);
        }
        self.measured += 1;
        let size = measure_shape(&shape)?;
        self.sizes.insert(node.block_id.clone(), (shape, size));
        Ok(size)
    }

    fn build(&mut self, root: &UiNode, structural: bool) -> Result<Frame> {
        self.measured = 0;
        self.arranged = 0;
        self.settle_passes = usize::from(!structural);
        self.slices = usize::from(!structural);
        let mut items = Vec::<Item<'_>>::new();
        let mut work = vec![(root, None::<usize>, 0, root_constraints(self.viewport))];
        let mut used = HashSet::new();
        while let Some((node, parent, slot, constraints)) = work.pop() {
            if items.len() >= MAX_LAYOUT_VISITS {
                bail!("expanded UI layout exceeds 10000 nodes");
            }
            if node.slots.len() != node.children().len() {
                bail!("invalid child slot table");
            }
            let key = Key {
                parent: parent.map(|i| items[i].id),
                slot,
                block: node.block_id.clone(),
            };
            let id = if let Some(id) = self.ids.get(&key) {
                *id
            } else {
                self.next_id = self
                    .next_id
                    .checked_add(1)
                    .context("Mount identity exhausted")?;
                self.ids.insert(key.clone(), self.next_id);
                self.next_id
            };
            if !used.insert(key) {
                bail!("duplicate child slot in one presentation");
            }
            let index = items.len();
            if let Some(parent) = parent {
                items[parent].children.push(index);
            }
            items.push(Item {
                node,
                parent,
                children: vec![],
                id,
                constraints,
            });
            for (child, slot) in node.children().iter().zip(&node.slots).rev() {
                work.push((
                    child.as_ref(),
                    Some(index),
                    *slot,
                    child_constraints(node, constraints),
                ));
            }
        }
        let mut sizes = vec![Size::default(); items.len()];
        if !structural {
            for (index, item) in items.iter().enumerate().rev() {
                sizes[index] = constrained(
                    self.measure_node(
                        item.node,
                        item.children.iter().map(|i| sizes[*i]).collect(),
                        item.constraints,
                    )?,
                    item.constraints,
                );
            }
        }
        let mut points = vec![(0.0_f32, 0.0_f32); items.len()];
        let mut mounts = Vec::with_capacity(items.len());
        let mut commands = Vec::new();
        let mut siblings = vec![None; items.len()];
        for item in &items {
            for pair in item.children.windows(2) {
                siblings[pair[0]] = Some(items[pair[1]].id);
            }
        }
        for (index, item) in items.iter().enumerate() {
            self.arranged += usize::from(!structural);
            let (x, y) = points[index];
            let size = sizes[index];
            bounded(x)?;
            bounded(y)?;
            bounded(x + size.width)?;
            bounded(y + size.height)?;
            mounts.push(MountGeometry {
                id: item.id,
                clip: None,
                scroll: None,
                block_id: item.node.block_id.clone(),
                parent: item.parent.map(|i| items[i].id),
                first_child: item.children.first().map(|i| items[*i].id),
                next_sibling: siblings[index],
                x,
                y,
                width: size.width,
                height: size.height,
                interactive: item.node.interactive,
            });
            if !structural && let Some(command) = paint_native(item.node, x, y, size) {
                commands.push(command);
            }
            if structural {
                continue;
            }
            let child_sizes: Vec<_> = item.children.iter().map(|child| sizes[*child]).collect();
            let child_sizes = flex_sizes(item.node, size, &child_sizes, &child_sizes);
            for (child, allocated) in item.children.iter().zip(&child_sizes) {
                sizes[*child] = *allocated;
            }

            for (child, point) in
                item.children
                    .iter()
                    .zip(arrange_native(item.node, x, y, size, &child_sizes)?)
            {
                points[*child] = point;
            }
        }
        self.ids.retain(|key, _| used.contains(key));
        let blocks: HashSet<_> = items
            .iter()
            .map(|item| item.node.block_id.as_str())
            .collect();
        self.sizes
            .retain(|block, _| blocks.contains(block.as_str()));
        self.version = self
            .version
            .checked_add(1)
            .context("layout version exhausted")?;
        Ok(Frame {
            commands,
            mounts,
            layout_version: self.version,
            ..Frame::default()
        })
    }
}

pub fn frame(root: &UiNode) -> Result<Frame> {
    LayoutRuntime::default().frame(root)
}

pub fn frame_with_viewport(root: &UiNode, viewport: Size) -> Result<Frame> {
    LayoutRuntime {
        viewport: Some(viewport),
        ..Default::default()
    }
    .frame(root)
}