firebird

defmodule Firebird.Compiler.WATGen do
  @moduledoc """
  Generates WebAssembly Text Format (WAT) from typed Firebird IR.

  Produces valid WAT that can be assembled to WASM binary by `wat2wasm`.

  ## Output Structure

  ```wat
  (module
    ;; Function declarations
    (func $add (export "add") (param $p0 i64) (param $p1 i64) (result i64)
      local.get $p0
      local.get $p1
      i64.add
    )
    ;; ... more functions
  )
  ```

  ## Code Generation Strategy

  - Uses WASM stack machine model (operands pushed, operations consume them)
  - Pattern matching compiled to `if/else` chains with `i64.eq` comparisons
  - Branchless `select` for simple if/else with trivial branches (avoids branch misprediction)
  - Recursive calls use WASM `call` instruction
  - Integer operations use `i64`, float operations use `f64`
  - Arithmetic ops select i64/f64 instructions based on operand types
  - Boolean results use `i32` with `i64.extend_i32_s` for promotion
  """

  alias Firebird.Compiler.IR

  @doc """
  Generate WAT string from a typed IR module.

  ## Returns

  - `{:ok, wat_string}` on success
  - `{:error, reason}` on failure
  """
  @spec generate(IR.Module.t()) :: {:ok, String.t()} | {:error, term()}
  def generate(%IR.Module{} = module) do
    # Build a map of function name → return type for cross-function type awareness
    func_types = build_func_type_map(module.functions)

    functions_wat =
      module.functions
      |> Enum.map(&generate_function(&1, func_types))
      |> Enum.join("\n\n")

    wat = """
    (module
    #{functions_wat}
    )
    """

    {:ok, String.trim(wat)}
  rescue
    e -> {:error, {:wat_gen_error, Exception.message(e)}}
  end

  # Build a map of function name → %IR.FunctionType{} for all module functions
  defp build_func_type_map(functions) do
    Map.new(functions, fn func ->
      type =
        func.type ||
          %IR.FunctionType{
            params: List.duplicate(:i64, func.arity),
            return: :i64
          }

      {func.name, type}
    end)
  end

  @doc """
  Generate WAT for a single function.
  """
  @spec generate_function(IR.Function.t(), map()) :: String.t()
  def generate_function(%IR.Function{} = func, func_types \\ %{}) do
    type =
      func.type ||
        %IR.FunctionType{
          params: List.duplicate(:i64, func.arity),
          return: :i64
        }

    # Build type environment: param name → wasm type
    type_env = Map.new(Enum.zip(func.params, type.params))

    params_wat =
      Enum.zip(func.params, type.params)
      |> Enum.map(fn {name, ptype} ->
        "(param $#{name} #{wasm_type(ptype)})"
      end)
      |> Enum.join(" ")

    result_wat = "(result #{wasm_type(type.return)})"
    export_wat = "(export \"#{func.name}\")"

    # Collect local variables from let bindings and case subjects
    locals = collect_locals(func.body, func.params)

    # Infer per-local types by analyzing what each let binding stores.
    # This correctly handles cross-function calls (e.g., i64 function storing
    # f64 result from a callee) instead of using a single default type.
    local_types = collect_local_types(func.body, type_env, func_types)

    # TCO temp locals match their corresponding parameter types
    tco_temp_types =
      func.params
      |> Enum.zip(type.params)
      |> Map.new(fn {p, pt} -> {String.to_atom("__tco_tmp_#{p}"), pt} end)

    # __case_subject is always i64 (case subjects are evaluated as i64)
    special_types = Map.put(tco_temp_types, :__case_subject, :i64)

    all_local_types = Map.merge(special_types, local_types)

    locals_wat =
      locals
      |> Enum.map(fn name ->
        lt = Map.get(all_local_types, name, :i64)
        "    (local $#{name} #{wasm_type(lt)})"
      end)
      |> Enum.join("\n")

    # Merge local types into type_env so code generation resolves
    # let-bound variables to their correct WASM types
    full_type_env = Map.merge(type_env, all_local_types)

    body_wat = generate_expr(func.body, type.return, type, full_type_env, func_types)

    params_section = if params_wat == "", do: "", else: " #{params_wat}"
    locals_section = if locals_wat == "", do: "", else: "\n#{locals_wat}"

    "  (func $#{func.name} #{export_wat}#{params_section} #{result_wat}#{locals_section}\n#{indent(body_wat, 4)}\n  )"
  end

  # Collect all local variable names from let bindings and tail loop temps
  defp collect_locals({:tail_loop, _loop_params, body}, params) do
    # For tail loops, we need temp variables for parameter swaps
    inner_locals = collect_locals(body, params)
    temp_locals = Enum.map(params, fn p -> String.to_atom("__tco_tmp_#{p}") end)
    (inner_locals ++ temp_locals) |> Enum.sort() |> Enum.uniq()
  end

  defp collect_locals(expr, params) do
    collect_let_vars(expr, MapSet.new(params))
    |> MapSet.to_list()
    |> Enum.sort()
  end

  defp collect_let_vars({:let, name, value}, params) do
    inner = collect_let_vars(value, params)

    if name in params do
      inner
    else
      MapSet.put(inner, name)
    end
  end

  defp collect_let_vars({:block, exprs}, params) do
    Enum.reduce(exprs, MapSet.new(), fn expr, acc ->
      MapSet.union(acc, collect_let_vars(expr, params))
    end)
  end

  defp collect_let_vars({:if, cond_expr, then_expr, else_expr}, params) do
    MapSet.union(
      collect_let_vars(cond_expr, params),
      MapSet.union(
        collect_let_vars(then_expr, params),
        collect_let_vars(else_expr, params)
      )
    )
  end

  defp collect_let_vars({:binop, _, left, right}, params) do
    MapSet.union(collect_let_vars(left, params), collect_let_vars(right, params))
  end

  defp collect_let_vars({:call, _, args}, params) do
    Enum.reduce(args, MapSet.new(), fn arg, acc ->
      MapSet.union(acc, collect_let_vars(arg, params))
    end)
  end

  defp collect_let_vars({:case, subject, clauses}, params) do
    subject_locals = collect_let_vars(subject, params)

    clause_locals =
      Enum.reduce(clauses, MapSet.new(), fn {pat, guard, body}, acc ->
        # Variable patterns with guards need a local for the bound variable
        pat_locals =
          case pat do
            {:var_pat, name} when guard != nil ->
              if name in params, do: MapSet.new(), else: MapSet.new([name])

            _ ->
              MapSet.new()
          end

        guard_locals =
          case guard do
            nil -> MapSet.new()
            g -> collect_let_vars(g, params)
          end

        acc
        |> MapSet.union(pat_locals)
        |> MapSet.union(guard_locals)
        |> MapSet.union(collect_let_vars(body, params))
      end)

    # Also need __case_subject local
    MapSet.union(subject_locals, clause_locals)
    |> MapSet.put(:__case_subject)
  end

  defp collect_let_vars({:unaryop, _, expr}, params) do
    collect_let_vars(expr, params)
  end

  defp collect_let_vars({:tail_loop, _loop_params, body}, params) do
    collect_let_vars(body, params)
  end

  defp collect_let_vars({:tail_call, _, args}, params) do
    Enum.reduce(args, MapSet.new(), fn arg, acc ->
      MapSet.union(acc, collect_let_vars(arg, params))
    end)
  end

  defp collect_let_vars(_, _params), do: MapSet.new()

  # Infer the WASM type for each local variable by analyzing let binding values.
  # Threads the type environment through sequential bindings in blocks so that
  # later bindings correctly resolve the types of earlier ones.
  # Returns a map of local_name → wasm_type.
  defp collect_local_types(expr, te, ftm) do
    {_updated_te, types} = collect_local_types_acc(expr, te, ftm, %{})
    types
  end

  defp collect_local_types_acc({:block, exprs}, te, ftm, acc) do
    Enum.reduce(exprs, {te, acc}, fn expr, {cur_te, cur_acc} ->
      collect_local_types_acc(expr, cur_te, ftm, cur_acc)
    end)
  end

  defp collect_local_types_acc({:let, name, value}, te, ftm, acc) do
    vtype = infer_expr_wasm_type(value, te, ftm)
    updated_te = Map.put(te, name, vtype)
    updated_acc = Map.put(acc, name, vtype)
    {updated_te, updated_acc}
  end

  defp collect_local_types_acc({:if, cond_expr, then_b, else_b}, te, ftm, acc) do
    {_, acc} = collect_local_types_acc(cond_expr, te, ftm, acc)
    {_, acc} = collect_local_types_acc(then_b, te, ftm, acc)
    {_, acc} = collect_local_types_acc(else_b, te, ftm, acc)
    {te, acc}
  end

  defp collect_local_types_acc({:case, subject, clauses}, te, ftm, acc) do
    {_, acc} = collect_local_types_acc(subject, te, ftm, acc)

    acc =
      Enum.reduce(clauses, acc, fn {pat, guard, body}, cur_acc ->
        # Variable patterns with guards introduce a local typed as i64 (case subject type)
        cur_acc =
          case pat do
            {:var_pat, name} when guard != nil -> Map.put(cur_acc, name, :i64)
            _ -> cur_acc
          end

        cur_acc =
          case guard do
            nil ->
              cur_acc

            g ->
              {_, cur_acc} = collect_local_types_acc(g, te, ftm, cur_acc)
              cur_acc
          end

        {_, cur_acc} = collect_local_types_acc(body, te, ftm, cur_acc)
        cur_acc
      end)

    {te, acc}
  end

  defp collect_local_types_acc({:tail_loop, _, body}, te, ftm, acc) do
    {_, acc} = collect_local_types_acc(body, te, ftm, acc)
    {te, acc}
  end

  defp collect_local_types_acc({:tail_call, _, args}, te, ftm, acc) do
    acc =
      Enum.reduce(args, acc, fn arg, cur_acc ->
        {_, cur_acc} = collect_local_types_acc(arg, te, ftm, cur_acc)
        cur_acc
      end)

    {te, acc}
  end

  defp collect_local_types_acc({:binop, _, left, right}, te, ftm, acc) do
    {_, acc} = collect_local_types_acc(left, te, ftm, acc)
    {_, acc} = collect_local_types_acc(right, te, ftm, acc)
    {te, acc}
  end

  defp collect_local_types_acc({:unaryop, _, expr}, te, ftm, acc) do
    collect_local_types_acc(expr, te, ftm, acc)
  end

  defp collect_local_types_acc({:call, _, args}, te, ftm, acc) do
    acc =
      Enum.reduce(args, acc, fn arg, cur_acc ->
        {_, cur_acc} = collect_local_types_acc(arg, te, ftm, cur_acc)
        cur_acc
      end)

    {te, acc}
  end

  defp collect_local_types_acc(_, te, _ftm, acc), do: {te, acc}

  # Generate WAT for an expression, with target type for coercion.
  #
  # Parameters:
  #   - expr: the IR expression node
  #   - target_type: the WASM type the result should be (:i32, :i64, :f64)
  #   - ft: the current function's %IR.FunctionType{}
  #   - type_env: map of variable name → WASM type (for locals/params)
  #   - func_types: map of function name → %IR.FunctionType{} (for calls)

  # Note: 3-arity generate_expr/3 removed — all callers use generate_expr/5

  defp generate_expr({:literal, n}, target_type, _ft, _te, _ftm) when is_integer(n) do
    case target_type do
      :i32 -> "i32.const #{n}"
      :f64 -> "f64.const #{n}.0"
      _ -> "i64.const #{n}"
    end
  end

  defp generate_expr({:literal, n}, target_type, _ft, _te, _ftm) when is_float(n) do
    case target_type do
      :i32 -> "i32.const #{trunc(n)}"
      :i64 -> "i64.const #{trunc(n)}"
      _ -> "f64.const #{n}"
    end
  end

  defp generate_expr({:var, name}, target_type, _ft, te, _ftm) do
    var_type = Map.get(te, name, :i64)
    wat = "local.get $#{name}"
    coerce(wat, var_type, target_type)
  end

  # --- Peephole: i64.eqz for zero-equality checks ---
  # `expr == 0` → `expr; i64.eqz` (saves one i64.const 0 instruction)
  # `0 == expr` → `expr; i64.eqz` (commutative)
  # Only applies when the non-zero operand is i64 (not f64).
  defp generate_expr({:binop, :eq, expr, {:literal, 0}}, target_type, ft, te, ftm)
       when not is_float(0) do
    expr_type = infer_expr_wasm_type(expr, te, ftm)

    if expr_type in [:i64, :i32] do
      inner = generate_expr(expr, :i64, ft, te, ftm)
      result = "#{inner}\ni64.eqz"
      coerce(result, :i32, target_type)
    else
      generate_binop_expr(:eq, expr, {:literal, 0}, target_type, ft, te, ftm)
    end
  end

  defp generate_expr({:binop, :eq, {:literal, 0}, expr}, target_type, ft, te, ftm)
       when not is_float(0) do
    expr_type = infer_expr_wasm_type(expr, te, ftm)

    if expr_type in [:i64, :i32] do
      inner = generate_expr(expr, :i64, ft, te, ftm)
      result = "#{inner}\ni64.eqz"
      coerce(result, :i32, target_type)
    else
      generate_binop_expr(:eq, {:literal, 0}, expr, target_type, ft, te, ftm)
    end
  end

  # --- Peephole: i64.eqz + i32.eqz for nonzero checks ---
  # `expr != 0` → `expr; i64.eqz; i32.eqz` isn't better (same instruction count).
  # However, for i32 target type, `expr; i64.eqz; i32.eqz` can avoid a coerce.
  # We skip this pattern and let the normal path handle it.

  defp generate_expr({:binop, op, left, right}, target_type, ft, te, ftm) do
    generate_binop_expr(op, left, right, target_type, ft, te, ftm)
  end

  defp generate_expr({:unaryop, :not, expr}, target_type, ft, te, ftm) do
    inner = generate_expr(expr, :i32, ft, te, ftm)
    result = "#{inner}\ni32.eqz"
    coerce(result, :i32, target_type)
  end

  defp generate_expr({:unaryop, :negate, expr}, target_type, ft, te, ftm) do
    expr_type = infer_expr_wasm_type(expr, te, ftm)

    if expr_type == :f64 do
      inner = generate_expr(expr, :f64, ft, te, ftm)
      result = "#{inner}\nf64.neg"
      coerce(result, :f64, target_type)
    else
      # negate as 0 - x
      inner = generate_expr(expr, :i64, ft, te, ftm)
      result = "i64.const 0\n#{inner}\ni64.sub"
      coerce(result, :i64, target_type)
    end
  end

  # Integer counting intrinsics — map to single WASM instructions
  defp generate_expr({:unaryop, op, expr}, target_type, ft, te, ftm)
       when op in [:clz, :ctz, :popcnt] do
    inner = generate_expr(expr, :i64, ft, te, ftm)
    wasm_op = "i64.#{op}"
    result = "#{inner}\n#{wasm_op}"
    coerce(result, :i64, target_type)
  end

  # Float math intrinsics — map to single WASM f64.* instructions
  defp generate_expr({:unaryop, op, expr}, target_type, ft, te, ftm)
       when op in [:f_abs, :f_ceil, :f_floor, :f_trunc, :f_nearest, :f_sqrt] do
    inner = generate_expr(expr, :f64, ft, te, ftm)

    wasm_op =
      case op do
        :f_abs -> "f64.abs"
        :f_ceil -> "f64.ceil"
        :f_floor -> "f64.floor"
        :f_trunc -> "f64.trunc"
        :f_nearest -> "f64.nearest"
        :f_sqrt -> "f64.sqrt"
      end

    result = "#{inner}\n#{wasm_op}"
    coerce(result, :f64, target_type)
  end

  # Type conversion: to_float(x) → f64.convert_i64_s
  defp generate_expr({:unaryop, :to_float, expr}, target_type, ft, te, ftm) do
    inner = generate_expr(expr, :i64, ft, te, ftm)
    result = "#{inner}\nf64.convert_i64_s"
    coerce(result, :f64, target_type)
  end

  # Type conversion: trunc_to_int(x) → i64.trunc_f64_s
  defp generate_expr({:unaryop, :trunc_to_int, expr}, target_type, ft, te, ftm) do
    inner = generate_expr(expr, :f64, ft, te, ftm)
    result = "#{inner}\ni64.trunc_f64_s"
    coerce(result, :i64, target_type)
  end

  # Type conversion: round_to_int(x) → f64.nearest + i64.trunc_f64_s
  defp generate_expr({:unaryop, :round_to_int, expr}, target_type, ft, te, ftm) do
    inner = generate_expr(expr, :f64, ft, te, ftm)
    result = "#{inner}\nf64.nearest\ni64.trunc_f64_s"
    coerce(result, :i64, target_type)
  end

  # Bitwise NOT: bnot(x) = x XOR -1 (all bits flipped)
  defp generate_expr({:unaryop, :bnot, expr}, target_type, ft, te, ftm) do
    inner = generate_expr(expr, :i64, ft, te, ftm)
    result = "#{inner}\ni64.const -1\ni64.xor"
    coerce(result, :i64, target_type)
  end

  defp generate_expr({:call, name, args}, target_type, ft, te, ftm) do
    # Look up the callee's type signature
    callee_type = Map.get(ftm, name)

    args_wat =
      if callee_type do
        args
        |> Enum.zip(callee_type.params)
        |> Enum.map(fn {arg, param_type} -> generate_expr(arg, param_type, ft, te, ftm) end)
        |> Enum.join("\n")
      else
        args
        |> Enum.map(fn arg -> generate_expr(arg, :i64, ft, te, ftm) end)
        |> Enum.join("\n")
      end

    return_type = if callee_type, do: callee_type.return, else: :i64

    result =
      if args_wat == "" do
        "call $#{name}"
      else
        "#{args_wat}\ncall $#{name}"
      end

    coerce(result, return_type, target_type)
  end

  defp generate_expr({:if, condition, then_body, else_body}, target_type, ft, te, ftm) do
    # Use branchless `select` for simple if/else where both branches are
    # trivial expressions (literals or variables). This avoids branch
    # misprediction penalties and produces smaller WASM output.
    if safe_for_select?(then_body) and safe_for_select?(else_body) do
      then_wat = generate_expr(then_body, target_type, ft, te, ftm)
      else_wat = generate_expr(else_body, target_type, ft, te, ftm)
      cond_wat = generate_expr(condition, :i32, ft, te, ftm)

      # select takes [then_val, else_val, cond] and returns then_val if cond != 0
      "#{then_wat}\n#{else_wat}\n#{cond_wat}\nselect"
    else
      cond_wat = generate_expr(condition, :i32, ft, te, ftm)
      then_wat = generate_expr(then_body, target_type, ft, te, ftm)
      else_wat = generate_expr(else_body, target_type, ft, te, ftm)

      result_type_str = wasm_type(target_type)

      """
      #{cond_wat}
      if (result #{result_type_str})
        #{then_wat}
      else
        #{else_wat}
      end\
      """
    end
  end

  defp generate_expr({:case, subject, clauses}, target_type, ft, te, ftm) do
    subject_local = "$__case_subject"
    subject_wat = generate_expr(subject, :i64, ft, te, ftm)

    clauses_wat =
      case maybe_br_table(clauses, subject_local, target_type, ft, te, ftm, :expr) do
        {:ok, wat} -> wat
        :fallback -> generate_case_clauses(clauses, subject_local, target_type, ft, te, ftm)
      end

    """
    #{subject_wat}
    local.set #{subject_local}
    #{clauses_wat}\
    """
  end

  defp generate_expr({:tail_loop, _loop_params, body}, target_type, ft, te, ftm) do
    body_wat = generate_tco_body(body, target_type, ft, te, ftm)
    result_type = wasm_type(target_type)

    """
    (block $__tco_exit (result #{result_type})
    (loop $__tco_loop
    #{body_wat}
    br $__tco_loop
    )
    unreachable
    )\
    """
  end

  defp generate_expr({:let, name, value}, _target_type, ft, te, ftm) do
    # Determine the type for the let binding value
    val_type = infer_expr_wasm_type(value, te, ftm)
    store_type = if val_type == :f64, do: :f64, else: :i64
    value_wat = generate_expr(value, store_type, ft, te, ftm)
    "#{value_wat}\nlocal.tee $#{name}"
  end

  defp generate_expr({:block, exprs}, target_type, ft, te, ftm) do
    {init, [last]} = Enum.split(exprs, -1)

    init_wat =
      init
      |> Enum.map(fn expr ->
        case expr do
          {:let, name, value} ->
            # In block init position, let bindings don't need their value
            # on the stack — use local.set directly (avoids tee + drop).
            val_type = infer_expr_wasm_type(value, te, ftm)
            store_type = if val_type == :f64, do: :f64, else: :i64
            value_wat = generate_expr(value, store_type, ft, te, ftm)
            "#{value_wat}\nlocal.set $#{name}"

          _ ->
            wat = generate_expr(expr, :i64, ft, te, ftm)
            "#{wat}\ndrop"
        end
      end)
      |> Enum.join("\n")

    last_wat = generate_expr(last, target_type, ft, te, ftm)

    if init_wat == "" do
      last_wat
    else
      "#{init_wat}\n#{last_wat}"
    end
  end

  defp generate_expr({:error_node, _}, target_type, _ft, _te, _ftm) do
    "#{wasm_type(target_type)}.const 0  ;; error: unsupported expression"
  end

  defp generate_expr(_, target_type, _ft, _te, _ftm) do
    "#{wasm_type(target_type)}.const 0  ;; error: unknown expression"
  end

  # Maximum IR node count for a branch to be eligible for branchless select.
  # Both branches are computed speculatively, so we cap size to avoid wasting
  # cycles on the discarded branch. A limit of 6 covers common patterns like
  # `x + 1`, `a * b + c`, `x &&& mask`, etc.
  @max_select_size 6

  # An expression is eligible for branchless `select` if:
  # 1. It cannot trap (no div-by-zero, no int truncation of out-of-range floats)
  # 2. It is small enough that speculative evaluation of the untaken branch is cheap
  # 3. It contains no function calls (which may be expensive or recursive)
  #
  # This extends the original simple_select_candidate? to handle arithmetic,
  # comparisons, bitwise ops, and unary operations — all of which are single
  # WASM instructions that cannot trap. The branchless `select` avoids branch
  # misprediction and produces smaller code.
  defp safe_for_select?(expr) do
    ir_node_count(expr) <= @max_select_size and cannot_trap?(expr)
  end

  # Check that an expression contains no trapping operations.
  # In WASM, the only trapping numeric instructions are:
  #   - i64.div_s/u, i64.rem_s/u (division by zero or INT_MIN/-1)
  #   - i64.trunc_f64_s/u, i32.trunc_f64_s/u (out-of-range float)
  # All other arithmetic, comparison, bitwise, and float ops are safe.
  defp cannot_trap?({:literal, _}), do: true
  defp cannot_trap?({:var, _}), do: true

  # Safe binary ops: arithmetic (except div/rem), comparisons, bitwise, float
  @safe_binops [
    :add,
    :sub,
    :mul,
    :eq,
    :ne,
    :lt_s,
    :gt_s,
    :le_s,
    :ge_s,
    :lt_u,
    :gt_u,
    :le_u,
    :ge_u,
    :and_,
    :or_,
    :shl,
    :shr_s,
    :shr_u,
    :band,
    :bor,
    :bxor,
    :rotl,
    :rotr,
    :f_min,
    :f_max,
    :f_copysign,
    :f_div
  ]

  defp cannot_trap?({:binop, op, left, right}) when op in @safe_binops do
    cannot_trap?(left) and cannot_trap?(right)
  end

  # Safe unary ops: all except trunc_to_int and round_to_int (which use i64.trunc_f64_s)
  @safe_unaryops [
    :not,
    :negate,
    :bnot,
    :clz,
    :ctz,
    :popcnt,
    :f_abs,
    :f_ceil,
    :f_floor,
    :f_trunc,
    :f_nearest,
    :f_sqrt,
    :to_float
  ]

  defp cannot_trap?({:unaryop, op, inner}) when op in @safe_unaryops do
    cannot_trap?(inner)
  end

  defp cannot_trap?(_), do: false

  # Count IR nodes (simple size metric for select eligibility)
  defp ir_node_count({:literal, _}), do: 1
  defp ir_node_count({:var, _}), do: 1
  defp ir_node_count({:binop, _, l, r}), do: 1 + ir_node_count(l) + ir_node_count(r)
  defp ir_node_count({:unaryop, _, e}), do: 1 + ir_node_count(e)
  defp ir_node_count(_), do: 100

  # --- TCO body generation ---
  # Return values use `br $__tco_exit`, tail calls use `br $__tco_loop`

  defp generate_tco_body({:tail_call, _name, args}, _target_type, ft, te, ftm) do
    param_names = ft.params |> Enum.with_index() |> Enum.map(fn {_type, i} -> "p#{i}" end)

    compute_wat =
      args
      |> Enum.zip(ft.params)
      |> Enum.with_index()
      |> Enum.map(fn {{arg, param_type}, i} ->
        arg_wat = generate_expr(arg, param_type, ft, te, ftm)
        "#{arg_wat}\nlocal.set $__tco_tmp_p#{i}"
      end)
      |> Enum.join("\n")

    copy_wat =
      param_names
      |> Enum.map(fn pname ->
        "local.get $__tco_tmp_#{pname}\nlocal.set $#{pname}"
      end)
      |> Enum.join("\n")

    """
    #{compute_wat}
    #{copy_wat}
    br $__tco_loop\
    """
  end

  defp generate_tco_body({:if, cond_expr, then_body, else_body}, target_type, ft, te, ftm) do
    cond_wat = generate_expr(cond_expr, :i32, ft, te, ftm)
    then_wat = generate_tco_body(then_body, target_type, ft, te, ftm)
    else_wat = generate_tco_body(else_body, target_type, ft, te, ftm)

    """
    #{cond_wat}
    if
      #{then_wat}
    else
      #{else_wat}
    end\
    """
  end

  # Blocks in TCO bodies: init exprs are normal, last expr may be a tail_call
  defp generate_tco_body({:block, exprs}, target_type, ft, te, ftm)
       when is_list(exprs) and length(exprs) > 0 do
    {init, [last]} = Enum.split(exprs, -1)

    init_wat =
      init
      |> Enum.map(fn expr ->
        case expr do
          {:let, name, value} ->
            val_type = infer_expr_wasm_type(value, te, ftm)
            store_type = if val_type == :f64, do: :f64, else: :i64
            value_wat = generate_expr(value, store_type, ft, te, ftm)
            "#{value_wat}\nlocal.set $#{name}"

          _ ->
            wat = generate_expr(expr, :i64, ft, te, ftm)
            "#{wat}\ndrop"
        end
      end)
      |> Enum.join("\n")

    # Last expression gets TCO treatment (may contain tail_call or nested if/block)
    last_wat = generate_tco_body(last, target_type, ft, te, ftm)

    if init_wat == "" do
      last_wat
    else
      "#{init_wat}\n#{last_wat}"
    end
  end

  # Case expressions in TCO bodies: clause bodies may contain tail calls
  defp generate_tco_body({:case, subject, clauses}, target_type, ft, te, ftm) do
    subject_local = "$__case_subject"
    subject_wat = generate_expr(subject, :i64, ft, te, ftm)

    clauses_wat =
      case maybe_br_table(clauses, subject_local, target_type, ft, te, ftm, :tco) do
        {:ok, wat} -> wat
        :fallback -> generate_tco_case_clauses(clauses, subject_local, target_type, ft, te, ftm)
      end

    """
    #{subject_wat}
    local.set #{subject_local}
    #{clauses_wat}\
    """
  end

  defp generate_tco_body(expr, target_type, ft, te, ftm) do
    result_wat = generate_expr(expr, target_type, ft, te, ftm)

    """
    #{result_wat}
    br $__tco_exit\
    """
  end

  # --- br_table optimization for dense integer case expressions ---
  #
  # When a case has all integer literal patterns (with an optional wildcard/var
  # default), and the range is dense enough (max - min < 2 * count), we emit a
  # WASM `br_table` instruction for O(1) dispatch instead of a linear if/else
  # chain.  This is the same optimization that C compilers use for switch
  # statements with contiguous cases.
  #
  # Structure of the emitted code:
  #
  #   (block $default             ;; outermost — target for default/fallthrough
  #     (block $case_0            ;; one nested block per literal pattern
  #       (block $case_1
  #         ...
  #         local.get $subject
  #         i64.const <min>
  #         i64.sub               ;; normalize to 0-based index
  #         i32.wrap_i64
  #         br_table $case_0 $case_1 ... $default
  #       )
  #       ;; body for case_0 ...  br $exit
  #     )
  #     ;; body for case_1 ...    br $exit
  #   )
  #   ;; default body

  @br_table_density_factor 2

  defp maybe_br_table(clauses, subject_local, target_type, ft, te, ftm, mode) do
    {literal_clauses, default_clause} = partition_case_clauses(clauses)

    count = length(literal_clauses)

    if count >= 3 do
      values = Enum.map(literal_clauses, fn {val, _body} -> val end)
      min_val = Enum.min(values)
      max_val = Enum.max(values)
      range = max_val - min_val + 1

      if range <= count * @br_table_density_factor and range <= 256 and min_val >= 0 do
        {:ok,
         gen_br_table(
           literal_clauses,
           default_clause,
           min_val,
           range,
           subject_local,
           target_type,
           ft,
           te,
           ftm,
           mode
         )}
      else
        :fallback
      end
    else
      :fallback
    end
  end

  # Partition clauses into {integer_value, body} pairs and an optional default body.
  # Clauses with guards cannot use br_table (they need conditional checks),
  # so any guarded clause forces a fallback to the if/else chain.
  defp partition_case_clauses(clauses) do
    has_guards = Enum.any?(clauses, fn {_pat, guard, _body} -> guard != nil end)

    if has_guards do
      # Fall back to if/else chain when any clause has a guard
      {[], nil}
    else
      Enum.reduce(Enum.reverse(clauses), {[], nil}, fn {pat, _guard, body}, {lits, default} ->
        case pat do
          {:literal_pat, val} when is_integer(val) ->
            {[{val, body} | lits], default}

          {:var_pat, _} ->
            {lits, body}

          :wildcard ->
            {lits, body}

          _ ->
            # Non-integer literal — can't use br_table
            {lits, default}
        end
      end)
    end
  end

  # Generate a br_table dispatch for dense integer case expressions.
  #
  # WASM br_table structure (for case with patterns 0, 1, 2 + default):
  #
  #   (block $__jt_exit (result i64)      ;; expr mode: typed result block
  #     (block $__jt_default              ;; default landing
  #       (block $__jt_2                  ;; clause 2 landing
  #         (block $__jt_1                ;; clause 1 landing
  #           (block $__jt_0              ;; clause 0 landing (innermost)
  #             subject - min_val
  #             i32.wrap_i64
  #             br_table $__jt_0 $__jt_1 $__jt_2 $__jt_default
  #           )                           ;; br to $__jt_0 lands here
  #           body_0
  #           br $__jt_exit
  #         )                             ;; br to $__jt_1 lands here
  #         body_1
  #         br $__jt_exit
  #       )                               ;; br to $__jt_2 lands here
  #       body_2
  #       br $__jt_exit
  #     )                                 ;; br to $__jt_default lands here
  #     default_body
  #   )
  #
  # br_table index 0 → innermost block ($__jt_0), index N → $__jt_default.
  defp gen_br_table(
         literal_clauses,
         default_body,
         min_val,
         _range,
         subject_local,
         target_type,
         ft,
         te,
         ftm,
         mode
       ) do
    value_to_idx =
      literal_clauses
      |> Enum.with_index()
      |> Map.new(fn {{val, _body}, idx} -> {val, idx} end)

    n = length(literal_clauses)
    result_type_str = wasm_type(target_type)

    # Generate body WAT for each literal clause
    clause_bodies =
      Enum.map(literal_clauses, fn {_val, body} ->
        case mode do
          :expr -> generate_expr(body, target_type, ft, te, ftm)
          :tco -> generate_tco_body(body, target_type, ft, te, ftm)
        end
      end)

    # Generate default body WAT
    default_body_wat =
      case {default_body, mode} do
        {nil, :expr} -> "#{wasm_type(target_type)}.const 0  ;; unreachable default"
        {nil, :tco} -> "unreachable  ;; no matching case clause"
        {body, :expr} -> generate_expr(body, target_type, ft, te, ftm)
        {body, :tco} -> generate_tco_body(body, target_type, ft, te, ftm)
      end

    # Build the br_table label list for each slot in [min_val .. max_val]
    max_val = literal_clauses |> Enum.map(fn {v, _} -> v end) |> Enum.max()
    range = max_val - min_val + 1

    br_labels =
      0..(range - 1)
      |> Enum.map(fn offset ->
        val = min_val + offset

        case Map.get(value_to_idx, val) do
          nil -> "$__jt_default"
          idx -> "$__jt_#{idx}"
        end
      end)
      |> Enum.join(" ")

    # Dispatch: normalize subject to 0-based index, then br_table
    subtract_wat = if min_val == 0, do: "", else: "i64.const #{min_val}\ni64.sub\n"

    dispatch =
      String.trim("""
      local.get #{subject_local}
      #{subtract_wat}i32.wrap_i64
      br_table #{br_labels} $__jt_default\
      """)

    # Build from innermost outward: start with dispatch, wrap in blocks,
    # interleave clause bodies after each block close.
    exit_br =
      case mode do
        :expr -> "\nbr $__jt_exit"
        :tco -> ""
      end

    result =
      Enum.reduce(0..(n - 1), dispatch, fn idx, code ->
        body_wat = Enum.at(clause_bodies, idx)
        "(block $__jt_#{idx}\n#{code}\n)\n#{body_wat}#{exit_br}"
      end)

    # Wrap in $__jt_default block, add default body
    result = "(block $__jt_default\n#{result}\n)\n#{default_body_wat}#{exit_br}"

    # For expr mode, wrap in typed $__jt_exit block
    case mode do
      :expr -> "(block $__jt_exit (result #{result_type_str})\n#{result}\n)"
      :tco -> result
    end
  end

  # Generate case clause chain as nested if/else
  defp generate_case_clauses([], _subject, target_type, _ft, _te, _ftm) do
    "#{wasm_type(target_type)}.const 0  ;; unreachable"
  end

  defp generate_case_clauses([{pattern, guard, body} | rest], subject, target_type, ft, te, ftm) do
    case pattern do
      {:literal_pat, val} ->
        body_wat = generate_expr(body, target_type, ft, te, ftm)
        else_wat = generate_case_clauses(rest, subject, target_type, ft, te, ftm)

        # Peephole: use i64.eqz for pattern matching against 0
        eq_wat =
          if val == 0 do
            "local.get #{subject}\ni64.eqz"
          else
            "local.get #{subject}\ni64.const #{val}\ni64.eq"
          end

        cond_wat =
          case guard do
            nil ->
              eq_wat

            guard_expr ->
              guard_wat = generate_expr(guard_expr, :i32, ft, te, ftm)
              "#{eq_wat}\n#{guard_wat}\ni32.and"
          end

        """
        #{cond_wat}
        if (result #{wasm_type(target_type)})
          #{body_wat}
        else
          #{else_wat}
        end\
        """

      {:var_pat, var_name} ->
        # Variable pattern always matches the pattern, but may have a guard
        case guard do
          nil ->
            generate_expr(body, target_type, ft, te, ftm)

          guard_expr ->
            # Bind the subject to the variable for use in guard and body
            # The guard references the variable, which is the case subject
            te_with_var = Map.put(te, var_name, :i64)
            guard_wat = generate_expr(guard_expr, :i32, ft, te_with_var, ftm)
            body_wat = generate_expr(body, target_type, ft, te_with_var, ftm)
            else_wat = generate_case_clauses(rest, subject, target_type, ft, te, ftm)

            """
            local.get #{subject}
            local.set $#{var_name}
            #{guard_wat}
            if (result #{wasm_type(target_type)})
              #{body_wat}
            else
              #{else_wat}
            end\
            """
        end

      :wildcard ->
        # Wildcard with guard
        case guard do
          nil ->
            generate_expr(body, target_type, ft, te, ftm)

          guard_expr ->
            guard_wat = generate_expr(guard_expr, :i32, ft, te, ftm)
            body_wat = generate_expr(body, target_type, ft, te, ftm)
            else_wat = generate_case_clauses(rest, subject, target_type, ft, te, ftm)

            """
            #{guard_wat}
            if (result #{wasm_type(target_type)})
              #{body_wat}
            else
              #{else_wat}
            end\
            """
        end
    end
  end

  # Generate case clause chain for TCO bodies — clause bodies use generate_tco_body
  # so tail calls inside case branches correctly emit loop branches.
  defp generate_tco_case_clauses([], _subject, _target_type, _ft, _te, _ftm) do
    "unreachable  ;; no matching case clause"
  end

  defp generate_tco_case_clauses(
         [{pattern, guard, body} | rest],
         subject,
         target_type,
         ft,
         te,
         ftm
       ) do
    case pattern do
      {:literal_pat, val} ->
        body_wat = generate_tco_body(body, target_type, ft, te, ftm)
        else_wat = generate_tco_case_clauses(rest, subject, target_type, ft, te, ftm)

        # Peephole: use i64.eqz for pattern matching against 0
        eq_wat =
          if val == 0 do
            "local.get #{subject}\ni64.eqz"
          else
            "local.get #{subject}\ni64.const #{val}\ni64.eq"
          end

        cond_wat =
          case guard do
            nil ->
              eq_wat

            guard_expr ->
              guard_wat = generate_expr(guard_expr, :i32, ft, te, ftm)
              "#{eq_wat}\n#{guard_wat}\ni32.and"
          end

        """
        #{cond_wat}
        if
          #{body_wat}
        else
          #{else_wat}
        end\
        """

      {:var_pat, var_name} ->
        case guard do
          nil ->
            generate_tco_body(body, target_type, ft, te, ftm)

          guard_expr ->
            te_with_var = Map.put(te, var_name, :i64)
            guard_wat = generate_expr(guard_expr, :i32, ft, te_with_var, ftm)
            body_wat = generate_tco_body(body, target_type, ft, te_with_var, ftm)
            else_wat = generate_tco_case_clauses(rest, subject, target_type, ft, te, ftm)

            """
            local.get #{subject}
            local.set $#{var_name}
            #{guard_wat}
            if
              #{body_wat}
            else
              #{else_wat}
            end\
            """
        end

      :wildcard ->
        case guard do
          nil ->
            generate_tco_body(body, target_type, ft, te, ftm)

          guard_expr ->
            guard_wat = generate_expr(guard_expr, :i32, ft, te, ftm)
            body_wat = generate_tco_body(body, target_type, ft, te, ftm)
            else_wat = generate_tco_case_clauses(rest, subject, target_type, ft, te, ftm)

            """
            #{guard_wat}
            if
              #{body_wat}
            else
              #{else_wat}
            end\
            """
        end
    end
  end

  # --- Expression type inference for WAT generation ---
  # Determines the natural WASM type of an IR expression without generating code.
  # Used to select i64 vs f64 instructions for arithmetic operations.

  defp infer_expr_wasm_type({:literal, n}, _te, _ftm) when is_integer(n), do: :i64
  defp infer_expr_wasm_type({:literal, n}, _te, _ftm) when is_float(n), do: :f64
  defp infer_expr_wasm_type({:var, name}, te, _ftm), do: Map.get(te, name, :i64)

  defp infer_expr_wasm_type({:binop, op, _l, _r}, _te, _ftm)
       when op in [:eq, :ne, :lt_s, :gt_s, :le_s, :ge_s, :lt_u, :gt_u, :le_u, :ge_u, :and_, :or_],
       do: :i32

  # Bitwise operations always produce i64
  defp infer_expr_wasm_type({:binop, op, _l, _r}, _te, _ftm)
       when op in [:band, :bor, :bxor, :shl, :shr_s, :shr_u, :rotl, :rotr],
       do: :i64

  # Unsigned arithmetic always produces i64
  defp infer_expr_wasm_type({:binop, op, _l, _r}, _te, _ftm)
       when op in [:div_u, :rem_u],
       do: :i64

  # Float binary intrinsics always return f64
  defp infer_expr_wasm_type({:binop, op, _l, _r}, _te, _ftm)
       when op in [:f_min, :f_max, :f_copysign, :f_div],
       do: :f64

  defp infer_expr_wasm_type({:binop, _op, left, right}, te, ftm) do
    widen_numeric(infer_expr_wasm_type(left, te, ftm), infer_expr_wasm_type(right, te, ftm))
  end

  defp infer_expr_wasm_type({:unaryop, :not, _}, _te, _ftm), do: :i32
  defp infer_expr_wasm_type({:unaryop, :bnot, _}, _te, _ftm), do: :i64
  # Integer counting intrinsics always return i64
  defp infer_expr_wasm_type({:unaryop, op, _}, _te, _ftm)
       when op in [:clz, :ctz, :popcnt],
       do: :i64

  # Float math intrinsics always return f64
  defp infer_expr_wasm_type({:unaryop, op, _}, _te, _ftm)
       when op in [:f_abs, :f_ceil, :f_floor, :f_trunc, :f_nearest, :f_sqrt],
       do: :f64

  # Type conversion intrinsics
  defp infer_expr_wasm_type({:unaryop, :to_float, _}, _te, _ftm), do: :f64
  defp infer_expr_wasm_type({:unaryop, :trunc_to_int, _}, _te, _ftm), do: :i64
  defp infer_expr_wasm_type({:unaryop, :round_to_int, _}, _te, _ftm), do: :i64

  defp infer_expr_wasm_type({:unaryop, :negate, expr}, te, ftm),
    do: infer_expr_wasm_type(expr, te, ftm)

  defp infer_expr_wasm_type({:call, name, _args}, _te, ftm) do
    case Map.get(ftm, name) do
      %{return: ret} -> ret
      _ -> :i64
    end
  end

  defp infer_expr_wasm_type({:if, _, then_b, _}, te, ftm),
    do: infer_expr_wasm_type(then_b, te, ftm)

  defp infer_expr_wasm_type({:case, _, [{_, _, body} | _]}, te, ftm),
    do: infer_expr_wasm_type(body, te, ftm)

  defp infer_expr_wasm_type({:let, _, value}, te, ftm), do: infer_expr_wasm_type(value, te, ftm)

  defp infer_expr_wasm_type({:block, exprs}, te, ftm) when is_list(exprs) and length(exprs) > 0 do
    infer_expr_wasm_type(List.last(exprs), te, ftm)
  end

  defp infer_expr_wasm_type({:tail_loop, _, body}, te, ftm),
    do: infer_expr_wasm_type(body, te, ftm)

  defp infer_expr_wasm_type({:tail_call, name, _}, _te, ftm) do
    case Map.get(ftm, name) do
      %{return: ret} -> ret
      _ -> :i64
    end
  end

  defp infer_expr_wasm_type(_, _te, _ftm), do: :i64

  # Widen two numeric types: if either is f64, result is f64
  defp widen_numeric(:f64, _), do: :f64
  defp widen_numeric(_, :f64), do: :f64
  defp widen_numeric(a, _), do: a

  # General binary operation code generation (extracted for peephole fallback)
  defp generate_binop_expr(op, left, right, target_type, ft, te, ftm) do
    if op in [:and_, :or_] do
      # Boolean ops always use i32
      {operand_type, result_type, wasm_op} = binop_info(op)
      left_wat = generate_expr(left, operand_type, ft, te, ftm)
      right_wat = generate_expr(right, operand_type, ft, te, ftm)
      result = "#{left_wat}\n#{right_wat}\n#{wasm_op}"
      coerce(result, result_type, target_type)
    else
      # Determine the natural type of each operand to select i64 vs f64 ops
      left_type = infer_expr_wasm_type(left, te, ftm)
      right_type = infer_expr_wasm_type(right, te, ftm)
      effective_type = widen_numeric(left_type, right_type)

      {operand_type, result_type, wasm_op} = binop_info(op, effective_type)

      left_wat = generate_expr(left, operand_type, ft, te, ftm)
      right_wat = generate_expr(right, operand_type, ft, te, ftm)

      result = "#{left_wat}\n#{right_wat}\n#{wasm_op}"

      coerce(result, result_type, target_type)
    end
  end

  # Binary operation info: {operand_type, result_type, wasm_instruction}
  # Type-aware version: selects i64 or f64 instructions based on effective_type
  defp binop_info(op, effective_type \\ :i64)

  # Arithmetic operations - type-aware
  defp binop_info(:add, :f64), do: {:f64, :f64, "f64.add"}
  defp binop_info(:add, _), do: {:i64, :i64, "i64.add"}
  defp binop_info(:sub, :f64), do: {:f64, :f64, "f64.sub"}
  defp binop_info(:sub, _), do: {:i64, :i64, "i64.sub"}
  defp binop_info(:mul, :f64), do: {:f64, :f64, "f64.mul"}
  defp binop_info(:mul, _), do: {:i64, :i64, "i64.mul"}
  defp binop_info(:div_s, :f64), do: {:f64, :f64, "f64.div"}
  defp binop_info(:div_s, _), do: {:i64, :i64, "i64.div_s"}
  # Float division (/) — always operates on f64, regardless of operand types.
  # Integer operands are coerced to f64 by the operand_type :f64.
  defp binop_info(:f_div, _), do: {:f64, :f64, "f64.div"}
  defp binop_info(:rem_s, _), do: {:i64, :i64, "i64.rem_s"}

  # Comparison operations - type-aware
  defp binop_info(:eq, :f64), do: {:f64, :i32, "f64.eq"}
  defp binop_info(:eq, _), do: {:i64, :i32, "i64.eq"}
  defp binop_info(:ne, :f64), do: {:f64, :i32, "f64.ne"}
  defp binop_info(:ne, _), do: {:i64, :i32, "i64.ne"}
  defp binop_info(:lt_s, :f64), do: {:f64, :i32, "f64.lt"}
  defp binop_info(:lt_s, _), do: {:i64, :i32, "i64.lt_s"}
  defp binop_info(:gt_s, :f64), do: {:f64, :i32, "f64.gt"}
  defp binop_info(:gt_s, _), do: {:i64, :i32, "i64.gt_s"}
  defp binop_info(:le_s, :f64), do: {:f64, :i32, "f64.le"}
  defp binop_info(:le_s, _), do: {:i64, :i32, "i64.le_s"}
  defp binop_info(:ge_s, :f64), do: {:f64, :i32, "f64.ge"}
  defp binop_info(:ge_s, _), do: {:i64, :i32, "i64.ge_s"}

  # Bitwise/shift operations - always i64
  # Unsigned arithmetic
  defp binop_info(:div_u, _), do: {:i64, :i64, "i64.div_u"}
  defp binop_info(:rem_u, _), do: {:i64, :i64, "i64.rem_u"}
  defp binop_info(:shr_u, _), do: {:i64, :i64, "i64.shr_u"}

  # Unsigned comparisons
  defp binop_info(:lt_u, _), do: {:i64, :i32, "i64.lt_u"}
  defp binop_info(:gt_u, _), do: {:i64, :i32, "i64.gt_u"}
  defp binop_info(:le_u, _), do: {:i64, :i32, "i64.le_u"}
  defp binop_info(:ge_u, _), do: {:i64, :i32, "i64.ge_u"}

  # Bitwise/shift operations - always i64
  defp binop_info(:shl, _), do: {:i64, :i64, "i64.shl"}
  defp binop_info(:shr_s, _), do: {:i64, :i64, "i64.shr_s"}
  defp binop_info(:band, _), do: {:i64, :i64, "i64.and"}
  defp binop_info(:bor, _), do: {:i64, :i64, "i64.or"}
  defp binop_info(:bxor, _), do: {:i64, :i64, "i64.xor"}
  defp binop_info(:rotl, _), do: {:i64, :i64, "i64.rotl"}
  defp binop_info(:rotr, _), do: {:i64, :i64, "i64.rotr"}

  # Float binary intrinsics - always f64
  defp binop_info(:f_min, _), do: {:f64, :f64, "f64.min"}
  defp binop_info(:f_max, _), do: {:f64, :f64, "f64.max"}
  defp binop_info(:f_copysign, _), do: {:f64, :f64, "f64.copysign"}

  # Boolean operations - always i32
  defp binop_info(:and_, _), do: {:i32, :i32, "i32.and"}
  defp binop_info(:or_, _), do: {:i32, :i32, "i32.or"}

  # Type coercion between WASM types
  defp coerce(wat, same, same), do: wat

  defp coerce(wat, :i32, :i64), do: "#{wat}\ni64.extend_i32_s"
  defp coerce(wat, :i64, :i32), do: "#{wat}\ni32.wrap_i64"
  defp coerce(wat, :i32, :f64), do: "#{wat}\nf64.convert_i32_s"
  defp coerce(wat, :i64, :f64), do: "#{wat}\nf64.convert_i64_s"
  defp coerce(wat, :f64, :i64), do: "#{wat}\ni64.trunc_f64_s"
  defp coerce(wat, :f64, :i32), do: "#{wat}\ni32.trunc_f64_s"
  defp coerce(wat, _, _), do: wat

  defp wasm_type(:i32), do: "i32"
  defp wasm_type(:i64), do: "i64"
  defp wasm_type(:f64), do: "f64"
  defp wasm_type(_), do: "i64"

  defp indent(text, spaces) do
    pad = String.duplicate(" ", spaces)

    text
    |> String.split("\n")
    |> Enum.map(fn line ->
      if String.trim(line) == "" do
        ""
      else
        pad <> line
      end
    end)
    |> Enum.join("\n")
  end
end