(** Pure flexbox layout engine.

    Two-pass algorithm:
    1. Measure (bottom-up): compute intrinsic sizes from text extents
       and children accumulation.
    2. Layout (top-down): allocate space starting from window bounds,
       distributing via flex-grow and handling alignment/gap. *)

open Sx_native_types

(* -- Text measurement -- *)

let measure_text (cr : Cairo.context) (family : [`Sans | `Mono]) (weight : [`Normal | `Bold])
    (slant : [`Normal | `Italic]) (size : float) (text : string) : float * float =
  let font_name = match family with `Sans -> "sans-serif" | `Mono -> "monospace" in
  let cairo_weight = match weight with `Normal -> Cairo.Normal | `Bold -> Cairo.Bold in
  let cairo_slant = match slant with `Normal -> Cairo.Upright | `Italic -> Cairo.Italic in
  Cairo.select_font_face cr ~slant:cairo_slant ~weight:cairo_weight font_name;
  Cairo.set_font_size cr size;
  let fe = Cairo.font_extents cr in
  if String.length text = 0 then (0., fe.ascent +. fe.descent)
  else begin
    (* Word wrap not needed for POC -- measure as single line *)
    let te = Cairo.text_extents cr text in
    (te.Cairo.width +. te.Cairo.x_bearing, fe.ascent +. fe.descent)
  end

(* -- Measure pass (bottom-up) -- *)

(** Set intrinsic [box.w] and [box.h] on each node based on text extents
    and children accumulation. Does NOT set x/y. *)
let rec measure (cr : Cairo.context) (node : node) : unit =
  (* Measure children first *)
  List.iter (measure cr) node.children;

  let pad = node.style.padding in
  let pad_h = pad.left +. pad.right in
  let pad_v = pad.top +. pad.bottom in

  match node.text with
  | Some txt ->
    (* Leaf text node: measure the text *)
    let (tw, th) = measure_text cr node.style.font_family node.style.font_weight
        node.style.font_style node.style.font_size txt in
    node.box.w <- tw +. pad_h;
    node.box.h <- th +. pad_v
  | None ->
    if node.style.display = `None then begin
      node.box.w <- 0.;
      node.box.h <- 0.
    end else begin
      let visible_children = List.filter (fun c -> c.style.display <> `None) node.children in
      let n_children = List.length visible_children in
      let total_gap = if n_children > 1 then node.style.gap *. float_of_int (n_children - 1) else 0. in
      match node.style.flex_direction with
      | `Column ->
        (* Stack vertically: width = max child width, height = sum of child heights + gaps *)
        let max_w = List.fold_left (fun acc c ->
          let cm = c.style.margin in
          Float.max acc (c.box.w +. cm.left +. cm.right)
        ) 0. visible_children in
        let sum_h = List.fold_left (fun acc c ->
          let cm = c.style.margin in
          acc +. c.box.h +. cm.top +. cm.bottom
        ) 0. visible_children in
        node.box.w <- max_w +. pad_h;
        node.box.h <- sum_h +. total_gap +. pad_v
      | `Row ->
        (* Stack horizontally: height = max child height, width = sum of child widths + gaps *)
        let sum_w = List.fold_left (fun acc c ->
          let cm = c.style.margin in
          acc +. c.box.w +. cm.left +. cm.right
        ) 0. visible_children in
        let max_h = List.fold_left (fun acc c ->
          let cm = c.style.margin in
          Float.max acc (c.box.h +. cm.top +. cm.bottom)
        ) 0. visible_children in
        node.box.w <- sum_w +. total_gap +. pad_h;
        node.box.h <- max_h +. pad_v
    end;

    (* Apply explicit width/height constraints *)
    (match node.style.width with
     | `Px w -> node.box.w <- w
     | `Full | `Auto -> ());
    (match node.style.height with
     | `Px h -> node.box.h <- h
     | `Full | `Auto -> ())


(* -- Layout pass (top-down) -- *)

(** Position all nodes within the given bounds [x, y, w, h].
    Distributes space according to flex-grow and handles alignment. *)
let rec layout (node : node) (x : float) (y : float) (avail_w : float) (avail_h : float) : unit =
  let margin = node.style.margin in
  let x = x +. margin.left in
  let y = y +. margin.top in
  let avail_w = avail_w -. margin.left -. margin.right in
  let avail_h = avail_h -. margin.top -. margin.bottom in

  node.box.x <- x;
  node.box.y <- y;

  (* Determine actual width/height.
     Container nodes with Auto width stretch to fill available space
     (like CSS block-level elements), while text nodes keep intrinsic width. *)
  let is_text_node = node.text <> None in
  let w = match node.style.width with
    | `Full -> avail_w
    | `Px pw -> Float.min pw avail_w
    | `Auto ->
      if is_text_node then Float.min node.box.w avail_w
      else avail_w  (* containers expand to fill *)
  in
  let h = match node.style.height with
    | `Full -> avail_h
    | `Px ph -> Float.min ph avail_h
    | `Auto -> node.box.h  (* Use intrinsic height *)
  in

  node.box.w <- w;
  node.box.h <- h;

  if node.style.display = `None then ()
  else begin
    let pad = node.style.padding in
    let inner_x = x +. pad.left in
    let inner_y = y +. pad.top in
    let inner_w = w -. pad.left -. pad.right in
    let inner_h = h -. pad.top -. pad.bottom in

    let visible_children = List.filter (fun c -> c.style.display <> `None) node.children in

    match visible_children with
    | [] -> ()  (* Leaf or empty container *)
    | children ->
      let n_children = List.length children in
      let total_gap = if n_children > 1 then node.style.gap *. float_of_int (n_children - 1) else 0. in

      match node.style.flex_direction with
      | `Column ->
        (* Calculate total intrinsic height and flex-grow sum *)
        let total_intrinsic = List.fold_left (fun acc c ->
          let cm = c.style.margin in
          acc +. c.box.h +. cm.top +. cm.bottom
        ) 0. children in
        let total_grow = List.fold_left (fun acc c -> acc +. c.style.flex_grow) 0. children in
        let remaining = Float.max 0. (inner_h -. total_intrinsic -. total_gap) in

        (* justify-content: space-between *)
        let (start_offset, between_extra) = match node.style.justify_content with
          | `Between when n_children > 1 ->
            (0., remaining /. float_of_int (n_children - 1))
          | `Center -> (remaining /. 2., 0.)
          | `End -> (remaining, 0.)
          | _ -> (0., 0.)
        in

        let cur_y = ref (inner_y +. start_offset) in
        List.iter (fun child ->
          let cm = child.style.margin in
          let child_w = match child.style.width with
            | `Full -> inner_w -. cm.left -. cm.right
            | _ -> Float.min child.box.w (inner_w -. cm.left -. cm.right)
          in
          let extra_h = if total_grow > 0. then remaining *. child.style.flex_grow /. total_grow else 0. in
          let child_h = child.box.h +. extra_h in

          (* align-items: cross-axis alignment *)
          let child_x = match node.style.align_items with
            | `Center -> inner_x +. (inner_w -. child_w -. cm.left -. cm.right) /. 2.
            | `End -> inner_x +. inner_w -. child_w -. cm.right
            | `Stretch ->
              (* Stretch: child takes full width *)
              layout child (inner_x) !cur_y (inner_w) child_h;
              cur_y := !cur_y +. child.box.h +. cm.top +. cm.bottom +. node.style.gap +. between_extra;
              (* skip the normal layout below *)
              inner_x  (* dummy, won't be used *)
            | _ -> inner_x
          in

          if node.style.align_items <> `Stretch then begin
            layout child child_x !cur_y child_w child_h;
            cur_y := !cur_y +. child.box.h +. cm.top +. cm.bottom +. node.style.gap +. between_extra
          end
        ) children

      | `Row ->
        (* Calculate total intrinsic width and flex-grow sum *)
        let total_intrinsic = List.fold_left (fun acc c ->
          let cm = c.style.margin in
          acc +. c.box.w +. cm.left +. cm.right
        ) 0. children in
        let total_grow = List.fold_left (fun acc c -> acc +. c.style.flex_grow) 0. children in
        let remaining = Float.max 0. (inner_w -. total_intrinsic -. total_gap) in

        let (start_offset, between_extra) = match node.style.justify_content with
          | `Between when n_children > 1 ->
            (0., remaining /. float_of_int (n_children - 1))
          | `Center -> (remaining /. 2., 0.)
          | `End -> (remaining, 0.)
          | _ -> (0., 0.)
        in

        let cur_x = ref (inner_x +. start_offset) in
        List.iter (fun child ->
          let cm = child.style.margin in
          let extra_w = if total_grow > 0. then remaining *. child.style.flex_grow /. total_grow else 0. in
          let child_w = child.box.w +. extra_w in
          let child_h = match child.style.height with
            | `Full -> inner_h -. cm.top -. cm.bottom
            | _ -> Float.min child.box.h (inner_h -. cm.top -. cm.bottom)
          in

          (* align-items: cross-axis alignment *)
          let child_y = match node.style.align_items with
            | `Center -> inner_y +. (inner_h -. child_h -. cm.top -. cm.bottom) /. 2.
            | `End -> inner_y +. inner_h -. child_h -. cm.bottom
            | `Stretch ->
              layout child !cur_x inner_y child_w inner_h;
              cur_x := !cur_x +. child.box.w +. cm.left +. cm.right +. node.style.gap +. between_extra;
              inner_y  (* dummy *)
            | _ -> inner_y
          in

          if node.style.align_items <> `Stretch then begin
            layout child !cur_x child_y child_w child_h;
            cur_x := !cur_x +. child.box.w +. cm.left +. cm.right +. node.style.gap +. between_extra
          end
        ) children
  end