import type { ReActRound, ReActStep } from '../types'

export function groupStepsIntoRounds(steps: ReActStep[]): ReActRound[] {
  const map = new Map<number, ReActRound>()

  for (const step of steps) {
    if (step.type === 'input' || step.type === 'answer' || step.type === 'plain') continue
    if (step.round <= 0) continue

    let round = map.get(step.round)
    if (!round) {
      round = { round: step.round }
      map.set(step.round, round)
    }

    if (step.type === 'thought') round.thought = step
    if (step.type === 'action') round.action = step
    if (step.type === 'observation') round.observation = step
  }

  return [...map.values()].sort((a, b) => a.round - b.round)
}

export function countToolCalls(steps: ReActStep[]): number {
  return steps.filter((s) => s.type === 'action').length
}