const pickPlaceholderRegex = /%%pick::([\w\W]*?)%%/g

export const replacePickPlaceholders = (text: string) => {
  return text.replaceAll(pickPlaceholderRegex, (_, optionsString: string) => {
    const options = optionsString.split('::')
    if (options.length === 0) {
      return ''
    }
    return options[Math.floor(Math.random() * options.length)] ?? ''
  })
}

export const setupPickPlaceholders = () => {
  const { eventTypes, eventSource } = SillyTavern.getContext()
  eventSource.on(
    eventTypes.GENERATE_AFTER_DATA,
    (
      chat: { prompt: string } | { prompt: { role: string; content: string }[] }
    ) => {
      if (typeof chat.prompt === 'string') {
        chat.prompt = replacePickPlaceholders(chat.prompt)
      } else if (Array.isArray(chat.prompt)) {
        chat.prompt = chat.prompt.map((prompt) => {
          prompt.content = replacePickPlaceholders(prompt.content)
          return prompt
        })
      }
    }
  )
}