ai_sdk_ex

defmodule AI.StreamText do
  @moduledoc """
  Streams text and tool calls from a provider-backed model.
  """

  @doc """
  Streams a request and returns `{:ok, stream}`.

  Required options:

    - `:model` - an `AI.Model` from `AI.OpenAI`, `AI.Anthropic`, or `AI.OpenRouter`
    - `:prompt` or `:messages` - user input

  Optional options:

    - `:system` - system prompt (string)
    - `:tools` - map of tool name to `AI.Tool`
    - `:tool_choice` - `:auto | :none | :required | %{tool: name}`
    - `:max_steps` - tool loop limit (default: 10)
    - `:temperature`, `:max_tokens`, `:stop`, `:headers`, `:timeout`
  """
  def stream_text(opts) do
    opts = normalize_opts(opts)

    _ = AI.Env.load()

    model = Map.fetch!(opts, :model)
    tools = Map.get(opts, :tools, %{})
    tool_choice = Map.get(opts, :tool_choice, :auto)
    settings = Map.take(opts, [:temperature, :max_tokens, :stop, :headers, :timeout])
    max_steps = Map.get(opts, :max_steps, 10)

    {messages, system} = AI.Prompt.standardize(opts)

    state = %{
      model: model,
      tools: tools,
      tool_choice: tool_choice,
      settings: settings,
      messages: messages,
      system: system,
      step: 0,
      max_steps: max_steps,
      response_id: nil
    }

    parent = self()

    {:ok, pid} =
      Task.start_link(fn ->
        stream_loop(state, parent)
      end)

    stream = Stream.resource(fn -> pid end, &next_event/1, fn _ -> :ok end)

    {:ok, stream}
  end

  defp next_event(pid) do
    receive do
      {:ai_event, event} -> {[event], pid}
      {:ai_done, ^pid} -> {:halt, pid}
      {:ai_error, ^pid, error} -> {[{:error, error}], pid}
    end
  end

  defp stream_loop(state, parent) do
    do_stream_loop(state, parent)
    send(parent, {:ai_done, self()})
  rescue
    error ->
      send(parent, {:ai_error, self(), error})
  end

  defp do_stream_loop(%{step: step, max_steps: max_steps}, parent)
       when step >= max_steps do
    send(parent, {:ai_event, {:finish, %{reason: "max_steps", usage: nil}}})
  end

  defp do_stream_loop(state, parent) do
    step_tool_choice =
      if state.step > 0 and state.tool_choice == :required do
        :none
      else
        state.tool_choice
      end

    provider_opts = %{
      model: state.model,
      model_id: state.model.model_id,
      options: state.model.options,
      messages: AI.Messages.to_provider_messages(state.messages, state.system),
      system: state.system,
      tools: state.tools,
      tool_choice: step_tool_choice,
      settings: state.settings,
      previous_response_id: state.response_id
    }

    case state.model.provider.do_stream(provider_opts) do
      {:ok, provider_stream, _req, _resp} ->
        acc =
          Enum.reduce(provider_stream, %{tool_calls: [], response_id: state.response_id}, fn
            event, acc_state ->
              add_events(AI.StreamEvents.normalize(event), acc_state, parent)
          end)

        tool_calls = acc.tool_calls

        if tool_calls == [] do
          :ok
        else
          tool_results =
            AI.ToolExecutor.execute_parallel(tool_calls, state.tools, %{messages: state.messages})

          Enum.each(tool_results, fn result ->
            send(parent, {:ai_event, {:tool_result, result}})
          end)

          next_messages =
            state.messages
            |> AI.Messages.append_tool_calls(tool_calls)
            |> AI.Messages.append_tool_results(tool_results)

          do_stream_loop(
            %{
              state
              | messages: next_messages,
                step: state.step + 1,
                response_id: acc.response_id
            },
            parent
          )
        end

      {:error, error} ->
        send(parent, {:ai_event, {:error, error}})
    end
  end

  defp add_events(normalized, acc, parent) do
    normalized
    |> List.wrap()
    |> Enum.reject(&is_nil/1)
    |> Enum.reduce(acc, fn event, acc_state ->
      case event do
        {:response_id, id} ->
          %{acc_state | response_id: id}

        {:tool_call, call} ->
          send(parent, {:ai_event, event})
          %{acc_state | tool_calls: [call | acc_state.tool_calls]}

        _ ->
          send(parent, {:ai_event, event})
          acc_state
      end
    end)
  end

  defp normalize_opts(opts) when is_map(opts), do: opts
  defp normalize_opts(opts) when is_list(opts), do: Enum.into(opts, %{})
end