class RubyLLM::StreamAccumulator

def accumulate_tool_calls(new_tool_calls)
  new_tool_calls.each_value do |tool_call|
    if tool_call.id
      @tool_calls[tool_call.id] = ToolCall.new(
        id: tool_call.id,
        name: tool_call.name,
        arguments: String.new
      )
      @latest_tool_call_id = tool_call.id
    else
      existing = @tool_calls[@latest_tool_call_id]
      existing.arguments << tool_call.arguments if existing
    end
  end
end

def add(chunk)
  RubyLLM.logger.debug chunk.inspect
  @model_id ||= chunk.model_id
  if chunk.tool_call?
    accumulate_tool_calls chunk.tool_calls
  else
    @content << (chunk.content || '')
  end
  count_tokens chunk
  RubyLLM.logger.debug inspect
end

def count_tokens(chunk)
  @input_tokens += chunk.input_tokens if chunk.input_tokens
  @output_tokens += chunk.output_tokens if chunk.output_tokens
end

def find_tool_call(tool_call_id)
  if tool_call_id.nil?
    @tool_calls[@latest_tool_call]
  else
    @latest_tool_call_id = tool_call_id
    @tool_calls[tool_call_id]
  end
end

def initialize
  @content = String.new
  @tool_calls = {}
  @input_tokens = 0
  @output_tokens = 0
  @latest_tool_call_id = nil
end

def to_message
  Message.new(
    role: :assistant,
    content: content,
    model_id: model_id,
    tool_calls: tool_calls_from_stream,
    input_tokens: @input_tokens.positive? ? @input_tokens : nil,
    output_tokens: @output_tokens.positive? ? @output_tokens : nil
  )
end

def tool_calls_from_stream
  tool_calls.transform_values do |tc|
    ToolCall.new(
      id: tc.id,
      name: tc.name,
      arguments: JSON.parse(tc.arguments)
    )
  end
end