Source file server_reason_react_ppx.ml

open Ppxlib
open Ast_builder.Default
module List = ListLabels

type target = Native | Js

(* Since ppxlib doesn't provide a way to get the submodules, we need to keep track of them manually *)
let mode = ref Native
let shared_folder_prefix = ref None
let repo_url = "https://github.com/ml-in-barcelona/server-reason-react"
let issues_url = Printf.sprintf "%s/issues" repo_url

let match_substring string substring =
  try
    Str.search_forward (Str.regexp_string substring) string 0 |> ignore;
    true
  with Not_found -> false

(* There's no Ppxlib.pexp_list since isn't a parsetree constructor *)
let pexp_list ~loc xs =
  List.fold_left (List.rev xs) ~init:[%expr []] ~f:(fun xs x ->
      let loc = x.pexp_loc in
      [%expr [%e x] :: [%e xs]])

exception Error of expression

let raise_errorf ~loc fmt =
  Printf.ksprintf
    (fun msg ->
      let expr = pexp_extension ~loc (Location.error_extensionf ~loc "%s" msg) in
      raise (Error expr))
    fmt

let longident ~loc txt = { txt = Lident txt; loc }
let ident ~loc txt = pexp_ident ~loc (longident ~loc txt)
let make_string ~loc str = Ast_helper.Exp.constant ~loc (Ast_helper.Const.string str)
let react_dot_component = "react.component"
let react_dot_async_dot_component = "react.async.component"
let react_dot_client_dot_component = "react.client.component"
let react_dot_server_dot_function = "react.server.function"
let hasAttr { attr_name; _ } comparable = attr_name.txt = comparable

let hasAnyReactComponentAttribute { attr_name; _ } =
  attr_name.txt = react_dot_component
  || attr_name.txt = react_dot_async_dot_component
  || attr_name.txt = react_dot_client_dot_component

let nonReactAttributes { attr_name; _ } =
  attr_name.txt <> react_dot_component
  && attr_name.txt <> react_dot_async_dot_component
  && attr_name.txt <> react_dot_client_dot_component

let hasAttrOnBinding { pvb_attributes } comparable =
  List.find_opt ~f:(fun attr -> hasAttr attr comparable) pvb_attributes <> None

let isReactComponentBinding vb = hasAttrOnBinding vb react_dot_component
let isReactAsyncComponentBinding vb = hasAttrOnBinding vb react_dot_async_dot_component
let isReactClientComponentBinding vb = hasAttrOnBinding vb react_dot_client_dot_component
let isReactServerFunctionBinding vb = hasAttrOnBinding vb react_dot_server_dot_function

let isClientComponentBinding value_bindings =
  let first_binding = List.hd value_bindings in
  isReactClientComponentBinding first_binding

let contains_client_component structure =
  List.exists
    ~f:(fun structure_item ->
      match structure_item.pstr_desc with
      | Pstr_value (_, value_bindings) -> List.exists ~f:isReactClientComponentBinding value_bindings
      | _ -> false)
    structure

let rec unwrap_children children = function
  | { pexp_desc = Pexp_construct ({ txt = Lident "[]"; _ }, None); _ } -> List.rev children
  | { pexp_desc = Pexp_construct ({ txt = Lident "::"; _ }, Some { pexp_desc = Pexp_tuple [ child; next ]; _ }); _ } ->
      unwrap_children (child :: children) next
  | e -> raise_errorf ~loc:e.pexp_loc "jsx: children prop should be a list"

let is_jsx = function { attr_name = { txt = "JSX"; _ }; _ } -> true | _ -> false
let has_jsx_attr attrs = List.exists ~f:is_jsx attrs

let strip_unit_args args =
  List.filter args ~f:(fun (label, expr) ->
      match (label, expr.pexp_desc) with Nolabel, Pexp_construct ({ txt = Lident "()"; _ }, None) -> false | _ -> true)

let component_make_props_ident tag =
  match tag with
  | { txt = Lident name; loc } -> { txt = Lident (name ^ "Props"); loc }
  | { txt = Ldot (path, name); loc } -> { txt = Ldot (path, name ^ "Props"); loc }
  | { txt = Lapply _; loc } -> raise_errorf ~loc "jsx: component props can't be created from functor applications"

let is_key_arg (label, _) = match label with Optional "key" | Labelled "key" -> true | _ -> false

let rewrite_component ~loc tag args children =
  let component = pexp_ident ~loc tag in
  let props_args =
    match children with
    | None -> args
    | Some [ children ] -> (Labelled "children", children) :: args
    | Some children -> (Labelled "children", [%expr React.list [%e pexp_list ~loc children]]) :: args
  in
  let key_args, non_key_args = List.partition ~f:is_key_arg props_args in
  let make_props = pexp_ident ~loc (component_make_props_ident tag) in
  let non_key_args = strip_unit_args non_key_args in
  let props = pexp_apply ~loc make_props (non_key_args @ [ (Nolabel, [%expr ()]) ]) in
  let make_args =
    match key_args with [] -> [ (Nolabel, props) ] | (label, expr) :: _ -> [ (label, expr); (Nolabel, props) ]
  in
  pexp_apply ~loc component make_args

let validate_prop ~loc id name =
  match DomProps.findByJsxName ~tag:id name with
  | Ok p -> p
  | Error `ElementNotFound ->
      raise_errorf ~loc "jsx: HTML tag '%s' doesn't exist.\nIf this isn't correct, please open an issue at %s" id
        issues_url
  | Error `AttributeNotFound -> (
      match DomProps.findClosestName name with
      | None ->
          raise_errorf ~loc
            "jsx: prop '%s' isn't valid on a '%s' element.\nIf this isn't correct, please open an issue at %s." name id
            issues_url
      | Some suggestion ->
          raise_errorf ~loc
            "jsx: prop '%s' isn't valid on a '%s' element.\n\
             Hint: Maybe you mean '%s'?\n\n\
             If this isn't correct, please open an issue at %s."
            name id suggestion issues_url)

let make_prop ~is_optional ~prop attribute_value =
  let loc = attribute_value.pexp_loc in
  let open DomProps in
  match (prop, is_optional) with
  | Attribute { type_ = DomProps.Action; name; jsxName }, false ->
      [%expr
        match ([%e attribute_value] : [ `String of string | `Function of 'a Runtime.server_function ]) with
        | `String s -> Some (React.JSX.String ([%e estring ~loc name], [%e estring ~loc jsxName], (s : string)))
        | `Function f ->
            Some
              (React.JSX.Action ([%e estring ~loc name], [%e estring ~loc jsxName], (f : 'a Runtime.server_function)))]
  | Attribute { type_ = DomProps.Action; name; jsxName }, true ->
      [%expr
        match ([%e attribute_value] : [ `String of string | `Function of 'a Runtime.server_function ] option) with
        | None -> None
        | Some v -> Some (React.JSX.Action ([%e estring ~loc name], [%e estring ~loc jsxName], v))]
  | Attribute { type_ = DomProps.String; name; jsxName }, false ->
      [%expr
        Some (React.JSX.String ([%e estring ~loc name], [%e estring ~loc jsxName], ([%e attribute_value] : string)))]
  | Attribute { type_ = DomProps.String; name; jsxName }, true ->
      [%expr
        match ([%e attribute_value] : string option) with
        | None -> None
        | Some v -> Some (React.JSX.String ([%e estring ~loc name], [%e estring ~loc jsxName], v))]
  | Attribute { type_ = DomProps.Int; name; jsxName }, false ->
      [%expr Some (React.JSX.Int ([%e estring ~loc name], [%e estring ~loc jsxName], ([%e attribute_value] : int)))]
  | Attribute { type_ = DomProps.Int; name; jsxName }, true ->
      [%expr
        match ([%e attribute_value] : int option) with
        | None -> None
        | Some v -> Some (React.JSX.Int ([%e estring ~loc name], [%e estring ~loc jsxName], v))]
  | Attribute { type_ = DomProps.Float; name; jsxName }, false ->
      [%expr Some (React.JSX.Float ([%e estring ~loc name], [%e estring ~loc jsxName], ([%e attribute_value] : float)))]
  | Attribute { type_ = DomProps.Float; name; jsxName }, true ->
      [%expr
        match ([%e attribute_value] : float option) with
        | None -> None
        | Some v -> Some (React.JSX.Float ([%e estring ~loc name], [%e estring ~loc jsxName], v))]
  | Attribute { type_ = DomProps.Bool; name; jsxName }, false ->
      [%expr Some (React.JSX.Bool ([%e estring ~loc name], [%e estring ~loc jsxName], ([%e attribute_value] : bool)))]
  | Attribute { type_ = DomProps.Bool; name; jsxName }, true ->
      [%expr
        match ([%e attribute_value] : bool option) with
        | None -> None
        | Some v -> Some (React.JSX.Bool ([%e estring ~loc name], [%e estring ~loc jsxName], v))]
  (* BooleanishString stays a boolean until the serialization seams: HTML renders "true"/"false", Flight keeps the raw
     JSON boolean *)
  | Attribute { type_ = DomProps.BooleanishString; name; jsxName }, false ->
      [%expr
        Some
          (React.JSX.BooleanishString ([%e estring ~loc name], [%e estring ~loc jsxName], ([%e attribute_value] : bool)))]
  | Attribute { type_ = DomProps.BooleanishString; name; jsxName }, true ->
      [%expr
        match ([%e attribute_value] : bool option) with
        | None -> None
        | Some v -> Some (React.JSX.BooleanishString ([%e estring ~loc name], [%e estring ~loc jsxName], v))]
  | Attribute { type_ = DomProps.Style; _ }, false ->
      [%expr Some (React.JSX.Style ([%e attribute_value] : ReactDOM.Style.t))]
  | Attribute { type_ = DomProps.Style; _ }, true ->
      [%expr
        match ([%e attribute_value] : ReactDOM.Style.t option) with None -> None | Some v -> Some (React.JSX.Style v)]
  | Attribute { type_ = DomProps.Ref; _ }, false -> [%expr Some (React.JSX.Ref ([%e attribute_value] : React.domRef))]
  | Attribute { type_ = DomProps.Ref; _ }, true ->
      [%expr match ([%e attribute_value] : React.domRef option) with None -> None | Some v -> Some (React.JSX.Ref v)]
  | Attribute { type_ = DomProps.InnerHtml; _ }, false ->
      [%expr Some (React.JSX.dangerouslyInnerHtml [%e attribute_value])]
  | Attribute { type_ = DomProps.InnerHtml; _ }, true ->
      [%expr match [%e attribute_value] with None -> None | Some v -> Some (React.JSX.dangerouslyInnerHtml v)]
  | Event { type_ = Mouse; jsxName }, false ->
      [%expr
        Some
          (React.JSX.Event
             ([%e make_string ~loc jsxName], React.JSX.Mouse ([%e attribute_value] : React.Event.Mouse.t -> unit)))]
  | Event { type_ = Mouse; jsxName }, true ->
      [%expr
        match ([%e attribute_value] : (React.Event.Mouse.t -> unit) option) with
        | None -> None
        | Some v -> Some (React.JSX.Event ([%e make_string ~loc jsxName], React.JSX.Mouse v))]
  | Event { type_ = Selection; jsxName }, false ->
      [%expr
        Some
          (React.JSX.Event
             ([%e make_string ~loc jsxName], React.JSX.Selection ([%e attribute_value] : React.Event.Mouse.t -> unit)))]
  | Event { type_ = Selection; jsxName }, true ->
      [%expr
        match ([%e attribute_value] : (React.Event.Selection.t -> unit) option) with
        | None -> None
        | Some v -> Some (React.JSX.Event ([%e make_string ~loc jsxName], React.JSX.Selection v))]
  | Event { type_ = Touch; jsxName }, false ->
      [%expr
        Some
          (React.JSX.Event
             ([%e make_string ~loc jsxName], React.JSX.Touch ([%e attribute_value] : React.Event.Touch.t -> unit)))]
  | Event { type_ = Touch; jsxName }, true ->
      [%expr
        match ([%e attribute_value] : (React.Event.Touch.t -> unit) option) with
        | None -> None
        | Some v -> Some (React.JSX.Event ([%e make_string ~loc jsxName], React.JSX.Touch v))]
  | Event { type_ = UI; jsxName }, false ->
      [%expr
        Some
          (React.JSX.Event
             ([%e make_string ~loc jsxName], React.JSX.UI ([%e attribute_value] : React.Event.UI.t -> unit)))]
  | Event { type_ = UI; jsxName }, true ->
      [%expr
        match ([%e attribute_value] : (React.Event.UI.t -> unit) option) with
        | None -> None
        | Some v -> Some (React.JSX.Event ([%e make_string ~loc jsxName], React.JSX.UI v))]
  | Event { type_ = Wheel; jsxName }, false ->
      [%expr
        Some
          (React.JSX.Event
             ([%e make_string ~loc jsxName], React.JSX.Wheel ([%e attribute_value] : React.Event.Wheel.t -> unit)))]
  | Event { type_ = Wheel; jsxName }, true ->
      [%expr
        match ([%e attribute_value] : (React.Event.Wheel.t -> unit) option) with
        | None -> None
        | Some v -> Some (React.JSX.Event ([%e make_string ~loc jsxName], React.JSX.Wheel v))]
  | Event { type_ = Clipboard; jsxName }, false ->
      [%expr
        Some
          (React.JSX.Event
             ( [%e make_string ~loc jsxName],
               React.JSX.Clipboard ([%e attribute_value] : React.Event.Clipboard.t -> unit) ))]
  | Event { type_ = Clipboard; jsxName }, true ->
      [%expr
        match ([%e attribute_value] : (React.Event.Clipboard.t -> unit) option) with
        | None -> None
        | Some v -> Some (React.JSX.Event ([%e make_string ~loc jsxName], React.JSX.Clipboard v))]
  | Event { type_ = Composition; jsxName }, false ->
      [%expr
        Some
          (React.JSX.Event
             ( [%e make_string ~loc jsxName],
               React.JSX.Composition ([%e attribute_value] : React.Event.Composition.t -> unit) ))]
  | Event { type_ = Composition; jsxName }, true ->
      [%expr
        match ([%e attribute_value] : (React.Event.Composition.t -> unit) option) with
        | None -> None
        | Some v -> Some (React.JSX.Event ([%e make_string ~loc jsxName], React.JSX.Composition v))]
  | Event { type_ = Keyboard; jsxName }, false ->
      [%expr
        Some
          (React.JSX.Event
             ([%e make_string ~loc jsxName], React.JSX.Keyboard ([%e attribute_value] : React.Event.Keyboard.t -> unit)))]
  | Event { type_ = Keyboard; jsxName }, true ->
      [%expr
        match ([%e attribute_value] : (React.Event.Keyboard.t -> unit) option) with
        | None -> None
        | Some v -> Some (React.JSX.Event ([%e make_string ~loc jsxName], React.JSX.Keyboard v))]
  | Event { type_ = Focus; jsxName }, false ->
      [%expr
        Some
          (React.JSX.Event
             ([%e make_string ~loc jsxName], React.JSX.Focus ([%e attribute_value] : React.Event.Focus.t -> unit)))]
  | Event { type_ = Focus; jsxName }, true ->
      [%expr
        match ([%e attribute_value] : (React.Event.Focus.t -> unit) option) with
        | None -> None
        | Some v -> Some (React.JSX.Event ([%e make_string ~loc jsxName], React.JSX.Focus v))]
  | Event { type_ = Form; jsxName }, false ->
      [%expr
        Some
          (React.JSX.Event
             ([%e make_string ~loc jsxName], React.JSX.Form ([%e attribute_value] : React.Event.Form.t -> unit)))]
  | Event { type_ = Form; jsxName }, true ->
      [%expr
        match ([%e attribute_value] : (React.Event.Form.t -> unit) option) with
        | None -> None
        | Some v -> Some (React.JSX.Event ([%e make_string ~loc jsxName], React.JSX.Form v))]
  | Event { type_ = Media; jsxName }, false ->
      [%expr
        Some
          (React.JSX.Event
             ([%e make_string ~loc jsxName], React.JSX.Media ([%e attribute_value] : React.Event.Media.t -> unit)))]
  | Event { type_ = Media; jsxName }, true ->
      [%expr
        match ([%e attribute_value] : (React.Event.Media.t -> unit) option) with
        | None -> None
        | Some v -> Some (React.JSX.Event ([%e make_string ~loc jsxName], React.JSX.Media v))]
  | Event { type_ = Inline; jsxName }, false ->
      [%expr Some (React.JSX.Event ([%e make_string ~loc jsxName], React.JSX.Inline ([%e attribute_value] : string)))]
  | Event { type_ = Inline; jsxName }, true ->
      [%expr
        match ([%e attribute_value] : string option) with
        | None -> None
        | Some v -> Some (React.JSX.Event ([%e make_string ~loc jsxName], React.JSX.Inline v))]
  | Event { type_ = Image; jsxName }, false ->
      [%expr
        Some
          (React.JSX.Event
             ( [%e make_string ~loc jsxName],
               React.JSX.Image ([%e attribute_value] : (React.Event.Image.t -> unit) option) ))]
  | Event { type_ = Image; jsxName }, true ->
      [%expr
        match ([%e attribute_value] : (React.Event.Image.t -> unit) option) with
        | None -> None
        | Some v -> Some (React.JSX.Event ([%e make_string ~loc jsxName], React.JSX.Image v))]
  | Event { type_ = Animation; jsxName }, false ->
      [%expr
        Some
          (React.JSX.Event
             ( [%e make_string ~loc jsxName],
               React.JSX.Animation ([%e attribute_value] : React.Event.Animation.t -> unit) ))]
  | Event { type_ = Animation; jsxName }, true ->
      [%expr
        match ([%e attribute_value] : (React.Event.Animation.t -> unit) option) with
        | None -> None
        | Some v -> Some (React.JSX.Event ([%e make_string ~loc jsxName], React.JSX.Animation v))]
  | Event { type_ = Transition; jsxName }, false ->
      [%expr
        Some
          (React.JSX.Event
             ( [%e make_string ~loc jsxName],
               React.JSX.Transition ([%e attribute_value] : React.Event.Transition.t -> unit) ))]
  | Event { type_ = Transition; jsxName }, true ->
      [%expr
        match ([%e attribute_value] : (React.Event.Transition.t -> unit) option) with
        | None -> None
        | Some v -> Some (React.JSX.Event ([%e make_string ~loc jsxName], React.JSX.Transition v))]
  | Event { type_ = Pointer; jsxName }, false ->
      [%expr
        Some
          (React.JSX.Event
             ([%e make_string ~loc jsxName], React.JSX.Pointer ([%e attribute_value] : React.Event.Pointer.t -> unit)))]
  | Event { type_ = Pointer; jsxName }, true ->
      [%expr
        match ([%e attribute_value] : (React.Event.Pointer.t -> unit) option) with
        | None -> None
        | Some v -> Some (React.JSX.Event ([%e make_string ~loc jsxName], React.JSX.Pointer v))]
  | Event { type_ = Drag; jsxName }, false ->
      [%expr
        Some
          (React.JSX.Event
             ([%e make_string ~loc jsxName], React.JSX.Drag ([%e attribute_value] : React.Event.Drag.t -> unit)))]
  | Event { type_ = Drag; jsxName }, true ->
      [%expr
        match ([%e attribute_value] : (React.Event.Drag.t -> unit) option) with
        | None -> None
        | Some v -> Some (React.JSX.Event ([%e make_string ~loc jsxName], React.JSX.Drag v))]

let is_optional = function Optional _ -> true | _ -> false
let get_label = function Nolabel -> "" | Optional name | Labelled name -> name

let transform_labelled ~loc ~tag_name (prop_label, (runtime_value : expression)) props =
  match prop_label with
  | Nolabel -> props
  | Optional name | Labelled name ->
      let is_optional = is_optional prop_label in
      let prop = validate_prop ~loc tag_name name in
      let new_prop = make_prop ~is_optional ~prop runtime_value in
      [%expr [%e new_prop] :: [%e props]]

let transform_lowercase_props ~loc ~tag_name args =
  match args with
  | [] -> [%expr []]
  | attrs -> (
      let list_of_attributes = attrs |> List.fold_right ~f:(transform_labelled ~loc ~tag_name) ~init:[%expr []] in
      match list_of_attributes with
      | [%expr []] -> [%expr []]
      | _ ->
          (* We need to filter attributes since optionals are represented as None *)
          [%expr Stdlib.List.filter_map Stdlib.Fun.id [%e list_of_attributes]])

let generate_create_element ~loc ~tag_name ~key ~props ~children =
  let dom_node_name = estring ~loc tag_name in
  match (key, children) with
  | Some key, Some children ->
      let childrens = pexp_list ~loc children in
      [%expr React.createElementWithKey ~key:[%e key] [%e dom_node_name] [%e props] [%e childrens]]
  | None, Some children ->
      let childrens = pexp_list ~loc children in
      [%expr React.createElement [%e dom_node_name] [%e props] [%e childrens]]
  | Some key, None -> [%expr React.createElementWithKey ~key:[%e key] [%e dom_node_name] [%e props] []]
  | None, None -> [%expr React.createElement [%e dom_node_name] [%e props] []]

(* Emit buffer writes that serialize a single attribute value of the given
   kind. Mirrors [ReactDOM.write_attribute_to_buffer] exactly so that PPX-
   lowered output is byte-identical to the variant-tree path.

   [value_expr] is the already-unwrapped (not an option) runtime expression
   that produces the value. Handles only the kinds declared lowerable in
   [Static_analysis.is_lowerable_kind]. Mirrored by
   [Static_analysis.render_static_attr_with_info] on the literal side. *)
let emit_attr_value_write ~loc ~info ~value_expr =
  let open Static_analysis in
  let name_expr = estring ~loc info.html_name in
  (* Wrap a value-writing sub-expression with the ` name="…"` skeleton. *)
  let quoted inner =
    [%expr
      Buffer.add_char __buf ' ';
      Buffer.add_string __buf [%e name_expr];
      Buffer.add_string __buf "=\"";
      [%e inner];
      Buffer.add_char __buf '"']
  in
  match info.kind with
  | DomProps.String -> quoted [%expr ReactDOM.escape_to_buffer __buf ([%e value_expr] : string)]
  | DomProps.Int ->
      (* [string_of_int] allocates a small short-lived string, but measures
         ~15% faster per write than [Printf.bprintf b "%d"], whose
         CamlinternalFormat dispatch outweighs the allocation it saves. *)
      quoted [%expr Buffer.add_string __buf (Stdlib.string_of_int ([%e value_expr] : int))]
  | DomProps.Bool ->
      (* Matches [Bool (name, _, true) -> " " ^ name] / [false -> nothing]. *)
      [%expr
        if ([%e value_expr] : bool) then begin
          Buffer.add_char __buf ' ';
          Buffer.add_string __buf [%e name_expr]
        end]
  | DomProps.BooleanishString ->
      quoted [%expr Buffer.add_string __buf (if ([%e value_expr] : bool) then "true" else "false")]
  | DomProps.Style ->
      (* Mirrors [ReactDOM.write_attribute_to_buffer]'s [Style] case exactly:
         serialize the style, then HTML-escape the result (so a quoted value
         such as a [font-family] cannot break out of the [style="…"] attribute),
         byte-identical output. *)
      [%expr
        Buffer.add_string __buf " style=\"";
        ReactDOM.escape_to_buffer __buf (ReactDOM.Style.to_string ([%e value_expr] : ReactDOM.Style.t));
        Buffer.add_char __buf '"']
  | DomProps.Float | DomProps.Action | DomProps.Ref | DomProps.InnerHtml ->
      (* Unreachable: [is_lowerable_kind] rejects these kinds before we ever
         reach emission. Fail loud at compile time if the invariant breaks. *)
      Location.raise_errorf ~loc "internal PPX error: attribute kind not lowerable but reached emission (name=%s)"
        info.html_name

(* The concrete (non-option) value type carried by a lowerable attribute
   kind. Used to annotate an optional attribute's scrutinee as
   [(expr : t option)] so OCaml's type-directed disambiguation resolves a
   bare [None]/[Some] in [expr] against [option] — even when a user type in
   scope shadows them (e.g. [type roundness = … | None]). This mirrors the
   annotations the variant-tree path emits in [make_prop]; without it the
   [Writer.emit] fast path is the only place that loses the expected type. *)
let attr_value_core_type ~loc (kind : DomProps.attributeType) =
  match kind with
  | DomProps.String -> [%type: string]
  | DomProps.Int -> [%type: int]
  | DomProps.Bool | DomProps.BooleanishString -> [%type: bool]
  | DomProps.Style -> [%type: ReactDOM.Style.t]
  | DomProps.Float | DomProps.Action | DomProps.Ref | DomProps.InnerHtml ->
      (* Unreachable for the same reason as [emit_attr_value_write]. *)
      Location.raise_errorf ~loc "internal PPX error: attribute kind not lowerable but reached emission"

(* Compile-time knowledge of whether the previously emitted node ended
   inside a text run. [Ct_unknown] appears only after a [Dynamic_element]
   hole, whose content is decided at render time. *)
type compile_time_textness = Ct_text | Ct_markup | Ct_unknown

(* Emit the [Buffer.t -> separators:bool -> unit] body for a
   [static_part list]. Writes the static skeleton inline; each dynamic hole
   becomes a [Buffer.add_*] / escape / attribute-emission call that runs at
   render time with no intermediate allocation. Shared by the
   [Needs_string_concat] and [Needs_buffer] tiers, which produce the same
   emit function; the tier names only record what the analysis found.

   Text-separator protocol ([renderToString] parity, required for
   hydration): adjacent text nodes must be delimited by [<!-- -->] when
   [separators] is true. The walk threads a compile-time textness state:
   when both sides of a boundary are known, the separator is emitted behind
   a plain [if separators] check (or not at all); when the left side is a
   [Dynamic_element] hole the generated code threads a [prev_text] ref,
   assigned from [ReactDOM.write_element_to_buffer]'s result. The ref is
   only read while the compile-time state is [Ct_unknown], so static parts
   and string holes never need to update it. *)
let emit_parts_emit_fn ~loc parts =
  let open Static_analysis in
  (* Emit one [Dynamic_attr_slot] hole. For optional slots we peek through
     an obvious [Some e]/[None] to skip the runtime [match]; common when a
     caller writes [?foo=Some x] or [?foo=None] literally. *)
  let write_attr_slot ~info ~expr ~is_optional =
    let loc = expr.pexp_loc in
    if not is_optional then emit_attr_value_write ~loc ~info ~value_expr:expr
    else
      match expr.pexp_desc with
      | Pexp_construct ({ txt = Lident "None"; _ }, None) -> [%expr ()]
      | Pexp_construct ({ txt = Lident "Some"; _ }, Some inner) -> emit_attr_value_write ~loc ~info ~value_expr:inner
      | _ ->
          let write_some = emit_attr_value_write ~loc ~info ~value_expr:[%expr v] in
          let value_ty = attr_value_core_type ~loc info.kind in
          [%expr match ([%e expr] : [%t value_ty] option) with None -> () | Some v -> [%e write_some]]
  in
  (* [Dynamic_attr_slot] lives inside the open tag: it doesn't touch the
     child text run. *)
  let after_part state = function
    | Static_str { ends_text; _ } -> if ends_text then Ct_text else Ct_markup
    | Dynamic_string _ | Dynamic_int _ -> Ct_text
    | Dynamic_element _ -> Ct_unknown
    | Dynamic_attr_slot _ -> state
  in
  let begins_with_text = function
    | Static_str { starts_text; _ } -> starts_text
    | Dynamic_string _ | Dynamic_int _ -> true
    | Dynamic_element _ | Dynamic_attr_slot _ -> false
  in
  (* Prepass: the [prev_text] ref is read when a text-beginning part or a
     [Dynamic_element] hole runs while the compile-time state is
     [Ct_unknown]. When that never happens, skip declaring the ref. *)
  let prev_text_ref_needed =
    let rec loop state = function
      | [] -> false
      | part :: rest ->
          let reads_ref =
            match part with Dynamic_element _ -> state = Ct_unknown | _ -> begins_with_text part && state = Ct_unknown
          in
          reads_ref || loop (after_part state part) rest
    in
    loop Ct_markup parts
  in
  let separator ~state =
    match state with
    | Ct_markup -> None
    | Ct_text -> Some [%expr if __separators then Buffer.add_string __buf "<!-- -->"]
    | Ct_unknown -> Some [%expr if !__prev_text && __separators then Buffer.add_string __buf "<!-- -->"]
  in
  let uses_separators = ref false in
  let writes =
    let rec loop state = function
      | [] -> []
      | part :: rest ->
          let maybe_separator =
            if begins_with_text part then (
              let sep = separator ~state in
              if Option.is_some sep then uses_separators := true;
              sep)
            else None
          in
          let write =
            match part with
            | Static_str { html; _ } when String.length html = 0 ->
                (* Nothing to write: an empty text chunk only participates
                   in the text-run tracking (compile-time state). *)
                None
            | Static_str { html; _ } -> Some [%expr Buffer.add_string __buf [%e estring ~loc html]]
            | Dynamic_string e ->
                let loc = e.pexp_loc in
                Some [%expr ReactDOM.escape_to_buffer __buf [%e e]]
            | Dynamic_int e ->
                let loc = e.pexp_loc in
                Some [%expr Buffer.add_string __buf (Stdlib.string_of_int [%e e])]
            | Dynamic_element e ->
                let loc = e.pexp_loc in
                uses_separators := true;
                let prev_text_arg =
                  match state with
                  | Ct_text -> [%expr true]
                  | Ct_markup -> [%expr false]
                  | Ct_unknown -> [%expr !__prev_text]
                in
                if prev_text_ref_needed then
                  Some
                    [%expr
                      __prev_text :=
                        ReactDOM.write_element_to_buffer __buf ~separators:__separators ~prev_text:[%e prev_text_arg]
                          [%e e]]
                else
                  Some
                    [%expr
                      let (_ : bool) =
                        ReactDOM.write_element_to_buffer __buf ~separators:__separators ~prev_text:[%e prev_text_arg]
                          [%e e]
                      in
                      ()]
            | Dynamic_attr_slot { info; expr; is_optional } -> Some (write_attr_slot ~info ~expr ~is_optional)
          in
          let stmts = List.filter_map ~f:(fun x -> x) [ maybe_separator; write ] in
          stmts @ loop (after_part state part) rest
    in
    loop Ct_markup parts
  in
  let body =
    List.fold_right writes ~init:[%expr ()] ~f:(fun w acc ->
        [%expr
          [%e w];
          [%e acc]])
  in
  let body =
    if prev_text_ref_needed then
      [%expr
        let __prev_text = ref false in
        [%e body]]
    else body
  in
  if !uses_separators then [%expr fun __buf ~separators:__separators -> [%e body]]
  else [%expr fun __buf ~separators:_ -> [%e body]]

let rewrite_lowercase ~loc tag_name args children =
  let key =
    args |> List.find_opt ~f:(fun (label, _) -> get_label label = "key") |> Option.map (fun (_, value) -> value)
  in
  let props = transform_lowercase_props ~loc ~tag_name args in
  match Static_analysis.analyze_element ~tag_name ~attrs:args ~children with
  | Static_analysis.Fully_static html ->
      let html_with_doctype = Static_analysis.maybe_add_doctype tag_name html in
      let html_expr = estring ~loc html_with_doctype in
      let original = generate_create_element ~loc ~tag_name ~key ~props ~children in
      [%expr React.Static { prerendered = [%e html_expr]; original = [%e original] }]
  | Static_analysis.Needs_string_concat parts | Static_analysis.Needs_buffer parts ->
      (* Emit a [Buffer.t -> separators:bool -> unit] that writes directly into the caller's
         buffer. Avoids the N-per-subtree [Buffer.create] + [Buffer.contents]
         that a [Static] wrapping would cost.

         [original] is a thunk that rebuilds the variant-tree on demand for
         [cloneElement] / RSC; zero-alloc unless called. *)
      let parts_with_doctype =
        match tag_name with "html" -> Static_analysis.static_markup "<!DOCTYPE html>" :: parts | _ -> parts
      in
      let emit_fn = emit_parts_emit_fn ~loc parts_with_doctype in
      let original_tree = generate_create_element ~loc ~tag_name ~key ~props ~children in
      let original_thunk = [%expr fun () -> [%e original_tree]] in
      [%expr React.Writer { emit = [%e emit_fn]; original = [%e original_thunk] }]
  | Static_analysis.Cannot_optimize -> generate_create_element ~loc ~tag_name ~key ~props ~children

let split_args args =
  let children = ref (Location.none, []) in
  let rest =
    List.filter_map args ~f:(function
      | Labelled "children", children_expression ->
          let children' = unwrap_children [] children_expression in
          children := (children_expression.pexp_loc, children');
          None
      | arg_label, e -> Some (arg_label, e))
  in
  let children_prop = match !children with _loc, [] -> None | _loc, children -> Some children in
  (children_prop, rest)

let reverse_pexp_list ~loc expr =
  let rec go acc = function
    | [%expr []] -> acc
    | [%expr [%e? hd] :: [%e? tl]] -> go [%expr [%e hd] :: [%e acc]] tl
    | expr -> expr
  in
  go [%expr []] expr

let list_have_tail expr =
  match expr with
  | Pexp_construct ({ txt = Lident "::"; _ }, Some { pexp_desc = Pexp_tuple _; _ })
  | Pexp_construct ({ txt = Lident "[]"; _ }, None) ->
      false
  | _ -> true

let transform_items_of_list ~loc children =
  let rec run_mapper children accum =
    match children with
    | [%expr []] -> reverse_pexp_list ~loc accum
    | [%expr [%e? v] :: [%e? acc]] when list_have_tail acc.pexp_desc -> [%expr [%e v]]
    | [%expr [%e? v] :: [%e? acc]] -> run_mapper acc [%expr [%e v] :: [%e accum]]
    | notAList -> notAList
  in
  run_mapper children [%expr []]

let remove_warning_16_optional_argument_cannot_be_erased ~loc =
  let open Ast_helper in
  {
    attr_name = { txt = "warning"; loc };
    attr_payload = PStr [ Str.eval (Exp.constant (Const.string "-16")) ];
    attr_loc = loc;
  }

let remove_warning_27_unused_var_strict ~loc =
  let open Ast_helper in
  {
    attr_name = { txt = "warning"; loc };
    attr_payload = PStr [ Str.eval (Exp.constant (Const.string "-27")) ];
    attr_loc = loc;
  }

(* Finds the name of the variable the binding is assigned to, otherwise raises *)
let get_function_name binding =
  match binding with
  | { pvb_pat = { ppat_desc = Ppat_var { txt } } } -> txt
  | _ -> raise_errorf ~loc:binding.pvb_loc "react.component calls cannot be destructured."

(* TODO: there are a few unsupported features inside of blocks - Pexp_letmodule , Pexp_letexception , Pexp_ifthenelse *)
let add_unit_at_the_last_argument expression =
  let loc = expression.pexp_loc in
  let has_final_unit params =
    match List.rev params with
    | {
        pparam_desc = Pparam_val (Nolabel, _, { ppat_desc = Ppat_construct ({ txt = Lident "()" }, _) | Ppat_any; _ });
        _;
      }
      :: _ ->
        true
    | _ -> false
  in
  let unit_param = { pparam_loc = loc; pparam_desc = Pparam_val (Nolabel, None, [%pat? ()]) } in
  let rec find_innermost_function_and_add_unit expression =
    match expression.pexp_desc with
    | Pexp_function (params, constraint_, Pfunction_body inner_body) -> (
        match inner_body.pexp_desc with
        | Pexp_function _ ->
            let modified_inner = find_innermost_function_and_add_unit inner_body in
            { expression with pexp_desc = Pexp_function (params, constraint_, Pfunction_body modified_inner) }
        | _ when (not (has_final_unit params)) && params <> [] ->
            {
              expression with
              pexp_attributes = remove_warning_16_optional_argument_cannot_be_erased ~loc :: expression.pexp_attributes;
              pexp_desc = Pexp_function (params @ [ unit_param ], constraint_, Pfunction_body inner_body);
            }
        | _ -> expression)
    | Pexp_function _ -> expression
    | _ -> expression
  in
  let rec inner expression =
    match expression.pexp_desc with
    | Pexp_function _ -> find_innermost_function_and_add_unit expression
    (* let make = {let foo = bar in (~prop) => ...} *)
    | Pexp_let (recursive, vbs, internalExpression) ->
        pexp_let ~loc:expression.pexp_loc recursive vbs (inner internalExpression)
    (* let make = React.forwardRef((~prop) => ...) *)
    | Pexp_apply (_, [ (Nolabel, internalExpression) ]) -> inner internalExpression
    (* let make = React.memoCustomCompareProps((~prop) => ..., (prevPros, nextProps) => true) *)
    | Pexp_apply (_, [ (Nolabel, internalExpression); ((Nolabel, { pexp_desc = Pexp_function _; _ }) as _compareProps) ])
      ->
        inner internalExpression
    | Pexp_sequence (wrapperExpression, internalExpression) ->
        pexp_sequence ~loc:expression.pexp_loc wrapperExpression (inner internalExpression)
    | _ -> expression
  in
  inner expression

let transform_fun_body_expression expr fn =
  let rec find_innermost_body_and_transform expr =
    match expr.pexp_desc with
    | Pexp_function (params, constraint_, Pfunction_body inner_body) -> (
        match inner_body.pexp_desc with
        | Pexp_function _ ->
            let transformed_inner = find_innermost_body_and_transform inner_body in
            { expr with pexp_desc = Pexp_function (params, constraint_, Pfunction_body transformed_inner) }
        | _ ->
            let transformed_body = fn inner_body in
            { expr with pexp_desc = Pexp_function (params, constraint_, Pfunction_body transformed_body) })
    | _ -> fn expr
  in
  find_innermost_body_and_transform expr

let transform_fun_arguments expr fn =
  match expr.pexp_desc with
  | Pexp_function (params, constraint_, Pfunction_body expression) ->
      let new_params =
        List.map
          ~f:(fun param ->
            match param.pparam_desc with
            | Pparam_val (label, def, patt) -> { param with pparam_desc = Pparam_val (label, def, fn patt) }
            | Pparam_newtype _ -> param)
          params
      in
      { expr with pexp_desc = Pexp_function (new_params, constraint_, Pfunction_body expression) }
  | _ -> expr

let transform_labelled_arguments_type (core_type : core_type) fn =
  let rec inner core_type =
    match core_type.ptyp_desc with
    | Ptyp_arrow (label, core_type_1, core_type_2) ->
        ptyp_arrow ~loc:core_type.ptyp_loc label (fn core_type_1) (inner core_type_2)
    | _ -> core_type
  in
  inner core_type

let get_label_or_empty = function Labelled str | Optional str -> str | Nolabel -> ""

let safe_type_from_label = function
  | (Labelled name | Optional name) when String.length name > 0 && name.[0] = '_' -> "T" ^ name
  | Labelled name | Optional name -> name
  | Nolabel -> "T"

(* Keep this keyword list and translate_mel_obj_label in sync with Melange's
   Lam_methname.translate implementation. *)
let mel_obj_keywords =
  [
    "and";
    "as";
    "assert";
    "begin";
    "class";
    "constraint";
    "do";
    "done";
    "downto";
    "else";
    "end";
    "exception";
    "external";
    "false";
    "for";
    "fun";
    "function";
    "functor";
    "if";
    "in";
    "include";
    "inherit";
    "initializer";
    "lazy";
    "let";
    "match";
    "method";
    "module";
    "mutable";
    "new";
    "nonrec";
    "object";
    "of";
    "open";
    "or";
    "private";
    "rec";
    "sig";
    "struct";
    "then";
    "to";
    "true";
    "try";
    "type";
    "val";
    "virtual";
    "when";
    "while";
    "with";
    "mod";
    "land";
    "lor";
    "lxor";
    "lsl";
    "lsr";
    "asr";
  ]

let find_double_underscore name =
  let rec go index =
    if index < 0 then -1
    else if index + 1 < String.length name && name.[index] = '_' && name.[index + 1] = '_' then index
    else go (index - 1)
  in
  go (String.length name - 2)

let translate_mel_obj_label name =
  let valid_start_char = function '_' | 'a' .. 'z' -> true | _ -> false in
  let double_underscore_index = find_double_underscore name in
  if double_underscore_index = 0 then name
  else if double_underscore_index > 0 then String.sub name 0 double_underscore_index
  else
    match name.[0] with
    | '_' when String.length name > 1 ->
        let candidate = String.sub name 1 (String.length name - 1) in
        if (not (valid_start_char candidate.[0])) || List.exists mel_obj_keywords ~f:(String.equal candidate) then
          candidate
        else name
    | _ -> name

type js_object_field = { method_name : string; js_name : string; present_expr : expression; value_expr : expression }

type component_arg = {
  public_label : arg_label;
  internal_label : arg_label;
  default_value : expression option;
  loc : Location.t;
  core_type : core_type option;
}

let option_is_some_expr ~loc expr = [%expr match [%e expr] with None -> false | Some _ -> true]

let js_obj_internal_expression ~loc name =
  pexp_ident ~loc { txt = Ldot (Ldot (Ldot (Lident "Js", "Obj"), "Internal"), name); loc }

let is_literal_bool_expr expr =
  match expr.pexp_desc with Pexp_construct ({ txt = Lident ("true" | "false"); _ }, None) -> true | _ -> false

(* Emits a registered [Js.t] object. Field values live in plain refs read by
   the object methods; the registry entry is a single [Deferred] thunk that
   builds the per-field [entry] records (a record plus two boxing closures
   each) only if [Js.Obj.keys]/[assign]/[merge] ever inspects the object.
   This keeps object construction — every [makeProps] call and [%mel.obj]
   literal — free of per-field entry allocation. *)
let build_registered_js_object_expression ?as_type ?(register_name = "register_deferred") ~loc fields =
  let generated_fields =
    List.mapi fields ~f:(fun index { method_name; js_name; present_expr; value_expr } ->
        let cell_name = Printf.sprintf "__js_obj_cell_%d" index in
        let cell_binding =
          value_binding ~loc ~pat:(ppat_var ~loc { loc; txt = cell_name }) ~expr:[%expr Stdlib.ref [%e value_expr]]
        in
        (* [present] must be evaluated at construction time (it reads the
           original argument), but literal booleans can be inlined into the
           thunk directly. *)
        let present_bindings, present_in_thunk =
          if is_literal_bool_expr present_expr then ([], present_expr)
          else
            let present_name = Printf.sprintf "__js_obj_present_%d" index in
            ( [ value_binding ~loc ~pat:(ppat_var ~loc { loc; txt = present_name }) ~expr:present_expr ],
              evar ~loc present_name )
        in
        let entry_expr =
          pexp_apply ~loc
            (js_obj_internal_expression ~loc "deferred_entry")
            [
              (Labelled "method_name", estring ~loc method_name);
              (Labelled "js_name", estring ~loc js_name);
              (Labelled "present", present_in_thunk);
              (Nolabel, evar ~loc cell_name);
            ]
        in
        let method_body = pexp_apply ~loc (evar ~loc "!") [ (Nolabel, evar ~loc cell_name) ] in
        let method_ = pcf_method ~loc (Located.mk method_name ~loc, Public, Cfk_concrete (Fresh, method_body)) in
        (cell_binding :: present_bindings, entry_expr, method_))
  in
  let field_bindings, entry_exprs, methods =
    List.fold_right generated_fields ~init:([], [], [])
      ~f:(fun (bindings, entry_expr, method_) (field_bindings, entry_exprs, methods) ->
        (bindings @ field_bindings, entry_expr :: entry_exprs, method_ :: methods))
  in
  let object_name = "__js_obj" in
  let object_binding =
    value_binding ~loc
      ~pat:(ppat_var ~loc { loc; txt = object_name })
      ~expr:(pexp_object ~loc (class_structure ~self:(ppat_any ~loc) ~fields:methods))
  in
  let entries_thunk = [%expr fun () -> [%e pexp_list ~loc entry_exprs]] in
  let register_call =
    pexp_apply ~loc
      (js_obj_internal_expression ~loc register_name)
      [ (Nolabel, evar ~loc object_name); (Nolabel, entries_thunk) ]
  in
  let register_call =
    match as_type with None -> register_call | Some core_type -> pexp_constraint ~loc register_call core_type
  in
  List.fold_right (field_bindings @ [ object_binding ]) ~init:register_call ~f:(fun binding acc ->
      pexp_let ~loc Nonrecursive [ binding ] acc)

let option_type ~loc inner = ptyp_constr ~loc { txt = Lident "option"; loc } [ inner ]

let strip_option_type core_type =
  match core_type.ptyp_desc with
  | Ptyp_constr (({ txt = Lident "option"; _ } | { txt = Ldot (Lident "*predef*", "option"); _ }), [ inner ]) ->
      Some inner
  | _ -> None

let component_arg_public_name arg = get_label_or_empty arg.public_label

let component_arg_type_var arg =
  let loc = arg.loc in
  ptyp_var ~loc (safe_type_from_label arg.public_label)

let component_arg_public_arg_type ~from_signature arg =
  match (arg.public_label, arg.core_type, arg.default_value) with
  | Optional _, Some core_type, _ when not from_signature -> (
      match strip_option_type core_type with Some inner -> inner | None -> core_type)
  | Optional _, Some core_type, _ -> core_type
  | _, Some core_type, Some _ -> core_type
  | Labelled _, Some core_type, _ -> core_type
  | (Labelled _ | Optional _), None, _ -> component_arg_type_var arg
  | Nolabel, _, _ -> assert false

let component_arg_field_type arg =
  match (arg.public_label, arg.core_type, arg.default_value) with
  | Optional _, Some core_type, _ -> (
      match strip_option_type core_type with
      | Some inner -> option_type ~loc:arg.loc inner
      | None -> option_type ~loc:arg.loc core_type)
  | _, Some core_type, Some _ -> option_type ~loc:arg.loc core_type
  | Labelled _, Some core_type, _ -> core_type
  | Optional _, None, _ -> option_type ~loc:arg.loc (component_arg_type_var arg)
  | Labelled _, None, _ -> component_arg_type_var arg
  | Nolabel, _, _ -> assert false

let make_props_name fn_name = fn_name ^ "Props"

let make_object_field ~loc (name, core_type) =
  { pof_desc = Otag ({ loc; txt = name }, core_type); pof_loc = loc; pof_attributes = [] }

let make_props_type ~loc args =
  let object_fields = List.map args ~f:(fun arg -> (component_arg_public_name arg, component_arg_field_type arg)) in
  ptyp_constr ~loc
    { txt = Ldot (Lident "Js", "t"); loc }
    [
      {
        ptyp_desc = Ptyp_object (List.map object_fields ~f:(make_object_field ~loc), Closed);
        ptyp_loc = loc;
        ptyp_loc_stack = [];
        ptyp_attributes = [];
      };
    ]

let react_component_like_type ~loc props_type return_type =
  ptyp_constr ~loc { txt = Ldot (Lident "React", "componentLike"); loc } [ props_type; return_type ]

let is_unit_or_any_pattern pattern =
  match pattern.ppat_desc with Ppat_construct ({ txt = Lident "()"; _ }, None) | Ppat_any -> true | _ -> false

let pattern_core_type pattern =
  match pattern.ppat_desc with Ppat_constraint (_, core_type) -> Some core_type | _ -> None

let rec unwrap_component_expression expr =
  match expr.pexp_desc with
  | Pexp_constraint (expr, _) -> unwrap_component_expression expr
  | Pexp_let (_, _, inner) -> unwrap_component_expression inner
  | Pexp_apply (_, [ (Nolabel, inner) ]) -> unwrap_component_expression inner
  | Pexp_apply (_, [ (Nolabel, inner); (Nolabel, { pexp_desc = Pexp_function _; _ }) ]) ->
      unwrap_component_expression inner
  | Pexp_sequence (_, inner) -> unwrap_component_expression inner
  | _ -> expr

let rec collect_component_fun_params acc expr =
  match expr.pexp_desc with
  | Pexp_constraint (expr, _) -> collect_component_fun_params acc expr
  | Pexp_function (params, _, Pfunction_body body) -> (
      match body.pexp_desc with
      | Pexp_function _ -> collect_component_fun_params (acc @ params) body
      | _ -> acc @ params)
  | _ -> acc

let extract_component_args expr =
  let params = collect_component_fun_params [] (unwrap_component_expression expr) in
  let last_index = List.length params - 1 in
  params
  |> List.mapi ~f:(fun index param ->
      match param.pparam_desc with
      | Pparam_newtype _ -> None
      | Pparam_val (Nolabel, _, pattern) when index = last_index && is_unit_or_any_pattern pattern -> None
      | Pparam_val ((Labelled "key" | Optional "key"), _, pattern) ->
          raise_errorf ~loc:pattern.ppat_loc
            "~key cannot be accessed from the component props. Please set the key where the component is being used."
      | Pparam_val (((Labelled _ | Optional _) as arg_label), default_value, pattern) ->
          Some
            {
              public_label = arg_label;
              internal_label = arg_label;
              default_value;
              loc = pattern.ppat_loc;
              core_type = pattern_core_type pattern;
            }
      | Pparam_val (Nolabel, _, pattern) ->
          raise_errorf ~loc:pattern.ppat_loc
            "props need to be labelled arguments. If your component doesn't have any props, use () or _ instead of a \
             name.")
  |> List.filter_map ~f:Fun.id

let build_make_props_binding ~loc ~fn_name args =
  let props_type = make_props_type ~loc args in
  let object_fields =
    List.map args ~f:(fun arg ->
        let field_name = component_arg_public_name arg in
        let value_expr = evar ~loc:arg.loc field_name in
        let present_expr =
          match arg.public_label with
          | Optional _ -> option_is_some_expr ~loc:arg.loc value_expr
          | _ -> ebool ~loc:arg.loc true
        in
        { method_name = field_name; js_name = translate_mel_obj_label field_name; present_expr; value_expr })
  in
  let body =
    build_registered_js_object_expression ~as_type:props_type ~register_name:"register_deferred_abstract" ~loc
      object_fields
  in
  let body = pexp_fun ~loc Nolabel None [%pat? ()] body in
  let body =
    List.fold_right args ~init:body ~f:(fun arg acc ->
        let arg_name = component_arg_public_name arg in
        let pattern =
          ppat_constraint ~loc:arg.loc
            (ppat_var ~loc:arg.loc { txt = arg_name; loc = arg.loc })
            (component_arg_field_type arg)
        in
        pexp_fun ~loc:arg.loc arg.public_label None pattern acc)
  in
  value_binding ~loc ~pat:(ppat_var ~loc { txt = make_props_name fn_name; loc }) ~expr:body

let build_public_component_binding ~loc ~fn_name args =
  let props_type = make_props_type ~loc args in
  let has_props = args <> [] in
  let props_name = if has_props then "Props" else "_Props" in
  let props_expr = evar ~loc (if has_props then "Props" else "_Props") in
  let key_arg = (Optional "key", evar ~loc "key") in
  let prop_args =
    List.map args ~f:(fun arg ->
        let value_expr = pexp_send ~loc props_expr { txt = component_arg_public_name arg; loc = arg.loc } in
        (arg.internal_label, value_expr))
  in
  let call_expr = pexp_apply ~loc (ident ~loc fn_name) ((key_arg :: prop_args) @ [ (Nolabel, [%expr ()]) ]) in
  let props_pattern = ppat_constraint ~loc (ppat_var ~loc { txt = props_name; loc }) props_type in
  let body = pexp_fun ~loc Nolabel None props_pattern call_expr in
  let key_pattern = ppat_constraint ~loc (ppat_var ~loc { txt = "key"; loc }) [%type: string option] in
  let body = pexp_fun ~loc (Optional "key") None key_pattern body in
  value_binding ~loc ~pat:(ppat_var ~loc { txt = fn_name; loc }) ~expr:body

let build_make_props_signature_item ~loc ~fn_name args =
  let props_type = make_props_type ~loc args in
  let make_props_type =
    List.fold_right args
      ~init:(ptyp_arrow ~loc Nolabel [%type: unit] props_type)
      ~f:(fun arg acc ->
        ptyp_arrow ~loc:arg.loc arg.public_label (component_arg_public_arg_type ~from_signature:true arg) acc)
  in
  psig_value ~loc
    {
      pval_name = { txt = make_props_name fn_name; loc };
      pval_type = make_props_type;
      pval_prim = [];
      pval_attributes = [];
      pval_loc = loc;
    }

let build_public_component_signature_item ~loc ~fn_name ~attributes args return_type =
  let props_type = make_props_type ~loc args in
  let make_type = react_component_like_type ~loc props_type return_type in
  psig_value ~loc
    {
      pval_name = { txt = fn_name; loc };
      pval_type = make_type;
      pval_prim = [];
      pval_attributes = attributes;
      pval_loc = loc;
    }

let extract_component_signature_args pval_type =
  let rec go args core_type =
    match core_type.ptyp_desc with
    | Ptyp_arrow (Nolabel, { ptyp_desc = Ptyp_constr ({ txt = Lident "unit"; _ }, []); _ }, rest) -> go args rest
    | Ptyp_arrow (((Labelled _ | Optional _) as arg_label), core_type, rest) ->
        go
          (args
          @ [
              {
                public_label = arg_label;
                internal_label = arg_label;
                default_value = None;
                loc = core_type.ptyp_loc;
                core_type = Some core_type;
              };
            ])
          rest
    | Ptyp_arrow (Nolabel, _, rest) -> go args rest
    | _ -> (args, core_type)
  in
  go [] pval_type

let expand_make_binding binding react_element_variant_wrapping =
  let attributers = binding.pvb_attributes |> List.filter ~f:nonReactAttributes in
  let loc = binding.pvb_loc in
  let ghost_loc = { binding.pvb_loc with loc_ghost = true } in
  let binding_with_unit = add_unit_at_the_last_argument binding.pvb_expr in
  let binding_expr = transform_fun_body_expression binding_with_unit react_element_variant_wrapping in
  let name = ppat_var ~loc:ghost_loc { txt = get_function_name binding; loc = ghost_loc } in
  let key_arg = Optional "key" in
  let default_value = None in
  let underscore = ppat_var ~loc:ghost_loc { txt = "_"; loc } in
  let core_type = [%type: string option] in
  let key_pattern = ppat_constraint ~loc underscore core_type in
  (* Append key argument since we want to allow users of this component to set key (and assign it to _ since it shouldn't be used) *)
  let function_body = pexp_fun ~loc:ghost_loc key_arg default_value key_pattern binding_expr in
  (* Since expand_make_binding is called on both native and js contexts, we need to keep the attributes *)
  { (value_binding ~loc:ghost_loc ~pat:name ~expr:function_body) with pvb_attributes = attributers }

let expand_make_binding_with_key binding react_element_variant_wrapping =
  let attributers = binding.pvb_attributes |> List.filter ~f:nonReactAttributes in
  let loc = binding.pvb_loc in
  let ghost_loc = { binding.pvb_loc with loc_ghost = true } in
  let binding_with_unit = add_unit_at_the_last_argument binding.pvb_expr in
  let key_expr = evar ~loc:ghost_loc "key" in
  let binding_expr = transform_fun_body_expression binding_with_unit (react_element_variant_wrapping ~key:key_expr) in
  let name = ppat_var ~loc:ghost_loc { txt = get_function_name binding; loc = ghost_loc } in
  let key_arg = Optional "key" in
  let default_value = None in
  let key_pattern =
    ppat_constraint ~loc (ppat_var ~loc:ghost_loc { txt = "key"; loc = ghost_loc }) [%type: string option]
  in
  let function_body = pexp_fun ~loc:ghost_loc key_arg default_value key_pattern binding_expr in
  { (value_binding ~loc:ghost_loc ~pat:name ~expr:function_body) with pvb_attributes = attributers }

let get_arguments pvb_expr =
  let rec go acc = function
    | Pexp_function (params, _, Pfunction_body expr) ->
        let args =
          List.filter_map
            ~f:(function
              | { pparam_desc = Pparam_val (label, default, patt); _ } -> Some (label, default, patt) | _ -> None)
            params
        in
        go (args @ acc) expr.pexp_desc
    | _ -> acc
  in
  go [] pvb_expr.pexp_desc

let string_of_core_type (core_type : core_type) =
  let formatter = Format.str_formatter in
  Astlib.Pprintast.core_type formatter core_type;
  Format.flush_str_formatter ()

let rec make_json_decoder ~loc (core_type : core_type) =
  match core_type with
  | [%type: string] -> [%expr Melange_json.Primitives.string_of_json]
  | [%type: bool] -> [%expr Melange_json.Primitives.bool_of_json]
  | [%type: int] -> [%expr Melange_json.Primitives.int_of_json]
  | [%type: float] -> [%expr Melange_json.Primitives.float_of_json]
  | [%type: int64] -> [%expr Melange_json.Primitives.int64_of_json]
  | [%type: char] ->
      [%expr
        fun json ->
          let s = Melange_json.Primitives.string_of_json json in
          if String.length s = 1 then String.get s 0
          else Melange_json.of_json_error ~json "expected a single-character string"]
  | [%type: unit] -> [%expr Melange_json.Primitives.unit_of_json]
  | [%type: [%t? inner_type] option] ->
      let decode = make_json_decoder ~loc inner_type in
      [%expr Melange_json.Primitives.option_of_json [%e decode]]
  | [%type: [%t? inner_type] list] ->
      let decode = make_json_decoder ~loc inner_type in
      [%expr Melange_json.Primitives.list_of_json [%e decode]]
  | [%type: [%t? inner_type] array] ->
      let decode = make_json_decoder ~loc inner_type in
      [%expr Melange_json.Primitives.array_of_json [%e decode]]
  | [%type: ([%t? ok_type], [%t? err_type]) result] ->
      let decode_ok = make_json_decoder ~loc ok_type in
      let decode_err = make_json_decoder ~loc err_type in
      [%expr Melange_json.Primitives.result_of_json [%e decode_ok] [%e decode_err]]
  | { ptyp_desc = Ptyp_tuple elements; _ } -> make_tuple_json_decoder ~loc elements
  | type_ -> (
      match type_.ptyp_desc with
      | Ptyp_arrow (_, _, _) -> [%expr fun json -> (json : [%t type_])]
      | _ -> [%expr [%of_json: [%t type_]]])

and make_tuple_json_decoder ~loc types =
  let n = List.length types in
  let vars = List.mapi ~f:(fun i _ -> Printf.sprintf "t%d" i) types in
  let decoders = List.map ~f:(make_json_decoder ~loc) types in
  (* Build pattern: `List [t0; t1; t2; ...] *)
  let list_elements_pat =
    List.fold_right ~f:(fun v acc -> [%pat? [%p ppat_var ~loc { txt = v; loc }] :: [%p acc]]) vars ~init:[%pat? []]
  in
  let match_pat = [%pat? `List [%p list_elements_pat]] in
  (* Build expression: (decoder0 t0, decoder1 t1, ...) *)
  let tuple_elements = List.map2 ~f:(fun decode v -> [%expr [%e decode] [%e evar ~loc v]]) decoders vars in
  let tuple_expr = pexp_tuple ~loc tuple_elements in
  let error_msg = Printf.sprintf "expected a JSON array of length %d" n in
  [%expr
    fun json ->
      match json with
      | [%p match_pat] -> [%e tuple_expr]
      | _ -> Melange_json.of_json_error ~json [%e estring ~loc error_msg]]

let make_of_json_argument ~loc (core_type : core_type) prop = [%expr [%e make_json_decoder ~loc core_type] [%e prop]]

let make_of_rsc ~loc (core_type : core_type) prop =
  let function_error loc =
    [%expr
      [%ocaml.error
        "server-reason-react: you can't pass plain functions into client components. Use Runtime.server_function for \
         server actions."]]
  in
  match core_type with
  | [%type: [%t? inner_type] option] as type_ -> (
      match inner_type.ptyp_desc with
      | Ptyp_arrow (_, _, _) -> function_error type_.ptyp_loc
      | _ -> [%expr [%of_rsc: [%t type_]] [%e prop]])
  | { ptyp_desc = Ptyp_arrow (_, _, _) } -> function_error core_type.ptyp_loc
  | type_ -> [%expr [%of_rsc: [%t type_]] [%e prop]]

let props_of_model ~loc (props : (arg_label * expression option * pattern) list) :
    value_binding list * (longident loc * expression) list =
  List.fold_right
    ~f:(fun (arg_label, default, pattern) (decoder_bindings, fields) ->
      match pattern.ppat_desc with
      | Ppat_construct ({ txt = Lident "()"; _ }, None) -> (decoder_bindings, fields)
      | Ppat_constraint (_, core_type) -> (
          match arg_label with
          | Nolabel ->
              (* This error is raised by reason-react-ppx as well *)
              let loc = pattern.ppat_loc in
              ( decoder_bindings,
                (longident ~loc "error", [%expr [%ocaml.error "props need to be labelled arguments"]]) :: fields )
          | Labelled label | Optional label ->
              let core_type = match default with Some _ -> [%type: [%t core_type] option] | None -> core_type in
              let prop = [%expr props##[%e ident ~loc label]] in
              let value = make_of_rsc ~loc core_type prop in
              (decoder_bindings, (longident ~loc label, value) :: fields))
      | _ ->
          let loc = pattern.ppat_loc in
          let expr =
            match arg_label with
            | Nolabel -> [%expr [%ocaml.error "server-reason-react: client components need type annotations"]]
            | Labelled label | Optional label ->
                let msg =
                  Printf.sprintf
                    "server-reason-react: client components need type annotations. Missing annotation for '%s'" label
                in
                let msg_expr = estring ~loc msg in
                [%expr [%ocaml.error [%e msg_expr]]]
          in
          (decoder_bindings, (longident ~loc "error", expr) :: fields))
    props ~init:([], [])

let react_component_attribute ~loc =
  { attr_name = { txt = "react.component"; loc }; attr_payload = PStr []; attr_loc = loc }

let mel_obj ~loc fields =
  match fields with
  (* QUESTION: Maybe unit would work here best, for correctness? *)
  | [] -> [%expr Js.Obj.empty ()]
  | _ ->
      let record = pexp_record ~loc fields None in
      let stri = pstr_eval ~loc record [] in
      [%expr [%mel.obj [%%i stri]]]

let expand_make_binding_to_client binding =
  let loc = binding.pvb_loc in
  let ghost_loc = { binding.pvb_loc with loc_ghost = true } in
  let arguments = get_arguments binding.pvb_expr in
  let promise_decoder_bindings, props_fields = props_of_model ~loc arguments in
  let props_as_object_with_decoders = mel_obj ~loc props_fields in
  let make_call = [%expr React.createElement make [%e props_as_object_with_decoders]] in
  let name = ppat_var ~loc:ghost_loc { txt = "make_client"; loc = ghost_loc } in
  let client_single_argument = ppat_var ~loc:ghost_loc { txt = "props"; loc } in
  let function_body = pexp_fun ~loc:ghost_loc Nolabel None client_single_argument make_call in
  let function_body =
    match promise_decoder_bindings with
    | [] -> function_body
    | _ -> pexp_let ~loc:ghost_loc Nonrecursive promise_decoder_bindings function_body
  in
  value_binding ~loc:ghost_loc ~pat:name ~expr:function_body

let rec add_unit_at_the_last_argument_in_core_type core_type =
  match core_type.ptyp_desc with
  | Ptyp_arrow (arg_label, core_type_1, core_type_2) ->
      {
        core_type with
        ptyp_desc = Ptyp_arrow (arg_label, core_type_1, add_unit_at_the_last_argument_in_core_type core_type_2);
      }
  | Ptyp_constr _ ->
      let loc = core_type.ptyp_loc in
      { core_type with ptyp_desc = Ptyp_arrow (Nolabel, [%type: unit], core_type) }
  | _ -> core_type

let rewrite_signature_item signature_item =
  match signature_item with
  | { psig_loc = loc; psig_desc = Psig_value { pval_name = { txt = fn_name; _ }; pval_attributes; pval_type } } -> (
      match List.filter ~f:hasAnyReactComponentAttribute pval_attributes with
      | [] -> signature_item
      | [ _ ] ->
          let args, return_type = extract_component_signature_args pval_type in
          let make_props_sig = build_make_props_signature_item ~loc ~fn_name args in
          let make_sig =
            build_public_component_signature_item ~loc ~fn_name
              ~attributes:(List.filter ~f:nonReactAttributes pval_attributes)
              args return_type
          in
          [%sigi:
            include sig
              [%%i make_props_sig]
              [%%i make_sig]
            end]
      | _ ->
          [%sigi:
            [%%ocaml.error "server-reason-react: there's seems to be an error in the signature of the component."]])
  | _ -> signature_item

let make_to_model ~loc (core_type : core_type) prop =
  let function_error loc =
    [%expr
      [%ocaml.error
        "server-reason-react: you can't pass plain functions into client components. Use Runtime.server_function for \
         server actions."]]
  in
  match core_type with
  | [%type: [%t? inner_type] option] as type_ -> (
      match inner_type.ptyp_desc with
      | Ptyp_arrow (_, _, _) -> function_error type_.ptyp_loc
      | _ -> [%expr RSC.to_model ([%to_rsc: [%t type_]] [%e prop])])
  | { ptyp_desc = Ptyp_arrow (_, _, _) } -> function_error core_type.ptyp_loc
  | type_ -> [%expr RSC.to_model ([%to_rsc: [%t type_]] [%e prop])]

let props_to_model ~loc (props : (arg_label * expression option * pattern) list) =
  List.fold_left ~init:[%expr []]
    ~f:(fun acc (arg_label, _default, pattern) ->
      match pattern.ppat_desc with
      | Ppat_construct ({ txt = Lident "()"; _ }, None) -> acc
      | Ppat_constraint (_, core_type) -> (
          match arg_label with
          | Nolabel ->
              (* This error is raised by reason-react-ppx as well *)
              let loc = pattern.ppat_loc in
              [%expr [%ocaml.error "props need to be labelled arguments"] :: [%e acc]]
          | Labelled label | Optional label ->
              let prop = ident ~loc label in
              let value = make_to_model ~loc core_type prop in
              let name = estring ~loc label in
              [%expr ([%e name], [%e value]) :: [%e acc]])
      (* TODO: Add all ppat_desc possibilities *)
      | _ ->
          let loc = pattern.ppat_loc in
          let expr =
            match arg_label with
            | Nolabel -> [%expr [%ocaml.error "server-reason-react: client components need type annotations"]]
            | Labelled label | Optional label ->
                let msg =
                  Printf.sprintf
                    "server-reason-react: client components need type annotations. Missing annotation for '%s'" label
                in
                let msg_expr = estring ~loc msg in
                [%expr [%ocaml.error [%e msg_expr]]]
          in
          [%expr [%e expr] :: [%e acc]])
    props

module ServerFunction = struct
  let rec last_expr_to_fn ~loc expr fn =
    match expr.pexp_desc with
    | Pexp_constraint (expr, _) -> last_expr_to_fn ~loc expr fn
    | Pexp_function (params, constraint_, Pfunction_body expression) when params <> [] -> (
        match expression.pexp_desc with
        | Pexp_function _ ->
            let transformed_inner = last_expr_to_fn ~loc expression fn in
            { expr with pexp_desc = Pexp_function (params, constraint_, Pfunction_body transformed_inner) }
        | _ -> { expr with pexp_desc = Pexp_function (params, constraint_, Pfunction_body fn) })
    | _ -> fn

  let generate_id ~loc name =
    let file_path = loc.loc_start.pos_fname in
    let replacement =
      match shared_folder_prefix.contents with
      | Some x ->
          if match_substring file_path x then x
          else raise_errorf ~loc "Prefix doesn't match the file path. Provide a prefix that matches the file path."
      | None -> raise_errorf ~loc "Found a server.function without --shared-folder-prefix argument. Provide one."
    in
    (* We need to add a nasty hack here, since have different files for native and melange.Assume that the file structure is native/lib and js, and replace the name directly. This is supposed to be temporal, until dune implements https://github.com/ocaml/dune/issues/10630 *)
    let file_path = Str.replace_first (Str.regexp replacement) "" file_path in
    let hash = Printf.sprintf "%s_%s_%d" name file_path loc.loc_start.pos_lnum |> Hashtbl.hash |> string_of_int in
    hash

  let get_arg_details (arg : arg_label * expression option * pattern) =
    let arg_label, default, pattern = arg in
    let loc = pattern.ppat_loc in
    match pattern.ppat_desc with
    | Ppat_construct ({ txt = Lident "()"; loc }, None) -> Ok (Nolabel, None, [%type: unit])
    | Ppat_constraint (pattern, core_type) -> (
        let loc = pattern.ppat_loc in
        let core_type = match default with Some _ -> [%type: [%t core_type] option] | None -> core_type in
        match pattern.ppat_desc with
        | Ppat_var { txt = label; _ } -> Ok (arg_label, Some label, core_type)
        | _ -> Error (loc, "server-reason-react: server function arguments must have a name"))
    | _ -> Error (loc, "server-reason-react: server function arguments must have type annotations")

  let get_response_type expr =
    let rec aux expr acc =
      match expr.pexp_desc with
      | Pexp_function (_, Some (Pconstraint core_type), Pfunction_body body) -> aux body (Some core_type)
      | Pexp_function (_, _, Pfunction_body body) -> aux body acc
      | Pexp_constraint (expr, core_type) -> aux expr (Some core_type)
      | _ -> acc
    in
    aux expr None

  let response_to_model ~loc core_type response =
    match core_type with
    | Some [%type: [%t? core_type] Js.Promise.t] -> [%expr RSC.to_model ([%to_rsc: [%t core_type]] [%e response])]
    | Some _ -> [%expr [%ocaml.error "server-reason-react: server functions must return a promise"]]
    | _ ->
        [%expr [%ocaml.error "server-reason-react: server functions must have a return type annotation (Js.Promise.t)"]]

  let map_arguments_to_expressions ~loc args =
    List.map
      ~f:(fun arg ->
        match arg with
        | Ok (arg_label, Some arg_name, _) -> (arg_label, [%expr [%e evar ~loc arg_name]])
        | Ok (arg_label, _, [%type: unit]) -> (arg_label, [%expr ()])
        | Ok _ ->
            ( Nolabel,
              [%expr
                [%ocaml.error
                  "server-reason-react: invalid argument, it must have a argument with name and type annotation"]] )
        | Error (loc, msg) -> (Nolabel, [%expr [%ocaml.error [%e estring ~loc msg]]]))
      args

  let encode_function_response ~loc ~response_expr ~core_type =
    [%expr
      try [%e response_expr] |> Lwt.map (fun response -> [%e response_to_model ~loc core_type [%expr response]])
      with e -> Lwt.fail e]

  let decode_arguments_vb ~loc args_to_decode =
    args_to_decode
    |> List.mapi ~f:(fun i (_, label, core_type) ->
        let core_type_string = string_of_core_type core_type in
        let of_json = make_of_json_argument ~loc core_type [%expr Stdlib.Array.unsafe_get args [%e eint ~loc i]] in
        value_binding ~loc
          ~pat:[%pat? [%p ppat_var ~loc { txt = label; loc }]]
          ~expr:
            [%expr
              try [%e of_json]
              with _ ->
                Stdlib.raise
                  (Invalid_argument
                     (Stdlib.Printf.sprintf
                        "server-reason-react: error on decoding argument '%s'. EXPECTED: %s, RECEIVED: %s"
                        [%e estring ~loc label] [%e estring ~loc core_type_string]
                        (Stdlib.Array.unsafe_get args [%e eint ~loc i] |> Yojson.Basic.to_string)))])

  let create_function_reference_registration ~loc ~id ~function_name ~args ~core_type =
    let apply_args = map_arguments_to_expressions ~loc args in
    let response_expr = pexp_apply ~loc [%expr [%e evar ~loc function_name].call] apply_args in

    let encoded_response_expr = encode_function_response ~loc ~response_expr ~core_type in
    let args_to_decode =
      List.filter_map
        ~f:(fun arg ->
          match arg with
          | Ok (_, _, [%type: Js.FormData.t]) -> None
          | Ok (arg_label, Some arg_name, core_type) -> Some (arg_label, arg_name, core_type)
          | Ok _ -> None
          | Error _ -> None)
        args
    in

    let args, formData =
      List.partition_map
        ~f:(fun arg ->
          match arg with Ok (_, _, [%type: Js.FormData.t]) -> Right arg | Ok _ -> Left arg | Error _ -> Left arg)
        args
    in

    let body_expr =
      match args_to_decode with
      | [] -> encoded_response_expr
      | args_to_decode ->
          let decoded_expr = decode_arguments_vb ~loc args_to_decode in
          pexp_let ~loc Nonrecursive decoded_expr encoded_response_expr
    in
    match (formData, args) with
    | [], _ -> [%stri FunctionReferences.register [%e estring ~loc id] (Body (fun args -> [%e body_expr]))]
    | [ _ ], [] ->
        [%stri FunctionReferences.register [%e estring ~loc id] (FormData (fun _ formData -> [%e body_expr]))]
    | _, [] ->
        [%stri [%ocaml.error "server-reason-react: server functions with form data must have at only one argument"]]
    | _ -> [%stri FunctionReferences.register [%e estring ~loc id] (FormData (fun args formData -> [%e body_expr]))]

  let create_server_function_record ~loc id expression =
    [%expr { Runtime.id = [%e estring ~loc id]; call = [%e expression] }]

  let rewrite_native_function ~vb ~rec_flag structure_item =
    let loc = structure_item.pstr_loc in
    let function_name = get_function_name vb in
    let args = get_arguments vb.pvb_expr |> List.map ~f:get_arg_details |> List.rev in
    let base_fn = vb.pvb_expr in
    let return_core_type = get_response_type base_fn in
    let id = generate_id ~loc:vb.pvb_loc function_name in
    let server_function_record_vb =
      value_binding ~loc:vb.pvb_loc ~pat:vb.pvb_pat ~expr:(create_server_function_record ~loc:vb.pvb_loc id base_fn)
    in
    let stri =
      [%stri
        include struct
          [%%i pstr_value ~loc rec_flag [ server_function_record_vb ]]
          [%%i create_function_reference_registration ~loc ~id ~function_name ~args ~core_type:return_core_type]
        end]
    in
    stri

  let response_of_rsc ~loc core_type response =
    match core_type with
    | Some [%type: [%t? core_type] Js.Promise.t] -> [%expr [%of_rsc: [%t core_type]] [%e response]]
    | Some _ -> [%expr [%ocaml.error "server-reason-react: server functions must return a promise"]]
    | _ ->
        [%expr [%ocaml.error "server-reason-react: server functions must have a return type annotation (Js.Promise.t)"]]

  let create_client_function ~loc ~return_core_type id args =
    let decode_response = response_of_rsc ~loc return_core_type in
    let apply_args = map_arguments_to_expressions ~loc args |> List.map ~f:(fun (_, expr) -> (Nolabel, expr)) in
    let fn =
      [%expr
        let action = ReactServerDOMEsbuild.createServerReference [%e estring ~loc id] in
        ([%e pexp_apply ~loc [%expr action] apply_args] [@u])
        |> Js.Promise.then_ (fun response -> Js.Promise.resolve [%e decode_response [%expr response]])]
    in
    fn

  let rewrite_client_function ~nested_module_names ~vb ~rec_flag structure_item =
    let loc = structure_item.pstr_loc in

    let function_name = get_function_name vb in
    let args = get_arguments vb.pvb_expr |> List.map ~f:get_arg_details |> List.rev in
    let base_fn = vb.pvb_expr in
    let return_core_type = get_response_type base_fn in
    let id = generate_id ~loc:vb.pvb_loc function_name in
    let server_function_record_vb =
      value_binding ~loc:vb.pvb_loc ~pat:vb.pvb_pat
        ~expr:
          (create_server_function_record ~loc:vb.pvb_loc id
             (last_expr_to_fn ~loc base_fn (create_client_function ~loc ~return_core_type id args)))
    in

    let loc = structure_item.pstr_loc in
    let module_name = String.concat "." nested_module_names in
    let _, formData =
      List.partition_map
        ~f:(fun arg ->
          match arg with Ok (_, _, [%type: Js.FormData.t]) -> Right arg | Ok _ -> Left arg | Error _ -> Left arg)
        args
    in
    let functionToCall = match formData with [] -> function_name | _ -> Printf.sprintf "%s.call" function_name in
    let comment = Printf.sprintf "// extract-server-function %s %s %s" id functionToCall module_name in
    let raw = estring ~loc comment in
    let extract_client_raw = [%stri [%%raw [%e raw]]] in
    [%stri
      include struct
        [%%i extract_client_raw]
        [%%i pstr_value ~loc:structure_item.pstr_loc rec_flag [ server_function_record_vb ]]
      end]
end

let rewrite_structure_item ~nested_module_names structure_item =
  match structure_item.pstr_desc with
  (* external *)
  | Pstr_primitive ({ pval_name = { txt = _fnName }; pval_attributes; pval_type = _ } as _value_description) -> (
      match
        List.filter
          ~f:(fun attr -> hasAttr attr react_dot_component || hasAttr attr react_dot_async_dot_component)
          pval_attributes
      with
      | [] -> structure_item
      | _ ->
          let loc = structure_item.pstr_loc in
          [%stri
            [%%ocaml.error
            "externals aren't supported on server-reason-react. externals are used to bind to React components defined \
             in JavaScript, in the server, that doesn't make sense. If you need to render this on the server, \
             implement a placeholder or an empty element"]])
  (* let make = ... *)
  | Pstr_value (rec_flag, value_bindings) when isReactServerFunctionBinding (List.hd value_bindings) ->
      let vb = List.hd value_bindings in
      let loc = structure_item.pstr_loc in
      if List.length value_bindings > 1 then
        [%stri
          [%%ocaml.error
          "server-reason-react: server functions don't support recursive bindings yet. If you need it, please open an \
           issue on https://github.com/reasonml-community/server-reason-react/issues"]]
      else ServerFunction.rewrite_native_function ~vb ~rec_flag structure_item
  | Pstr_value (rec_flag, value_bindings) -> (
      try
        let rewrite_component_binding vb =
          if isReactClientComponentBinding vb then
            let loc = vb.pvb_loc in
            let fn_name = get_function_name vb in
            let args = extract_component_args vb.pvb_expr in
            let internal_binding =
              expand_make_binding_with_key vb (fun ~key expr ->
                  let loc = expr.pexp_loc in
                  let fileName = expr.pexp_loc.loc_start.pos_fname in
                  let replacement =
                    match shared_folder_prefix.contents with
                    | Some prefix ->
                        if match_substring fileName prefix then prefix
                        else
                          raise_errorf ~loc
                            "Prefix doesn't match the file path. Provide a prefix that matches the file path."
                    | None ->
                        raise_errorf ~loc
                          "Found a react.client.component without --shared-folder-prefix argument. Provide one."
                  in
                  let file = fileName |> Str.replace_first (Str.regexp replacement) "" |> estring ~loc in
                  let import_module =
                    match nested_module_names with
                    | [] -> file
                    | _ ->
                        let submodule = estring ~loc (String.concat "." nested_module_names) in
                        [%expr Printf.sprintf "%s#%s" [%e file] [%e submodule]]
                  in
                  let arguments = get_arguments vb.pvb_expr in
                  let props = props_to_model ~loc arguments in
                  [%expr
                    React.Client_component
                      {
                        key = [%e key];
                        import_module = [%e import_module];
                        import_name = "";
                        props = [%e props];
                        client = React.Upper_case_component (Stdlib.__FUNCTION__, fun () -> [%e expr]);
                      }])
            in
            Some
              ( internal_binding,
                build_make_props_binding ~loc ~fn_name args,
                build_public_component_binding ~loc ~fn_name args )
          else if isReactComponentBinding vb then
            let loc = vb.pvb_loc in
            let fn_name = get_function_name vb in
            let args = extract_component_args vb.pvb_expr in
            let internal_binding =
              expand_make_binding vb (fun expr ->
                  let loc = expr.pexp_loc in
                  [%expr React.Upper_case_component (Stdlib.__FUNCTION__, fun () -> [%e expr])])
            in
            Some
              ( internal_binding,
                build_make_props_binding ~loc ~fn_name args,
                build_public_component_binding ~loc ~fn_name args )
          else if isReactAsyncComponentBinding vb then
            let loc = vb.pvb_loc in
            let fn_name = get_function_name vb in
            let args = extract_component_args vb.pvb_expr in
            let internal_binding =
              expand_make_binding vb (fun expr ->
                  let loc = expr.pexp_loc in
                  [%expr React.Async_component (Stdlib.__FUNCTION__, fun () -> [%e expr])])
            in
            Some
              ( internal_binding,
                build_make_props_binding ~loc ~fn_name args,
                build_public_component_binding ~loc ~fn_name args )
          else None
        in
        let rewrites = List.map value_bindings ~f:rewrite_component_binding in
        let make_props_bindings =
          List.filter_map rewrites ~f:(function
            | Some (_, make_props_binding, _) -> Some make_props_binding
            | None -> None)
        in
        let public_bindings =
          List.filter_map rewrites ~f:(function Some (_, _, public_binding) -> Some public_binding | None -> None)
        in
        let internal_bindings =
          List.map2 value_bindings rewrites ~f:(fun vb rewrite ->
              match rewrite with Some (internal_binding, _, _) -> internal_binding | None -> vb)
        in
        match make_props_bindings with
        | [] -> pstr_value ~loc:structure_item.pstr_loc rec_flag internal_bindings
        | _ -> (
            let loc = structure_item.pstr_loc in
            let is_react_attr attr =
              let name = attr.attr_name.txt in
              String.length name >= 6 && String.sub name 0 6 = "react."
            in
            let propagated_attrs =
              List.concat_map value_bindings ~f:(fun vb ->
                  List.filter vb.pvb_attributes ~f:(fun attr -> not (is_react_attr attr)))
            in
            let include_stri =
              [%stri
                include struct
                  [%%i pstr_value ~loc:structure_item.pstr_loc Nonrecursive make_props_bindings]
                  [%%i pstr_value ~loc:structure_item.pstr_loc rec_flag internal_bindings]
                  [%%i pstr_value ~loc:structure_item.pstr_loc Nonrecursive public_bindings]
                end]
            in
            match (propagated_attrs, include_stri.pstr_desc) with
            | _ :: _, Pstr_include incl ->
                { include_stri with pstr_desc = Pstr_include { incl with pincl_attributes = propagated_attrs } }
            | _ -> include_stri)
      with Error err -> pstr_eval ~loc:structure_item.pstr_loc err [])
  | _ -> structure_item

let rewrite_structure_item_for_js ~nested_module_names ctx structure_item =
  match structure_item.pstr_desc with
  (* external *)
  | Pstr_primitive ({ pval_name = { txt = _fnName }; pval_attributes; pval_type = _ } as _value_description) -> (
      match List.filter ~f:(fun attr -> hasAttr attr react_dot_client_dot_component) pval_attributes with
      | [] -> structure_item
      | _ ->
          let loc = structure_item.pstr_loc in
          [%stri [%%ocaml.error "server-reason-react: externals aren't supported on client components yet"]])
  | Pstr_value (rec_flag, value_bindings) when isReactServerFunctionBinding (List.hd value_bindings) ->
      let vb = List.hd value_bindings in
      ServerFunction.rewrite_client_function ~nested_module_names ~vb ~rec_flag structure_item
  (* let make = ... *)
  | Pstr_value (rec_flag, value_bindings) when isClientComponentBinding value_bindings ->
      let first_value_binding = List.hd value_bindings in
      let make_client = expand_make_binding_to_client first_value_binding in
      let make_client_binding = pstr_value ~loc:structure_item.pstr_loc rec_flag [ make_client ] in
      let original_value_binding =
        { first_value_binding with pvb_attributes = [ react_component_attribute ~loc:first_value_binding.pvb_loc ] }
      in
      let loc = structure_item.pstr_loc in
      let code_path = Expansion_context.Base.code_path ctx in
      let fileName = Code_path.file_path code_path in
      (* We need to add a nasty hack here, since have different files for native and melange.Assume that the file structure is /native/shared/ and js, and replace the name directly. This is supposed to be temporal, until dune implements https://github.com/ocaml/dune/issues/10630 *)
      let replacement =
        match shared_folder_prefix.contents with
        | Some prefix ->
            if match_substring fileName prefix then prefix
            else raise_errorf ~loc "Prefix doesn't match the file path. Provide a prefix that matches the file path."
        | None ->
            raise_errorf ~loc "Found a react.client.component without --shared-folder-prefix argument. Provide one."
      in
      let fileName = Str.replace_first (Str.regexp replacement) "" fileName in
      let comment =
        match nested_module_names with
        | [] -> estring ~loc (Printf.sprintf "// extract-client %s" fileName)
        | _ -> estring ~loc (Printf.sprintf "// extract-client %s %s" fileName (String.concat "." nested_module_names))
      in
      [%stri
        include struct
          [%%i [%stri [%%raw [%e comment]]]]
          [%%i pstr_value ~loc:structure_item.pstr_loc rec_flag [ original_value_binding ]]
          [%%i make_client_binding]
        end]
  | _ -> structure_item

let validate_tag_children tag children attributes : (unit, string) result =
  match Html.is_self_closing_tag tag with
  | true when Option.fold ~none:false ~some:(fun children -> List.length children > 0) children ->
      Error (Printf.sprintf {|"%s" is a self-closing tag and must not have "children".\n|} tag)
  | true
    when List.exists
           ~f:(fun (arg_label, _) ->
             match arg_label with
             | Labelled "dangerouslySetInnerHTML" | Optional "dangerouslySetInnerHTML" -> true
             | _ -> false)
           attributes ->
      Error (Printf.sprintf {|server-reason-react: "%s" is a self-closing tag and must not have "children".\n|} tag)
  | false -> Ok ()
  | true -> Ok ()

let traverse =
  object (_)
    inherit [Expansion_context.Base.t] Ast_traverse.map_with_context as super
    val mutable nested_module_names = []

    method! module_binding ctxt module_binding =
      (match module_binding.pmb_name.txt with
      | None -> ()
      | Some name -> nested_module_names <- nested_module_names @ [ name ]);
      let mapped = super#module_binding ctxt module_binding in
      let rec remove_last l = match l with [] -> [] | [ _ ] -> [] | hd :: tl -> hd :: remove_last tl in
      nested_module_names <- remove_last nested_module_names;
      mapped

    method! structure_item ctx structure_item =
      match mode.contents with
      | Native -> rewrite_structure_item ~nested_module_names (super#structure_item ctx structure_item)
      | Js -> rewrite_structure_item_for_js ~nested_module_names ctx (super#structure_item ctx structure_item)

    method! signature_item ctx signature_item =
      match mode.contents with
      | Native -> rewrite_signature_item (super#signature_item ctx signature_item)
      | Js -> super#signature_item ctx signature_item

    method! expression ctx expr =
      let expr = super#expression ctx expr in
      (* Rewrite ReactDOM.Style.make(~k:v, ..., ()) to a direct list literal at
         compile time. On stock OCaml the 347-optional-arg signature forces
         ~1460 words/call; the PPX elides this entirely when all args are
         labeled string values. Falls through if optional args or unknown
         names appear. Native-only: on the JS target [ReactDOM.Style.t] is
         reason-react's abstract type, so a list literal annotated as
         [ReactDOM.Style.t] fails to type-check. *)
      let expr = match mode.contents with Native -> Style_rewrite.rewrite_expression expr | Js -> expr in
      match mode.contents with
      | Js -> (
          (* In the case of expressions, it's the only transformation that needs to be done for JS. This expansion from "styles" prop into "className" and "style" props is a feature by styled-ppx. The existence of this here, is because dune/ppxlib doesn't allow more than one preprocess_impl and even that, the combination of styled-ppx and server-reason-react.ppx doesn't compose properly. *)
          try Styles_attribute.expand expr with Error err -> [%expr [%e err]])
      | Native -> (
          try
            match expr.pexp_desc with
            | Pexp_apply (({ pexp_desc = Pexp_ident _; pexp_loc = loc; _ } as tag), args)
              when has_jsx_attr expr.pexp_attributes -> (
                let children, rest_of_args = split_args args in
                match validate_tag_children (Pprintast.string_of_expression tag) children rest_of_args with
                | Error err -> [%expr [%ocaml.error [%e estring ~loc:expr.pexp_loc err]]]
                | Ok () -> (
                    match tag.pexp_desc with
                    (* div() [@JSX] *)
                    | Pexp_ident { txt = Lident name; loc = _name_loc } ->
                        (* This expansion from "styles" prop into "className" and "style" props is a feature by styled-ppx. The existence of this here, is because dune/ppxlib doesn't allow more than one preprocess_impl and even that, the combination of styled-ppx and server-reason-react.ppx doesn't compose properly. *)
                        let new_args = Styles_attribute.expand_attributes ~loc:tag.pexp_loc rest_of_args in
                        rewrite_lowercase ~loc:expr.pexp_loc name new_args children
                    (* Reason adds `createElement` as default when an uppercase is found,
                   we change it back to make *)
                    (* Foo.createElement() [@JSX] *)
                    | Pexp_ident { txt = Ldot (modulePath, ("createElement" | "make")); loc } ->
                        let id = { loc; txt = Ldot (modulePath, "make") } in
                        rewrite_component ~loc:expr.pexp_loc id rest_of_args children
                    (* local_function() [@JSX] *)
                    | Pexp_ident id -> rewrite_component ~loc:expr.pexp_loc id rest_of_args children
                    | _ -> assert false))
            (* div() [@JSX] *)
            | Pexp_apply (tag, _props) when has_jsx_attr expr.pexp_attributes ->
                raise_errorf ~loc:expr.pexp_loc "jsx: %s should be an identifier, not an expression"
                  (Pprintast.string_of_expression tag)
            (* <> </> is represented as a list in the Parsetree with [@JSX] *)
            | Pexp_construct ({ txt = Lident "::"; loc }, Some { pexp_desc = Pexp_tuple _; _ })
            | Pexp_construct ({ txt = Lident "[]"; loc }, None) -> (
                let jsx_attr, rest_attributes = List.partition ~f:is_jsx expr.pexp_attributes in
                match (jsx_attr, rest_attributes) with
                | [], _ -> expr
                | _, rest_attributes ->
                    let children = transform_items_of_list ~loc expr in
                    let new_expr = [%expr React.fragment (React.list [%e children])] in
                    { new_expr with pexp_attributes = rest_attributes })
            | _ -> expr
          with Error err -> [%expr [%e err]])
  end

let () =
  Driver.add_arg "-melange" (Unit (fun () -> mode := Js)) ~doc:"preprocess for js build";

  Driver.add_arg "-shared-folder-prefix"
    (String
       (fun str ->
         let components = String.split_on_char '/' str |> List.filter ~f:(fun x -> x <> "") in
         let prefix = String.concat "/" components in
         let prefix = if prefix = "" then "" else prefix ^ "/" in
         shared_folder_prefix := Some prefix))
    ~doc:"prefix of shared folder, used to replace the it in the file path";

  Ppxlib.Driver.V2.register_transformation "server-reason-react.ppx" ~preprocess_impl:traverse#structure
    ~preprocess_intf:traverse#signature