import type { AgentFileNode } from '../form-state'
import type { DraftFieldUpdate } from './utils'
import { atom } from 'jotai'
import { syncFileReferenceLabels } from '../reference-labels'
import { agentComposerDraftAtom } from '../store'
import { resolveDraftFieldUpdate } from './utils'

export const agentComposerFilesAtom = atom<AgentFileNode[], [DraftFieldUpdate<AgentFileNode[]>], void>(
  get => get(agentComposerDraftAtom).files,
  (get, set, filesUpdate: DraftFieldUpdate<AgentFileNode[]>) => {
    const draft = get(agentComposerDraftAtom)
    const files = resolveDraftFieldUpdate(draft.files, filesUpdate)

    set(agentComposerDraftAtom, {
      ...draft,
      prompt: syncFileReferenceLabels({
        prompt: draft.prompt,
        currentFiles: draft.files,
        nextFiles: files,
      }),
      files,
    })
  },
)

const removeAgentFileNode = (files: AgentFileNode[], fileId: string): AgentFileNode[] => files.flatMap((file) => {
  if (file.id === fileId)
    return []

  if (file.children)
    return [{ ...file, children: removeAgentFileNode(file.children, fileId) }]

  return [file]
})

export const upsertAgentFileAtom = atom(null, (_get, set, file: AgentFileNode) => {
  set(agentComposerFilesAtom, files => [
    ...removeAgentFileNode(files, file.id),
    file,
  ])
})

export const removeAgentFileAtom = atom(null, (_get, set, fileId: string) => {
  set(agentComposerFilesAtom, files => removeAgentFileNode(files, fileId))
})

export const clearAgentConfigNoteAtom = atom(null, (get, set) => {
  const draft = get(agentComposerDraftAtom)

  set(agentComposerDraftAtom, {
    ...draft,
    configNote: '',
  })
})