loom_ex

defmodule LoomEx.Message do
  @type role :: :system | :user | :assistant | :tool

  @type t :: %__MODULE__{
          role: role(),
          content: String.t() | nil,
          tool_calls: [LoomEx.ToolCall.t()] | nil,
          tool_call_id: String.t() | nil,
          name: String.t() | nil,
          reasoning_content: String.t() | nil
        }

  @enforce_keys [:role]
  defstruct [:role, :content, :tool_calls, :tool_call_id, :name, :reasoning_content]

  @spec system(String.t()) :: t()
  def system(content), do: %__MODULE__{role: :system, content: content}

  @spec user(String.t()) :: t()
  def user(content), do: %__MODULE__{role: :user, content: content}

  @spec assistant(String.t() | nil, [LoomEx.ToolCall.t()] | nil) :: t()
  def assistant(content, tool_calls \\ nil) do
    %__MODULE__{role: :assistant, content: content, tool_calls: tool_calls}
  end

  @spec tool_result(String.t(), term()) :: t()
  def tool_result(tool_call_id, result) do
    content = if is_binary(result), do: result, else: Jason.encode!(result)
    %__MODULE__{role: :tool, tool_call_id: tool_call_id, content: content}
  end

  @spec to_openai(t()) :: map()
  def to_openai(%__MODULE__{role: :system, content: content}) do
    %{"role" => "system", "content" => content}
  end

  def to_openai(%__MODULE__{role: :user, content: content}) do
    %{"role" => "user", "content" => content}
  end

  def to_openai(%__MODULE__{role: :assistant} = msg) do
    base = %{"role" => "assistant", "content" => msg.content || ""}

    base =
      if msg.reasoning_content do
        Map.put(base, "reasoning_content", msg.reasoning_content)
      else
        base
      end

    if msg.tool_calls && msg.tool_calls != [] do
      tool_calls =
        Enum.map(msg.tool_calls, fn tc ->
          %{
            "id" => tc.id,
            "type" => "function",
            "function" => %{
              "name" => tc.name,
              "arguments" =>
                if(is_binary(tc.arguments), do: tc.arguments, else: Jason.encode!(tc.arguments))
            }
          }
        end)

      Map.put(base, "tool_calls", tool_calls)
    else
      base
    end
  end

  def to_openai(%__MODULE__{role: :tool} = msg) do
    %{
      "role" => "tool",
      "tool_call_id" => msg.tool_call_id,
      "content" => msg.content
    }
  end

  @spec from_openai(map()) :: t()
  def from_openai(%{"role" => "system", "content" => content}) do
    system(content)
  end

  def from_openai(%{"role" => "user", "content" => content}) do
    user(content)
  end

  def from_openai(%{"role" => "assistant"} = msg) do
    tool_calls =
      case msg["tool_calls"] do
        nil ->
          nil

        [] ->
          nil

        tcs ->
          Enum.map(tcs, fn tc ->
            args = get_in(tc, ["function", "arguments"]) || ""

            parsed =
              case Jason.decode(args) do
                {:ok, map} -> map
                _ -> args
              end

            %LoomEx.ToolCall{
              id: tc["id"],
              name: get_in(tc, ["function", "name"]),
              arguments: parsed
            }
          end)
      end

    %__MODULE__{
      role: :assistant,
      content: msg["content"],
      tool_calls: tool_calls,
      reasoning_content: msg["reasoning_content"]
    }
  end

  def from_openai(%{"role" => "tool"} = msg) do
    %__MODULE__{
      role: :tool,
      tool_call_id: msg["tool_call_id"],
      content: msg["content"]
    }
  end
end