snakebridge

defmodule SnakeBridge.Generator do
  @moduledoc """
  Generates Elixir source files from introspection data.
  """

  @spec render_library(SnakeBridge.Config.Library.t(), list(), list(), keyword()) :: String.t()
  def render_library(library, functions, classes, opts \\ []) do
    version = Keyword.get(opts, :version, Application.spec(:snakebridge, :vsn) |> to_string())
    module_name = module_to_string(library.module_name)

    header = """
    # Generated by SnakeBridge v#{version} - DO NOT EDIT MANUALLY
    # Regenerate with: mix compile
    # Library: #{library.python_name} #{library.version || "unknown"}
    """

    moduledoc = """
      @moduledoc \"\"\"
      SnakeBridge bindings for `#{library.python_name}`.
      \"\"\"

      def __snakebridge_python_name__, do: "#{library.python_name}"
      def __snakebridge_library__, do: "#{library.python_name}"
    """

    functions_by_module =
      functions
      |> Enum.map(&Map.put_new(&1, "python_module", library.python_name))
      |> Enum.group_by(& &1["python_module"])

    base_functions = Map.get(functions_by_module, library.python_name, [])

    function_defs =
      base_functions
      |> Enum.sort_by(& &1["name"])
      |> Enum.map_join("\n\n", &render_function/1)

    submodule_defs =
      functions_by_module
      |> Map.drop([library.python_name])
      |> Enum.sort_by(fn {python_module, _} -> python_module end)
      |> Enum.map_join("\n\n", fn {python_module, funcs} ->
        render_submodule(python_module, funcs)
      end)

    class_defs =
      classes
      |> Enum.sort_by(&class_sort_key/1)
      |> Enum.map_join("\n\n", &render_class(&1, library))

    discovery = render_discovery(functions, classes)

    """
    #{header}
    defmodule #{module_name} do
    #{moduledoc}
    #{function_defs}

    #{submodule_defs}

    #{class_defs}

    #{discovery}
    end
    """
    |> Code.format_string!()
    |> IO.iodata_to_binary()
  end

  @spec generate_library(SnakeBridge.Config.Library.t(), list(), list(), SnakeBridge.Config.t()) ::
          :ok
  def generate_library(library, functions, classes, config) do
    File.mkdir_p!(config.generated_dir)
    path = Path.join(config.generated_dir, "#{library.python_name}.ex")

    source =
      render_library(library, functions, classes, version: Application.spec(:snakebridge, :vsn))

    File.write!(path, source)
  end

  defp render_function(info) do
    name = info["name"]
    params = info["parameters"] || []
    doc = info["docstring"] || ""

    {param_names, has_opts} = build_params(params)
    args = "[" <> Enum.join(param_names, ", ") <> "]"
    call = runtime_call(name, args, has_opts)
    spec = function_spec(name, param_names, has_opts, "term()")

    """
      @doc \"\"\"
      #{String.trim(doc)}
      \"\"\"
      #{spec}
      def #{name}(#{param_list(param_names, has_opts)}) do
        #{call}
      end
    """
  end

  defp render_submodule(python_module, functions) do
    module_name =
      python_module
      |> String.split(".")
      |> Enum.drop(1)
      |> Enum.map_join(".", &Macro.camelize/1)

    function_defs =
      functions
      |> Enum.sort_by(& &1["name"])
      |> Enum.map_join("\n\n", &render_function/1)

    """
      defmodule #{module_name} do
        def __snakebridge_python_name__, do: "#{python_module}"

    #{indent(function_defs, 4)}
      end
    """
  end

  defp render_class(class_info, library) do
    class_name = class_name(class_info)
    python_module = class_python_module(class_info, library)
    module_name = class_module_name(class_info, library)
    relative_module = relative_module_name(library, module_name)

    methods = class_info["methods"] || []
    attrs = class_info["attributes"] || []

    methods_source =
      methods
      |> Enum.map_join("\n\n", &render_method/1)

    attrs_source =
      attrs
      |> Enum.map_join("\n\n", &render_attribute/1)

    """
      defmodule #{relative_module} do
        def __snakebridge_python_name__, do: "#{python_module}"
        def __snakebridge_python_class__, do: "#{class_name}"

        @spec new(term(), keyword()) :: {:ok, Snakepit.PyRef.t()} | {:error, Snakepit.Error.t()}
        def new(arg, opts \\\\ []) do
          SnakeBridge.Runtime.call_class(__MODULE__, :__init__, [arg], opts)
        end

    #{indent(methods_source, 4)}

    #{indent(attrs_source, 4)}
      end
    """
  end

  defp render_method(%{"name" => "__init__"}), do: ""

  defp render_method(info) do
    name = info["name"]
    params = info["parameters"] || []
    {param_names, has_opts} = build_params(params)
    spec = method_spec(name, param_names, has_opts)
    call = runtime_method_call(name, param_names, has_opts)

    """
        #{spec}
        def #{name}(ref#{method_param_suffix(param_names, has_opts)}) do
          #{call}
        end
    """
  end

  defp render_attribute(attr) do
    """
        @spec #{attr}(Snakepit.PyRef.t()) :: {:ok, term()} | {:error, Snakepit.Error.t()}
        def #{attr}(ref) do
          SnakeBridge.Runtime.get_attr(ref, :#{attr})
        end
    """
  end

  defp render_discovery(functions, classes) do
    function_list =
      functions
      |> Enum.map_join(",\n      ", fn info ->
        arity = required_arity(info["parameters"] || [])
        summary = info["docstring"] |> to_string() |> String.split("\n") |> List.first() || ""
        "{:#{info["name"]}, #{arity}, __MODULE__, #{inspect(summary)}}"
      end)

    class_list =
      classes
      |> Enum.map_join(",\n      ", fn info ->
        module = class_module_name(info, nil)
        doc = info["docstring"] |> to_string()
        "{#{module}, #{inspect(doc)}}"
      end)

    """
      @doc false
      def __functions__ do
        [
          #{function_list}
        ]
      end

      @doc false
      def __classes__ do
        [
          #{class_list}
        ]
      end

      @doc false
      def __search__(query) do
        SnakeBridge.Docs.search(__MODULE__, query)
      end

      @doc false
      def doc(function) do
        SnakeBridge.Docs.get(__MODULE__, function)
      end
    """
  end

  defp build_params(params) do
    required =
      params
      |> Enum.filter(&(&1["kind"] in ["POSITIONAL_ONLY", "POSITIONAL_OR_KEYWORD"]))
      |> Enum.reject(&Map.has_key?(&1, "default"))

    optional =
      params
      |> Enum.filter(&(&1["kind"] in ["KEYWORD_ONLY", "VAR_KEYWORD"]))

    param_names = Enum.map(required, &sanitize_name/1)
    {param_names, optional != []}
  end

  defp required_arity(params) do
    params
    |> Enum.filter(&(&1["kind"] in ["POSITIONAL_ONLY", "POSITIONAL_OR_KEYWORD"]))
    |> Enum.reject(&Map.has_key?(&1, "default"))
    |> length()
  end

  defp param_list(param_names, true), do: Enum.join(param_names ++ ["opts \\\\ []"], ", ")
  defp param_list(param_names, false), do: Enum.join(param_names, ", ")

  defp runtime_call(name, args, true),
    do: "SnakeBridge.Runtime.call(__MODULE__, :#{name}, #{args}, opts)"

  defp runtime_call(name, args, false),
    do: "SnakeBridge.Runtime.call(__MODULE__, :#{name}, #{args})"

  defp function_spec(name, param_names, has_opts, return_type) do
    args = Enum.map(param_names, fn _ -> "term()" end)
    args = if has_opts, do: args ++ ["keyword()"], else: args

    "@spec #{name}(#{Enum.join(args, ", ")}) :: {:ok, #{return_type}} | {:error, Snakepit.Error.t()}"
  end

  defp method_spec(name, param_names, has_opts) do
    args = ["Snakepit.PyRef.t()"] ++ Enum.map(param_names, fn _ -> "term()" end)
    args = if has_opts, do: args ++ ["keyword()"], else: args

    "@spec #{name}(#{Enum.join(args, ", ")}) :: {:ok, Snakepit.PyRef.t()} | {:error, Snakepit.Error.t()}"
  end

  defp runtime_method_call(name, param_names, true) do
    args = "[" <> Enum.join(param_names, ", ") <> "]"
    "SnakeBridge.Runtime.call_method(ref, :#{name}, #{args}, opts)"
  end

  defp runtime_method_call(name, param_names, false) do
    args = "[" <> Enum.join(param_names, ", ") <> "]"
    "SnakeBridge.Runtime.call_method(ref, :#{name}, #{args})"
  end

  defp method_param_suffix([], true), do: ", opts \\\\ []"
  defp method_param_suffix([], false), do: ""
  defp method_param_suffix(params, true), do: ", " <> Enum.join(params, ", ") <> ", opts \\\\ []"
  defp method_param_suffix(params, false), do: ", " <> Enum.join(params, ", ")

  defp sanitize_name(%{"name" => name}), do: sanitize_name(name)

  defp sanitize_name(name) when is_binary(name) do
    name
    |> Macro.underscore()
    |> String.replace(~r/[^a-z0-9_]/, "_")
    |> ensure_identifier()
  end

  defp ensure_identifier(<<first::utf8, _rest::binary>> = name)
       when first in ?a..?z or first == ?_ do
    name
  end

  defp ensure_identifier(<<first::utf8, _rest::binary>> = name) when first in ?0..?9 do
    "_" <> name
  end

  defp ensure_identifier(""), do: "_unnamed"
  defp ensure_identifier(name), do: "_" <> name

  defp module_to_string(module) when is_atom(module),
    do: module |> Module.split() |> Enum.join(".")

  defp module_to_string(module) when is_binary(module), do: module

  defp class_sort_key(info) do
    {class_module_name(info, nil), class_name(info)}
  end

  defp class_name(info) do
    info["name"] || info["class"] || "Class"
  end

  defp class_python_module(info, library) do
    info["python_module"] || library.python_name
  end

  defp class_module_name(info, nil) do
    info["module"] || class_name(info)
  end

  defp class_module_name(info, library) do
    case info["module"] do
      module when is_binary(module) ->
        module

      _ ->
        python_module = class_python_module(info, library)
        python_parts = String.split(python_module, ".")
        library_parts = String.split(library.python_name, ".")
        extra_parts = Enum.drop(python_parts, length(library_parts))
        extra_parts = drop_class_suffix(extra_parts, class_name(info))

        library.module_name
        |> Module.split()
        |> Kernel.++(Enum.map(extra_parts, &Macro.camelize/1))
        |> Kernel.++([class_name(info)])
        |> Module.concat()
        |> module_to_string()
    end
  end

  defp relative_module_name(library, module_name) when is_binary(module_name) do
    base = module_to_string(library.module_name) <> "."

    if String.starts_with?(module_name, base) do
      String.replace_prefix(module_name, base, "")
    else
      module_name
    end
  end

  defp drop_class_suffix(parts, class_name) when is_list(parts) and is_binary(class_name) do
    class_suffix = Macro.underscore(class_name)

    case List.last(parts) do
      ^class_suffix -> Enum.drop(parts, -1)
      _ -> parts
    end
  end

  defp drop_class_suffix(parts, _class_name), do: parts

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

    text
    |> String.split("\n")
    |> Enum.map_join("\n", fn line -> prefix <> line end)
  end
end