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()
    ref = make_ref()

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

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

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

    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
              if Process.get(:sse_emitted?) != true do
                raw_buffer = Process.get(:sse_raw_buffer, "") <> chunk
                Process.put(:sse_raw_buffer, raw_buffer)
              end

              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, ref, error})
    end

    if Process.get(:sse_emitted?) != true do
      raw_buffer = Process.get(:sse_raw_buffer, "")
      trimmed = String.trim(raw_buffer)

      if trimmed not in ["", "[DONE]"] do
        case Jason.decode(trimmed) do
          {:ok, decoded} ->
            send(parent, {:sse_event, ref, decoded})

          _ ->
            :ok
        end
      end
    end

    send(parent, {:sse_done, ref})
  rescue
    error ->
      ref = Process.get(:sse_ref)
      send(parent, {:sse_error, ref, error})
  end

  defp next_event(ref) do
    receive do
      {:sse_event, ^ref, event} -> {[event], ref}
      {:sse_error, ^ref, error} -> {[{:error, error}], ref}
      {:sse_done, ^ref} -> {:halt, ref}
    end
  end

  defp parse_chunk(chunk, state, parent) when is_binary(chunk) do
    ref = Process.get(:sse_ref)
    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, ref})

        nil ->
          :ok

        data ->
          case Jason.decode(data) do
            {:ok, decoded} ->
              Process.put(:sse_emitted?, true)
              send(parent, {:sse_event, ref, decoded})

            {:error, error} ->
              send(parent, {:sse_error, ref, 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
    data_lines =
      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 data_lines do
      [:done | _] ->
        :done

      [] ->
        trimmed = String.trim(raw_event)

        if trimmed in ["", "[DONE]"] do
          nil
        else
          if String.starts_with?(trimmed, "{") or String.starts_with?(trimmed, "[") do
            trimmed
          else
            nil
          end
        end

      _ ->
        data_lines |> Enum.reverse() |> Enum.join("\n")
    end
  end
end