"""Top-down layout pass — assign (x, y, w, h) positions to every node."""

from __future__ import annotations

from uframe.ir import (
    IRNode, Page, Box, Row, Col, Spacer, Pad, Rect,
    Heading, Text, Label, Divider, Link, ListNode, ListItem,
    Gauge, Sparkline, Status, Table,
    Form, Field, Password, Radio, Checkbox, FormButton,
    Let, Source, IfBlock, ForLoop, CacheControl, OnSubmit, StateDecl,
    BigTitle, ImageNode, HNav, VNav,
)


def layout(node: IRNode, x: int, y: int, w: int, h: int) -> int:
    """Assign positions to a node and its children.

    Args:
        node: the IR node to lay out
        x, y: top-left position in the grid
        w: available width
        h: available height (advisory, may grow)

    Returns:
        The actual height consumed by this node.
    """
    node.rect = Rect(x=x, y=y, w=w, h=0)

    if isinstance(node, Page):
        cursor_y = y
        for child in node.children:
            child_h = layout(child, x, cursor_y, w, h - (cursor_y - y))
            cursor_y += child_h
        node.rect.h = cursor_y - y
        return node.rect.h

    elif isinstance(node, Box):
        # Border takes 1 char on each side
        inner_x = x + 1
        inner_y = y + 1
        inner_w = w - 2
        cursor_y = inner_y
        for child in node.children:
            child_h = layout(child, inner_x, cursor_y, inner_w,
                             h - 2 - (cursor_y - inner_y))
            cursor_y += child_h
        inner_h = cursor_y - inner_y
        node.rect.h = inner_h + 2  # +2 for top/bottom border
        return node.rect.h

    elif isinstance(node, Row):
        n = len(node.children)
        if n == 0:
            return 0

        gap_total = node.gap * (n - 1)
        usable = w - gap_total

        # Distribute width
        widths: list[int] = []
        fixed_total = 0
        flex_count = 0
        for child in node.children:
            if isinstance(child, Col) and child.col_width is not None:
                widths.append(child.col_width)
                fixed_total += child.col_width
            else:
                widths.append(0)
                flex_count += 1

        flex_each = max(1, (usable - fixed_total) // flex_count) if flex_count > 0 else 0
        remainder = (usable - fixed_total) - (flex_each * flex_count) if flex_count > 0 else 0

        for i, child in enumerate(node.children):
            if widths[i] == 0:
                widths[i] = flex_each
                if remainder > 0:
                    widths[i] += 1
                    remainder -= 1

        # Lay out each child at its column position
        max_h = 0
        col_x = x
        for i, child in enumerate(node.children):
            child_h = layout(child, col_x, y, widths[i], h)
            max_h = max(max_h, child_h)
            col_x += widths[i] + node.gap

        node.rect.h = max_h
        return max_h

    elif isinstance(node, Col):
        cursor_y = y
        col_w = node.col_width if node.col_width is not None else w
        col_w = min(col_w, w)
        for child in node.children:
            child_h = layout(child, x, cursor_y, col_w, h - (cursor_y - y))
            cursor_y += child_h
        node.rect.w = col_w
        node.rect.h = cursor_y - y
        return node.rect.h

    elif isinstance(node, Spacer):
        node.rect.h = node.lines
        return node.lines

    elif isinstance(node, Pad):
        cursor_y = y + node.top
        inner_w = w - node.left - node.right
        for child in node.children:
            child_h = layout(child, x + node.left, cursor_y, inner_w,
                             h - node.top - node.bottom - (cursor_y - y - node.top))
            cursor_y += child_h
        node.rect.h = (cursor_y - y) + node.bottom
        return node.rect.h

    elif isinstance(node, Form):
        cursor_y = y
        for child in node.children:
            child_h = layout(child, x, cursor_y, w, h - (cursor_y - y))
            cursor_y += child_h
        node.rect.h = cursor_y - y
        return node.rect.h

    elif isinstance(node, (Let, Source, CacheControl, StateDecl)):
        node.rect.h = 0
        return 0

    elif isinstance(node, (IfBlock, ForLoop, OnSubmit)):
        cursor_y = y
        for child in node.children:
            child_h = layout(child, x, cursor_y, w, h - (cursor_y - y))
            cursor_y += child_h
        node.rect.h = cursor_y - y
        return node.rect.h

    elif isinstance(node, (Heading, Text, Label, Divider, Link, ListItem,
                           Gauge, Sparkline, Status, Table,
                           Field, Password, Radio, Checkbox, FormButton,
                           BigTitle, ImageNode, HNav, VNav)):
        node.rect.h = node.pref_height
        return node.pref_height

    elif isinstance(node, ListNode):
        cursor_y = y
        for child in node.children:
            child_h = layout(child, x + 2, cursor_y, w - 2, h - (cursor_y - y))
            cursor_y += child_h
        node.rect.h = cursor_y - y
        return node.rect.h

    else:
        # Generic vertical stacking
        cursor_y = y
        for child in node.children:
            child_h = layout(child, x, cursor_y, w, h - (cursor_y - y))
            cursor_y += child_h
        node.rect.h = cursor_y - y
        return node.rect.h
