import {
  StreamableValue,
  createAI,
  createStreamableUI,
  createStreamableValue,
  getMutableAIState
} from 'ai/rsc'
import { ExperimentalMessage } from 'ai'
import { Spinner } from '@/components/ui-v2/spinner'
import { Section } from '@/components/mpv2/section'
import { FollowupPanel } from '@/components/mpv2/followup-panel'
import { inquire, researcher, taskManager, querySuggestor } from '@/lib/agents'

const initialAIState: ExperimentalMessage[] = []

const initialUIState: {
  id: number
  isGenerating: StreamableValue<boolean>
  component: React.ReactNode
}[] = []

const AIProvider = createAI({
  actions: {
    submit: async function(formData?: FormData, skip?: boolean) {
      'use server'
      const aiState = getMutableAIState<typeof AIProvider>()
      const uiStream = createStreamableUI()
      const isGenerating = createStreamableValue(true)
      const messages: ExperimentalMessage[] = aiState.get() as any
      const userInput = skip
        ? `{"action": "skip"}`
        : (formData?.get('input') as string)
      const content = skip
        ? userInput
        : formData
          ? JSON.stringify(Object.fromEntries(formData))
          : null

      if (content) {
        const message = { role: 'user', content }
        messages.push(message as ExperimentalMessage)
        aiState.update([...(aiState.get() as any), message])
      }

      async function processEvents() {
        uiStream.update(<Spinner />)
        let action: any = { object: { next: 'proceed' } }
        if (!skip) action = await taskManager(messages)
        if (action.object.next === 'inquire') {
          const inquiry = await inquire(uiStream, messages)
          uiStream.done()
          isGenerating.done()
          aiState.done([
            ...aiState.get(),
            { role: 'assistant', content: `inquiry: ${inquiry?.question}` }
          ])
          return
        }
        let answer = ''
        let errorOccurred = false
        const streamText = createStreamableValue<string>()
        while (answer.length === 0) {
          const { fullResponse, hasError } = await researcher(
            uiStream,
            streamText,
            messages
          )
          answer = fullResponse
          errorOccurred = hasError
        }
        streamText.done()
        if (!errorOccurred) {
          await querySuggestor(uiStream, messages)
          uiStream.append(
            <Section title="Follow-up">
              <FollowupPanel />
            </Section>
          )
        }
        isGenerating.done(false)
        uiStream.done()
        aiState.done([...aiState.get(), { role: 'assistant', content: answer }])
      }
      processEvents()
      return {
        id: Date.now(),
        isGenerating: isGenerating.value,
        component: uiStream.value
      }
    }
  },
  initialUIState,
  initialAIState
})

export default AIProvider