ai_sdk_ex

defmodule AI.HTTP.SSE do
  @moduledoc """
  Minimal SSE client for streaming JSON events.
  """

  @doc """
  Streams a request and emits decoded JSON events.
  """
  def stream_request(method, url, headers, body, opts \\ []) do
    parent = self()

    {:ok, _pid} =
      Task.start_link(fn ->
        request_loop(method, url, headers, body, opts, parent)
      end)

    Stream.resource(fn -> :ok end, fn _ -> next_event() end, fn _ -> :ok end)
  end

  defp request_loop(method, url, headers, body, opts, parent) do
    Process.put(:sse_parser, %{buffer: ""})
    Process.put(:sse_status, nil)
    Process.put(:sse_error_buffer, "")

    req_opts =
      [
        method: method,
        url: url,
        headers: headers,
        body: body,
        into: fn
          {:data, chunk}, {req, resp} ->
            status = Process.get(:sse_status) || resp.status
            Process.put(:sse_status, status)

            if status >= 400 do
              buffer = Process.get(:sse_error_buffer, "") <> chunk
              Process.put(:sse_error_buffer, buffer)
            else
              parser = Process.get(:sse_parser, %{buffer: ""})
              new_parser = parse_chunk(chunk, parser, parent)
              Process.put(:sse_parser, new_parser)
            end

            {:cont, {req, resp}}

          _other, {req, resp} ->
            {:cont, {req, resp}}
        end
      ]
      |> Keyword.merge(opts)

    response = Req.request!(req_opts)
    status = Process.get(:sse_status) || response.status

    if status && status >= 400 do
      body = Process.get(:sse_error_buffer, "")

      error =
        case Jason.decode(body) do
          {:ok, decoded} -> decoded
          _ -> body
        end

      send(parent, {:sse_error, error})
    end

    send(parent, :sse_done)
  rescue
    error ->
      send(parent, {:sse_error, error})
  end

  defp next_event do
    receive do
      {:sse_event, event} -> {[event], :ok}
      {:sse_error, error} -> {[{:error, error}], :ok}
      :sse_done -> {:halt, :ok}
    end
  end

  defp parse_chunk(chunk, state, parent) when is_binary(chunk) do
    buffer = state.buffer <> chunk
    {events, rest} = split_events(buffer)

    Enum.each(events, fn raw_event ->
      case extract_data(raw_event) do
        :done ->
          send(parent, :sse_done)

        nil ->
          :ok

        data ->
          case Jason.decode(data) do
            {:ok, decoded} -> send(parent, {:sse_event, decoded})
            {:error, error} -> send(parent, {:sse_error, error})
          end
      end
    end)

    %{state | buffer: rest}
  end

  defp split_events(buffer) do
    normalized = String.replace(buffer, "\r\n", "\n")
    parts = String.split(normalized, "\n\n")

    case Enum.split(parts, -1) do
      {events, [rest]} -> {events, rest}
      {events, []} -> {events, ""}
    end
  end

  defp extract_data(raw_event) do
    raw_event
    |> String.split("\n")
    |> Enum.reduce([], fn line, acc ->
      case line do
        "data: [DONE]" -> [:done]
        "data:" <> data -> [String.trim_leading(data) | acc]
        _ -> acc
      end
    end)
    |> case do
      [:done | _] -> :done
      [] -> nil
      data_lines -> data_lines |> Enum.reverse() |> Enum.join("\n")
    end
  end
end