defmodule Pi.LLM do
  @moduledoc "BEAM API for model calls backed by the active pi session."

  alias Pi.LLM.Broker
  alias Pi.Protocol.LLM.Message

  def complete(messages, opts \\ []) do
    Broker.complete(normalize_messages(messages), opts)
  end

  def stream(messages, opts \\ []) do
    messages
    |> normalize_messages()
    |> Broker.stream(opts)
  end

  def complete!(messages, opts \\ []) do
    case complete(messages, opts) do
      {:ok, result} -> result
      {:error, reason} -> raise RuntimeError, message: to_string(reason)
    end
  end

  defp normalize_messages(messages) when is_binary(messages),
    do: [Message.from_map!(%{role: :user, content: messages})]

  defp normalize_messages(messages) when is_list(messages),
    do: Enum.map(messages, &normalize_message/1)

  defp normalize_message(%Message{} = message), do: message
  defp normalize_message(%{} = message), do: Message.from_map!(message)

  defp normalize_message(message) when is_binary(message),
    do: Message.from_map!(%{role: :user, content: message})
end