class Riffer::Agent::Session

# File lib/riffer/agent/session.rb, line 13
def initialize(messages: [])
  @messages = messages
  @callbacks = [] #: Array[^(Riffer::Messages::Base) -> void]
end

# File lib/riffer/agent/session.rb, line 29
def add(message, silent: false)
  @messages.insert(insertion_index(message), message)
  # +silent+ is for non-inference inputs (e.g. user messages) that subscribers don't expect.
  @callbacks.each { |callback| callback.call(message) } unless silent
  message
end

# File lib/riffer/agent/session.rb, line 125
def discard_pending_tool_calls
  messages, filled = Repair.fill_orphans(@messages)
  set(messages)
  filled
end

# File lib/riffer/agent/session.rb, line 134
def each(&block)
  return @messages.each unless block

  @messages.each(&block)
end

# File lib/riffer/agent/session.rb, line 148
def final_assistant_message
  # TODO: Replace with rfind when minimum Ruby is 4.0+
  @messages.reverse_each.find { |m| m.is_a?(Riffer::Messages::Assistant) } #: Riffer::Messages::Assistant?
end

# File lib/riffer/agent/session.rb, line 20
def on_message(&block)
  raise Riffer::ArgumentError, "on_message requires a block" unless block_given?

  @callbacks << block
  self
end

# File lib/riffer/agent/session.rb, line 98
def orphaned_tool_call_ids
  result_ids = @messages.filter_map { |m| m.tool_call_id if m.is_a?(Riffer::Messages::Tool) }
  @messages.flat_map do |m|
    next [] unless m.is_a?(Riffer::Messages::Assistant)

    m.tool_calls.reject { |tc| result_ids.include?(tc.call_id) }.map(&:call_id)
  end
end

# File lib/riffer/agent/session.rb, line 109
def pending_tool_calls
  last_assistant_idx = @messages.rindex { |m| m.is_a?(Riffer::Messages::Assistant) }
  return [nil, []] unless last_assistant_idx

  assistant = @messages[last_assistant_idx] #: Riffer::Messages::Assistant
  return [assistant, []] if assistant.tool_calls.empty?

  executed_ids = (@messages[(last_assistant_idx + 1)..] || []).filter_map do |m|
    m.tool_call_id if m.is_a?(Riffer::Messages::Tool)
  end

  [assistant, assistant.tool_calls.reject { |tc| executed_ids.include?(tc.call_id) }]
end

# File lib/riffer/agent/session.rb, line 52
def remove(id:)
  idx = @messages.index { |m| m.id == id }
  return nil unless idx

  target = @messages[idx]
  if target.is_a?(Riffer::Messages::Tool)
    raise Riffer::ArgumentError,
          "remove cannot drop a Tool message (would orphan the parent's tool_use); use #update instead"
  end

  if target.is_a?(Riffer::Messages::Assistant) && !target.tool_calls.empty?
    child_ids = target.tool_calls.map(&:call_id)
    @messages.reject! { |m| m.is_a?(Riffer::Messages::Tool) && child_ids.include?(m.tool_call_id) }
    @messages.delete(target)
  else
    @messages.delete_at(idx)
  end
  target
end

# File lib/riffer/agent/session.rb, line 38
def set(messages)
  @messages = messages
  self
end

# File lib/riffer/agent/session.rb, line 142
def steps
  @messages.count { |m| m.is_a?(Riffer::Messages::Assistant) }
end

# File lib/riffer/agent/session.rb, line 45
def unset
  @messages = []
  self
end

# File lib/riffer/agent/session.rb, line 74
def update(id: nil, tool_call_id: nil, **attrs)
  raise Riffer::ArgumentError, "update requires either id: or tool_call_id:" if id.nil? && tool_call_id.nil?
  raise Riffer::ArgumentError, "update accepts id: or tool_call_id:, not both" if id && tool_call_id

  idx = if id
          @messages.index { |m| m.id == id }
        else
          @messages.index { |m| m.is_a?(Riffer::Messages::Tool) && m.tool_call_id == tool_call_id }
        end

  unless idx
    key = id ? "id #{id.inspect}" : "tool_call_id #{tool_call_id.inspect}"
    raise Riffer::ArgumentError, "no message found for #{key}"
  end

  old = @messages[idx] #: Riffer::Messages::Base
  replacement = rebuild_message(old, attrs)
  @messages[idx] = replacement
  cascade_dropped_tool_calls(old, replacement)
  replacement
end