import type {
  AgentFileNode,
  AgentKnowledgeRetrievalItem,
  AgentSkill,
  AgentTool,
} from './form-state'

const getKnowledgeRetrievalName = (item: AgentKnowledgeRetrievalItem) =>
  item.name ?? item.nameKey ?? item.id

const escapeRegExp = (value: string) => value.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')

const createReferenceToken = (kind: string, id: string, label: string) =>
  `[§${kind}:${id}${label ? `:${label}` : ''}§]`

const syncReferenceLabels = ({
  prompt,
  kind,
  currentItems,
  nextItems,
}: {
  prompt: string
  kind: string
  currentItems: Array<{ id: string; name: string }>
  nextItems: Array<{ id: string; name: string }>
}) => {
  const currentItemById = new Map(currentItems.map((item) => [item.id, item]))

  return nextItems.reduce((nextPrompt, nextItem) => {
    const currentItem = currentItemById.get(nextItem.id)
    if (!currentItem) return nextPrompt

    if (currentItem.name === nextItem.name) return nextPrompt

    return nextPrompt.replace(
      new RegExp(`\\[§${escapeRegExp(kind)}:${escapeRegExp(nextItem.id)}(?::[^§\\]]*)?§\\]`, 'g'),
      createReferenceToken(kind, nextItem.id, nextItem.name),
    )
  }, prompt)
}

const toReferenceLabelItems = <Item extends { id: string }>(
  items: Item[],
  getName: (item: Item) => string,
) =>
  items.map((item) => ({
    id: item.id,
    name: getName(item),
  }))

export const syncKnowledgeReferenceLabels = ({
  prompt,
  currentRetrievals,
  nextRetrievals,
}: {
  prompt: string
  currentRetrievals: AgentKnowledgeRetrievalItem[]
  nextRetrievals: AgentKnowledgeRetrievalItem[]
}) =>
  syncReferenceLabels({
    prompt,
    kind: 'knowledge',
    currentItems: toReferenceLabelItems(currentRetrievals, getKnowledgeRetrievalName),
    nextItems: toReferenceLabelItems(nextRetrievals, getKnowledgeRetrievalName),
  })

export const syncCliToolReferenceLabels = ({
  prompt,
  currentTools,
  nextTools,
}: {
  prompt: string
  currentTools: AgentTool[]
  nextTools: AgentTool[]
}) =>
  syncReferenceLabels({
    prompt,
    kind: 'cli_tool',
    currentItems: toReferenceLabelItems(
      currentTools.filter((tool) => tool.kind === 'cli'),
      (tool) => tool.name,
    ),
    nextItems: toReferenceLabelItems(
      nextTools.filter((tool) => tool.kind === 'cli'),
      (tool) => tool.name,
    ),
  })

export const syncSkillReferenceLabels = ({
  prompt,
  currentSkills,
  nextSkills,
}: {
  prompt: string
  currentSkills: AgentSkill[]
  nextSkills: AgentSkill[]
}) =>
  syncReferenceLabels({
    prompt,
    kind: 'skill',
    currentItems: toReferenceLabelItems(currentSkills, (skill) => skill.name),
    nextItems: toReferenceLabelItems(nextSkills, (skill) => skill.name),
  })

const flattenFileNodes = (files: AgentFileNode[]): AgentFileNode[] =>
  files.flatMap((file) => (file.children?.length ? flattenFileNodes(file.children) : [file]))

export const syncFileReferenceLabels = ({
  prompt,
  currentFiles,
  nextFiles,
}: {
  prompt: string
  currentFiles: AgentFileNode[]
  nextFiles: AgentFileNode[]
}) =>
  syncReferenceLabels({
    prompt,
    kind: 'file',
    currentItems: toReferenceLabelItems(flattenFileNodes(currentFiles), (file) => file.name),
    nextItems: toReferenceLabelItems(flattenFileNodes(nextFiles), (file) => file.name),
  })