eyeon

defmodule Eyeon.Macro do
  @moduledoc false

  alias Eyeon.Macro.{Context, Helpers, System}

  @system_symbols ~w($ion $ion_1_0 $ion_symbol_table name version imports symbols max_id $ion_shared_symbol_table $ion_encoding)

  @spec expand(term()) :: {:ok, term()} | {:error, String.t()}
  def expand(ast) when is_list(ast) do
    case expand_top_level(ast) do
      {:error, _} = err -> err
      expanded -> {:ok, expanded}
    end
  end

  def expand(other), do: {:ok, other}

  # Public for callback from System.parse_ion
  @doc false
  def expand_top_level(nodes), do: expand_top_level(nodes, %Context{})

  @doc false
  def expand_top_level(nodes, ctx) do
    flat = Enum.flat_map(nodes, &flatten_compound/1)

    Enum.reduce_while(flat, {[], ctx}, fn node, {acc, ctx} ->
      case expand_node(node, ctx) do
        {:error, _} = err ->
          {:halt, err}

        {:directive, :set_symbols, symbols} ->
          {:cont, {acc, %{ctx | symbols: symbols}}}

        {:directive, :add_symbols, symbols} ->
          {:cont, {acc, %{ctx | symbols: ctx.symbols ++ symbols}}}

        {:directive, :set_macros, macro_defs} ->
          macros =
            macro_defs
            |> Enum.with_index()
            |> Enum.reduce(%{}, fn {m, i}, acc ->
              acc |> Map.put(m.name, m) |> Map.put(i, m)
            end)

          {:cont, {acc, %{ctx | macros: macros, macro_count: length(macro_defs)}}}

        {:directive, :add_macros, macro_defs} ->
          macros =
            macro_defs
            |> Enum.with_index(ctx.macro_count)
            |> Enum.reduce(ctx.macros, fn {m, i}, acc ->
              acc |> Map.put(m.name, m) |> Map.put(i, m)
            end)

          {:cont,
           {acc, %{ctx | macros: macros, macro_count: ctx.macro_count + length(macro_defs)}}}

        {:splice, values} ->
          {:cont, {acc ++ values, ctx}}

        expanded ->
          {:cont, {acc ++ [expanded], ctx}}
      end
    end)
    |> case do
      {:error, _} = err -> err
      {acc, _ctx} -> acc
    end
  end

  # Public for callback from System.parse_ion
  @doc false
  def flatten_compound({first, second})
      when is_tuple(first) and tuple_size(first) == 3 do
    case first do
      {tag, line, _} when is_atom(tag) and is_integer(line) ->
        flatten_compound(first) ++ flatten_compound(second)

      _ ->
        [{first, second}]
    end
  end

  def flatten_compound(node), do: [node]

  # Macro invocation node
  defp expand_node({:macro, name_token, args}, ctx) do
    name = Helpers.extract_macro_name(name_token, ctx)

    # Check user macros first so we can choose the right arg expansion strategy
    case Map.get(ctx.macros, name) do
      nil ->
        # System macro — flatten expression groups into positional args
        case expand_args(args, ctx) do
          {:error, _} = err ->
            err

          expanded_args ->
            case System.dispatch(name, expanded_args, name_token) do
              {:error, _} = err -> err
              result -> result
            end
        end

      macro_def ->
        # User macro — preserve expression groups as single args per parameter
        case expand_args_for_user_macro(args, ctx) do
          {:error, _} = err -> err
          expanded_args -> expand_user_macro(macro_def, expanded_args, ctx)
        end
    end
  end

  # Symbol ID resolution against user symbol table
  defp expand_node({:symbol_id, line, chars}, ctx) when ctx.symbols != [] do
    id =
      chars
      |> List.to_string()
      |> String.trim_leading("$")
      |> String.to_integer()

    user_count = length(ctx.symbols)

    cond do
      id == 0 ->
        {:symbol_id, line, chars}

      id >= 1 and id <= user_count ->
        name = Enum.at(ctx.symbols, id - 1)
        {:id_symbol, line, String.to_charlist(name)}

      id > user_count and id <= user_count + length(@system_symbols) ->
        sys_idx = id - user_count - 1
        name = Enum.at(@system_symbols, sys_idx)
        {:id_symbol, line, String.to_charlist(name)}

      true ->
        {:symbol_id, line, chars}
    end
  end

  # List container
  defp expand_node({:list, items}, ctx) do
    case expand_and_flatten(items, ctx) do
      {:error, _} = err -> err
      expanded -> {:list, expanded}
    end
  end

  # Sexp containers
  defp expand_node({:sexp}, _ctx) do
    {:sexp}
  end

  defp expand_node({:sexp, values}, ctx) do
    case expand_sexp_values(values, ctx) do
      {:error, _} = err -> err
      expanded -> {:sexp, expanded}
    end
  end

  defp expand_node({:sexp, values, last}, ctx) do
    case expand_sexp_values(values, ctx) do
      {:error, _} = err ->
        err

      expanded_values ->
        case expand_node(last, ctx) do
          {:error, _} = err ->
            err

          {:splice, spliced} ->
            {:sexp, expanded_values ++ spliced}

          expanded_last ->
            {:sexp, expanded_values ++ [expanded_last]}
        end
    end
  end

  # Struct container
  defp expand_node({:struct, fields}, ctx) do
    case expand_struct_fields(fields, ctx) do
      {:error, _} = err -> err
      expanded -> {:struct, expanded}
    end
  end

  # Annotation
  defp expand_node({:annotated, annotation, value}, ctx) do
    case expand_node(value, ctx) do
      {:error, _} = err -> err
      expanded -> {:annotated, annotation, expanded}
    end
  end

  # Expression group — splice its contents
  defp expand_node({:expr_group, items}, ctx) do
    case expand_and_flatten(items, ctx) do
      {:error, _} = err -> err
      expanded -> {:splice, expanded}
    end
  end

  # Everything else (terminals, etc.) passes through unchanged
  defp expand_node(node, _ctx), do: node

  # Expand a user macro invocation
  defp expand_user_macro(macro_def, args, ctx) do
    eval_tdl(macro_def.body, args, ctx)
  end

  @doc false
  def expand_user_macro_public(macro_def, args, ctx) do
    eval_tdl(macro_def.body, args, ctx)
  end

  defp eval_tdl(nil, _args, _ctx), do: {:splice, []}
  defp eval_tdl({:splice, []}, _args, _ctx), do: {:splice, []}

  defp eval_tdl({:template_body, nodes}, args, ctx) do
    result =
      Enum.reduce_while(nodes, [], fn node, acc ->
        case eval_tdl(node, args, ctx) do
          {:error, _} = err -> {:halt, err}
          {:splice, values} -> {:cont, acc ++ values}
          expanded -> {:cont, acc ++ [expanded]}
        end
      end)

    case result do
      {:error, _} = err -> err
      values -> {:splice, values}
    end
  end

  defp eval_tdl({:param_ref, idx}, args, ctx) do
    case Enum.at(args, idx) do
      nil ->
        {:splice, []}

      {:expr_group, items} ->
        # Variadic parameter — expand each item and splice them all
        case expand_and_flatten(items, ctx) do
          {:error, _} = err -> err
          expanded -> {:splice, expanded}
        end

      arg ->
        expand_node(arg, ctx)
    end
  end

  # TDL conditionals
  defp eval_tdl({:if_none, test, true_branch, false_branch}, args, ctx) do
    case eval_tdl(test, args, ctx) do
      {:error, _} = err -> err
      {:splice, []} -> eval_tdl(true_branch, args, ctx)
      _ -> eval_tdl(false_branch, args, ctx)
    end
  end

  defp eval_tdl({:if_some, test, true_branch, false_branch}, args, ctx) do
    case eval_tdl(test, args, ctx) do
      {:error, _} = err -> err
      {:splice, []} -> eval_tdl(false_branch, args, ctx)
      _ -> eval_tdl(true_branch, args, ctx)
    end
  end

  defp eval_tdl({:if_single, test, true_branch, false_branch}, args, ctx) do
    case eval_tdl(test, args, ctx) do
      {:error, _} = err ->
        err

      {:splice, [single]} when not is_tuple(single) or elem(single, 0) != :splice ->
        eval_tdl(true_branch, args, ctx)

      result when not is_tuple(result) or elem(result, 0) != :splice ->
        eval_tdl(true_branch, args, ctx)

      _ ->
        eval_tdl(false_branch, args, ctx)
    end
  end

  defp eval_tdl({:if_multi, test, true_branch, false_branch}, args, ctx) do
    case eval_tdl(test, args, ctx) do
      {:error, _} = err -> err
      {:splice, vals} when length(vals) > 1 -> eval_tdl(true_branch, args, ctx)
      _ -> eval_tdl(false_branch, args, ctx)
    end
  end

  # TDL for — stub, returns empty splice (for loop expansion not yet supported)
  defp eval_tdl({:tdl_for, _compiled_args}, _args, _ctx) do
    {:splice, []}
  end

  # Regular AST nodes — expand normally
  defp eval_tdl(node, _args, ctx) do
    expand_node(node, ctx)
  end

  # Expand args for user macros, preserving expression groups as single args.
  # Each {:expr_group, items} becomes a single {:expr_group, expanded_items} in the result.
  # This ensures variadic parameters are passed as one arg that can be spliced by param_ref.
  defp expand_args_for_user_macro(args, ctx) do
    Enum.reduce_while(args, [], fn arg, acc ->
      case arg do
        {:expr_group, items} ->
          case expand_expr_group(items, ctx) do
            {:error, _} = err -> {:halt, err}
            expanded -> {:cont, acc ++ [{:expr_group, expanded}]}
          end

        _ ->
          case expand_node(arg, ctx) do
            {:error, _} = err -> {:halt, err}
            {:splice, values} -> {:cont, acc ++ values}
            expanded -> {:cont, acc ++ [expanded]}
          end
      end
    end)
  end

  # Expand a list of macro args, returning flat list of expanded nodes
  defp expand_args(args, ctx) do
    Enum.reduce_while(args, [], fn arg, acc ->
      case arg do
        {:expr_group, items} ->
          case expand_expr_group(items, ctx) do
            {:error, _} = err -> {:halt, err}
            expanded -> {:cont, acc ++ expanded}
          end

        _ ->
          case expand_node(arg, ctx) do
            {:error, _} = err -> {:halt, err}
            {:splice, values} -> {:cont, acc ++ values}
            expanded -> {:cont, acc ++ [expanded]}
          end
      end
    end)
  end

  # Expand items in a list/values context, flattening spliced results
  defp expand_and_flatten(items, ctx) do
    Enum.reduce_while(items, [], fn item, acc ->
      case item do
        {:expr_group, inner} ->
          case expand_expr_group(inner, ctx) do
            {:error, _} = err -> {:halt, err}
            expanded -> {:cont, acc ++ expanded}
          end

        _ ->
          case expand_node(item, ctx) do
            {:error, _} = err -> {:halt, err}
            {:splice, values} -> {:cont, acc ++ values}
            expanded -> {:cont, acc ++ [expanded]}
          end
      end
    end)
  end

  # Expression group - flatten items into argument list
  defp expand_expr_group(items, ctx) do
    Enum.reduce_while(items, [], fn item, acc ->
      case item do
        {:expr_group, inner} ->
          case expand_expr_group(inner, ctx) do
            {:error, _} = err -> {:halt, err}
            expanded -> {:cont, acc ++ expanded}
          end

        _ ->
          case expand_node(item, ctx) do
            {:error, _} = err -> {:halt, err}
            {:splice, values} -> {:cont, acc ++ values}
            expanded -> {:cont, acc ++ [expanded]}
          end
      end
    end)
  end

  # Expand sexp values
  defp expand_sexp_values(values, ctx) do
    Enum.reduce_while(values, [], fn value, acc ->
      case expand_sexp_entry(value, ctx) do
        {:error, _} = err -> {:halt, err}
        expanded_list -> {:cont, acc ++ expanded_list}
      end
    end)
  end

  defp expand_sexp_entry(value, ctx) when is_list(value) do
    case expand_and_flatten(value, ctx) do
      {:error, _} = err -> err
      expanded -> [expanded]
    end
  end

  defp expand_sexp_entry({first, second}, ctx) when not is_atom(first) do
    case expand_node(first, ctx) do
      {:error, _} = err ->
        err

      {:splice, values} ->
        case expand_sexp_entry(second, ctx) do
          {:error, _} = err -> err
          rest -> Enum.map(values, &[&1]) ++ rest
        end

      expanded_first ->
        case expand_sexp_entry(second, ctx) do
          {:error, _} = err -> err
          rest -> [[expanded_first] | rest]
        end
    end
  end

  defp expand_sexp_entry(value, ctx) do
    case expand_node(value, ctx) do
      {:error, _} = err -> err
      {:splice, values} -> Enum.map(values, &[&1])
      expanded -> [[expanded]]
    end
  end

  # Expand struct fields
  defp expand_struct_fields(fields, ctx) do
    Enum.reduce_while(fields, [], fn field, acc ->
      case expand_struct_field(field, ctx) do
        {:error, _} = err -> {:halt, err}
        expanded_fields -> {:cont, acc ++ expanded_fields}
      end
    end)
  end

  defp expand_struct_field({:macro_field, macro_node}, ctx) do
    case expand_node(macro_node, ctx) do
      {:error, _} = err ->
        err

      {:splice, values} ->
        extract_fields_from_spliced(values)

      {:struct, inner_fields} ->
        inner_fields

      _other ->
        []
    end
  end

  defp expand_struct_field({field_name, value}, ctx) do
    case expand_node(value, ctx) do
      {:error, _} = err ->
        err

      {:splice, []} ->
        []

      {:splice, [single]} ->
        [{field_name, single}]

      {:splice, multiple} ->
        Enum.map(multiple, fn v -> {field_name, v} end)

      expanded ->
        [{field_name, expanded}]
    end
  end

  defp extract_fields_from_spliced(values) do
    Enum.reduce_while(values, [], fn value, acc ->
      case value do
        {:struct, inner_fields} ->
          {:cont, acc ++ inner_fields}

        _ ->
          {:halt, {:error, "values macro in struct field position must contain structs"}}
      end
    end)
  end
end