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}`
    - `:tool_sequence` - list of tool names to force in order (optional)
    - `:max_steps` - tool loop limit (default: 10)
    - `:continue_prompt` - user prompt appended after tool results (optional)
    - `: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)
    tool_sequence = Map.get(opts, :tool_sequence)
    continue_prompt = Map.get(opts, :continue_prompt)
    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,
      tool_sequence: tool_sequence,
      continue_prompt: continue_prompt,
      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 = tool_choice_for_step(state)

    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)
            |> maybe_append_continue(state.continue_prompt)

          do_stream_loop(
            %{
              state
              | messages: next_messages,
                step: state.step + step_increment(state, tool_calls),
                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, %{})

  defp tool_choice_for_step(%{tool_sequence: tool_sequence, step: step} = state)
       when is_list(tool_sequence) do
    case Enum.at(tool_sequence, step) do
      nil -> default_tool_choice(state)
      name when is_binary(name) -> %{tool: name}
      other -> other
    end
  end

  defp tool_choice_for_step(state), do: default_tool_choice(state)

  defp default_tool_choice(%{step: step, tool_choice: :required}) when step > 0 do
    :none
  end

  defp default_tool_choice(%{tool_choice: tool_choice}), do: tool_choice

  defp step_increment(%{tool_sequence: tool_sequence}, tool_calls)
       when is_list(tool_sequence) do
    max(length(tool_calls), 1)
  end

  defp step_increment(_state, _tool_calls), do: 1

  defp maybe_append_continue(messages, nil), do: messages

  defp maybe_append_continue(messages, prompt) when is_binary(prompt) do
    messages ++ [%{role: "user", content: prompt}]
  end
end