snakebridge

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

  alias SnakeBridge.Docs.{MarkdownConverter, RstParser}
  alias SnakeBridge.Generator.TypeMapper

  @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
    start_time = System.monotonic_time()
    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))

    result = write_if_changed(path, source)

    bytes_written = if result == :written, do: byte_size(source), else: 0

    SnakeBridge.Telemetry.generate_stop(
      start_time,
      library.name,
      path,
      bytes_written,
      length(functions),
      length(classes)
    )

    :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"] || ""

    plan = build_params(params)
    param_names = Enum.map(plan.required, & &1.name)
    args_name = extra_args_name(param_names)
    return_type = info["return_type"] || %{"type" => "any"}

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

    normal = render_normal_function(name, plan, args_name, return_type, doc, params)

    if is_streaming do
      streaming = render_streaming_variant(name, plan, args_name)
      normal <> "\n\n" <> streaming
    else
      normal
    end
  end

  defp render_normal_function(name, plan, args_name, return_type, doc, params) do
    param_names = Enum.map(plan.required, & &1.name)
    args = args_expr(param_names, plan.has_args, args_name)
    call = runtime_call(name, args, plan.has_opts)
    spec = function_spec(name, plan.required, plan.has_args, return_type)
    formatted_doc = format_docstring(doc, params, return_type)
    normalize = normalize_args_line(plan.has_args, args_name, 8)

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

  defp render_streaming_variant(name, plan, args_name) do
    param_names = Enum.map(plan.required, & &1.name)
    args = args_expr(param_names, plan.has_args, args_name)

    stream_params =
      param_names
      |> maybe_add_args(plan.has_args, args_name)
      |> Kernel.++(["opts \\\\ []", "callback"])

    stream_params_str = Enum.join(stream_params, ", ")

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

    spec_args =
      plan.required
      |> Enum.map(&param_type_spec/1)
      |> maybe_add_args_spec(plan.has_args)
      |> Kernel.++(["keyword()", "(term() -> any())"])

    spec_args_str = Enum.join(spec_args, ", ")

    base_arity = length(param_names) + if(plan.has_args, do: 2, else: 1)
    normalize = normalize_args_line(plan.has_args, args_name, 8)

    """
      @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
    #{normalize}        #{stream_call}
      end
    """
  end

  defp render_submodule(python_module, functions, library) do
    module_name =
      python_module
      |> String.split(".")
      |> Enum.drop(length(String.split(library.python_name, ".")))
      |> 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}"
        def __snakebridge_library__, do: "#{library.python_name}"

    #{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: []
    plan = build_params(init_params)
    param_names = Enum.map(plan.required, & &1.name)
    args_name = extra_args_name(param_names)
    constructor = render_constructor(plan, args_name)

    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}"
        def __snakebridge_library__, do: "#{library.python_name}"
        @opaque t :: SnakeBridge.Ref.t()

    #{indent(constructor, 4)}

    #{indent(methods_source, 4)}

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

  defp render_constructor(plan, args_name) do
    param_names = Enum.map(plan.required, & &1.name)
    args = args_expr(param_names, plan.has_args, args_name)

    param_list = param_list(param_names, plan.has_args, plan.has_opts, args_name)

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

    spec_args =
      plan.required
      |> Enum.map(&param_type_spec/1)
      |> maybe_add_args_spec(plan.has_args)
      |> Kernel.++(["keyword()"])

    spec_args_str = Enum.join(spec_args, ", ")
    normalize = normalize_args_line(plan.has_args, args_name, 10)

    """
        @spec new(#{spec_args_str}) :: {:ok, SnakeBridge.Ref.t()} | {:error, Snakepit.Error.t()}
        def new(#{param_list}) do
    #{normalize}          #{call}
        end
    """
  end

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

  defp render_method(info) do
    name = info["name"]
    params = info["parameters"] || []
    plan = build_params(params)
    param_names = Enum.map(plan.required, & &1.name)
    args_name = extra_args_name(param_names)
    return_type = info["return_type"] || %{"type" => "any"}
    spec = method_spec(name, plan.required, plan.has_args, return_type)
    call = runtime_method_call(name, param_names, plan.has_args, args_name)
    normalize = normalize_args_line(plan.has_args, args_name, 10)

    """
        #{spec}
        def #{name}(ref#{method_param_suffix(param_names, plan.has_args, plan.has_opts, args_name)}) do
    #{normalize}          #{call}
        end
    """
  end

  defp render_attribute(attr) do
    """
        @spec #{attr}(SnakeBridge.Ref.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, list(), map() | nil) :: String.t()
  def format_docstring(raw_doc, params \\ [], return_type \\ nil)

  def format_docstring(nil, _params, _return_type), do: ""
  def format_docstring("", _params, _return_type), do: ""

  def format_docstring(raw_doc, params, return_type) when is_binary(raw_doc) do
    base =
      raw_doc
      |> RstParser.parse()
      |> MarkdownConverter.convert()

    extras = format_param_docs(params, return_type)

    if extras == "" do
      base
    else
      base <> "\n\n" <> extras
    end
  rescue
    _ ->
      extras = format_param_docs(params, return_type)

      if extras == "" do
        raw_doc
      else
        raw_doc <> "\n\n" <> extras
      end
  end

  @spec build_params(list()) :: %{
          required: list(map()),
          has_args: boolean(),
          has_opts: boolean()
        }
  def build_params(params) do
    required =
      params
      |> Enum.filter(&required_positional?/1)
      |> Enum.map(&param_entry/1)

    has_args =
      Enum.any?(params, fn param ->
        optional_positional?(param) or varargs?(param)
      end)

    %{required: required, has_args: has_args, has_opts: 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_args, has_opts, args_name) do
    param_names
    |> maybe_add_args(has_args, args_name)
    |> maybe_add_opts(has_opts)
    |> Enum.join(", ")
  end

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

  defp function_spec(name, param_entries, has_args, return_type) do
    args =
      param_entries
      |> Enum.map(&param_type_spec/1)
      |> maybe_add_args_spec(has_args)
      |> Kernel.++(["keyword()"])

    return_spec = type_spec_string(return_type)

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

  defp method_spec(name, param_entries, has_args, return_type) do
    args =
      [ref_type_spec()]
      |> Kernel.++(Enum.map(param_entries, &param_type_spec/1))
      |> maybe_add_args_spec(has_args)
      |> Kernel.++(["keyword()"])

    return_spec = type_spec_string(return_type)

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

  defp runtime_method_call(name, param_names, has_args, args_name) do
    args = args_expr(param_names, has_args, args_name)
    "SnakeBridge.Runtime.call_method(ref, :#{name}, #{args}, opts)"
  end

  defp method_param_suffix(param_names, has_args, has_opts, args_name) do
    suffix =
      param_names
      |> maybe_add_args(has_args, args_name)
      |> maybe_add_opts(has_opts)

    ", " <> Enum.join(suffix, ", ")
  end

  defp args_expr(param_names, true, args_name) do
    base = "[" <> Enum.join(param_names, ", ") <> "]"
    base <> " ++ List.wrap(" <> args_name <> ")"
  end

  defp args_expr(param_names, false, _args_name) do
    "[" <> Enum.join(param_names, ", ") <> "]"
  end

  defp maybe_add_args(items, true, args_name), do: items ++ ["#{args_name} \\\\ []"]
  defp maybe_add_args(items, false, _args_name), do: items

  defp maybe_add_opts(items, _has_opts), do: items ++ ["opts \\\\ []"]

  defp maybe_add_args_spec(items, true), do: items ++ ["list(term())"]
  defp maybe_add_args_spec(items, false), do: items

  defp normalize_args_line(true, args_name, indent) do
    String.duplicate(" ", indent) <>
      "{#{args_name}, opts} = SnakeBridge.Runtime.normalize_args_opts(#{args_name}, opts)\n"
  end

  defp normalize_args_line(false, _args_name, _indent), do: ""

  defp extra_args_name(param_names) do
    if "args" in param_names do
      "extra_args"
    else
      "args"
    end
  end

  defp param_entry(param) do
    %{
      name: sanitize_name(param),
      type: param_type(param)
    }
  end

  defp param_type(%{"type" => type}) when is_map(type), do: type
  defp param_type(%{type: type}) when is_map(type), do: type
  defp param_type(_), do: %{"type" => "any"}

  defp param_type_spec(%{type: type}), do: type_spec_string(type)
  defp param_type_spec(_), do: "term()"

  defp type_spec_string(type) do
    type
    |> TypeMapper.to_spec()
    |> Macro.to_string()
  end

  defp ref_type_spec do
    "SnakeBridge.Ref.t()"
  end

  defp required_positional?(param) do
    kind = param_kind(param)
    kind in ["POSITIONAL_ONLY", "POSITIONAL_OR_KEYWORD"] and not param_default?(param)
  end

  defp optional_positional?(param) do
    kind = param_kind(param)
    kind in ["POSITIONAL_ONLY", "POSITIONAL_OR_KEYWORD"] and param_default?(param)
  end

  defp varargs?(param), do: param_kind(param) == "VAR_POSITIONAL"

  defp param_kind(%{"kind" => kind}), do: kind
  defp param_kind(%{kind: kind}), do: kind
  defp param_kind(_), do: nil

  defp param_default?(%{"default" => _}), do: true
  defp param_default?(%{default: _}), do: true
  defp param_default?(_), do: false

  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

  defp format_param_docs(params, return_type) do
    param_lines =
      params
      |> Enum.map(&param_doc_line/1)
      |> Enum.reject(&is_nil/1)

    sections =
      []
      |> maybe_add_params_section(param_lines)
      |> maybe_add_return_section(return_type)

    Enum.join(sections, "\n\n")
  end

  defp maybe_add_params_section(sections, []), do: sections

  defp maybe_add_params_section(sections, lines) do
    sections ++ ["Parameters:\n" <> Enum.join(lines, "\n")]
  end

  defp maybe_add_return_section(sections, nil), do: sections

  defp maybe_add_return_section(sections, return_type) do
    return_spec = type_spec_string(return_type)
    sections ++ ["Returns:\n- `#{return_spec}`"]
  end

  defp param_doc_line(param) do
    name = param_name(param)

    if name do
      type = param_type(param) |> type_spec_string()
      default = param_default_value(param)

      default_fragment =
        if default do
          " default: #{default}"
        else
          ""
        end

      "- `#{name}` (#{type}#{default_fragment})"
    else
      nil
    end
  end

  defp param_name(%{"name" => name}) when is_binary(name), do: name
  defp param_name(%{name: name}) when is_binary(name), do: name
  defp param_name(_), do: nil

  defp param_default_value(%{"default" => default}), do: default
  defp param_default_value(%{default: default}), do: default
  defp param_default_value(_), do: nil
end