snakebridge

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

  alias SnakeBridge.Docs.{MarkdownConverter, RstParser}

  @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, library))

    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, library)
      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))

    write_if_changed(path, source)
    :ok
  end

  @spec write_if_changed(String.t(), String.t()) :: :written | :unchanged
  def write_if_changed(path, new_content) do
    case File.read(path) do
      {:ok, existing} when existing == new_content ->
        :unchanged

      _ ->
        temp_path = "#{path}.tmp.#{System.unique_integer([:positive])}"

        try do
          File.write!(temp_path, new_content)
          File.rename!(temp_path, path)
          :written
        after
          File.rm(temp_path)
        end
    end
  end

  @spec render_function(map(), SnakeBridge.Config.Library.t()) :: String.t()
  def render_function(info, library) do
    name = info["name"]
    params = info["parameters"] || []
    doc = info["docstring"] || ""

    {param_names, has_opts} = build_params(params)
    args = "[" <> Enum.join(param_names, ", ") <> "]"

    is_streaming = name in (library.streaming || [])

    normal = render_normal_function(name, param_names, has_opts, args, doc)

    if is_streaming do
      streaming = render_streaming_variant(name, param_names, args, has_opts)
      normal <> "\n\n" <> streaming
    else
      normal
    end
  end

  defp render_normal_function(name, param_names, has_opts, args, doc) do
    call = runtime_call(name, args, has_opts)
    spec = function_spec(name, param_names, has_opts, "term()")
    formatted_doc = format_docstring(doc)

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

  defp render_streaming_variant(name, param_names, args, _has_opts) do
    stream_params = param_names ++ ["opts \\\\ []", "callback"]
    stream_params_str = Enum.join(stream_params, ", ")

    stream_call = "SnakeBridge.Runtime.stream(__MODULE__, :#{name}, #{args}, opts, callback)"

    spec_args = Enum.map(param_names, fn _ -> "term()" end) ++ ["keyword()", "(term() -> any())"]
    spec_args_str = Enum.join(spec_args, ", ")

    base_arity = length(param_names) + 1

    """
      @doc \"\"\"
      Streaming variant of `#{name}/#{base_arity}`.

      The callback receives chunks as they arrive.
      \"\"\"
      @spec #{name}_stream(#{spec_args_str}) :: :ok | {:error, Snakepit.Error.t()}
      def #{name}_stream(#{stream_params_str}) when is_function(callback, 1) do
        #{stream_call}
      end
    """
  end

  defp render_submodule(python_module, functions, library) 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, library))

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

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

  @spec render_class(map(), SnakeBridge.Config.Library.t()) :: String.t()
  def 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"] || []

    init_method = Enum.find(methods, fn method -> method["name"] == "__init__" end)
    init_params = if init_method, do: init_method["parameters"] || [], else: []
    {param_names, has_opts} = build_params(init_params)
    constructor = render_constructor(param_names, has_opts)

    methods_source =
      methods
      |> Enum.reject(fn method -> method["name"] == "__init__" end)
      |> 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}"

    #{indent(constructor, 4)}

    #{indent(methods_source, 4)}

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

  defp render_constructor(param_names, _has_opts) do
    args = "[" <> Enum.join(param_names, ", ") <> "]"

    param_list = Enum.join(param_names ++ ["opts \\\\ []"], ", ")

    call = "SnakeBridge.Runtime.call_class(__MODULE__, :__init__, #{args}, opts)"

    spec_args = Enum.map(param_names, fn _ -> "term()" end)
    spec_args = spec_args ++ ["keyword()"]
    spec_args_str = Enum.join(spec_args, ", ")

    """
        @spec new(#{spec_args_str}) :: {:ok, Snakepit.PyRef.t()} | {:error, Snakepit.Error.t()}
        def new(#{param_list}) do
          #{call}
        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

  @spec format_docstring(String.t() | nil) :: String.t()
  def format_docstring(nil), do: ""
  def format_docstring(""), do: ""

  def format_docstring(raw_doc) when is_binary(raw_doc) do
    raw_doc
    |> RstParser.parse()
    |> MarkdownConverter.convert()
  rescue
    _ -> raw_doc
  end

  @spec build_params(list()) :: {[String.t()], boolean()}
  def build_params(params) do
    required =
      params
      |> Enum.filter(&(&1["kind"] in ["POSITIONAL_ONLY", "POSITIONAL_OR_KEYWORD"]))
      |> Enum.reject(&Map.has_key?(&1, "default"))

    param_names = Enum.map(required, &sanitize_name/1)
    {param_names, true}
  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, _has_opts) do
    Enum.join(param_names ++ ["opts \\\\ []"], ", ")
  end

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

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

    "@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) ++ ["keyword()"]

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

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

  defp method_param_suffix([], _has_opts), do: ", opts \\\\ []"

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

  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