import type { LLMNodeType } from '../types'
import type {
  PromptItem,
  RolePrefix,
} from '@/app/components/workflow/types'
import { produce } from 'immer'
import {
  useCallback,
  useEffect,
  useRef,
  useState,
} from 'react'

type CompletionPromptTemplate = {
  prompt: PromptItem
  conversation_histories_role: {
    user_prefix: string
    assistant_prefix: string
  }
}

export type LLMDefaultConfig = {
  prompt_templates?: {
    chat_model: {
      prompts: PromptItem[]
    }
    completion_model: CompletionPromptTemplate
  }
}

type Params = {
  inputs: LLMNodeType
  doSetInputs: (inputs: LLMNodeType) => void
  defaultConfig?: LLMDefaultConfig
  isChatModel: boolean
}

const useLLMInputManager = ({
  inputs,
  doSetInputs,
  defaultConfig,
  isChatModel,
}: Params) => {
  const [defaultRolePrefix, setDefaultRolePrefix] = useState<RolePrefix>({ user: '', assistant: '' })
  const inputRef = useRef(inputs)

  useEffect(() => {
    inputRef.current = inputs
  }, [inputs])

  const setInputs = useCallback((newInputs: LLMNodeType) => {
    if (newInputs.memory && !newInputs.memory.role_prefix) {
      const payloadWithRolePrefix = produce(newInputs, (draft) => {
        draft.memory!.role_prefix = defaultRolePrefix
      })
      doSetInputs(payloadWithRolePrefix)
      inputRef.current = payloadWithRolePrefix
      return
    }

    doSetInputs(newInputs)
    inputRef.current = newInputs
  }, [defaultRolePrefix, doSetInputs])

  const appendDefaultPromptConfig = useCallback((draft: LLMNodeType, nextDefaultConfig: LLMDefaultConfig, passInIsChatMode?: boolean) => {
    const promptTemplates = nextDefaultConfig.prompt_templates
    if (!promptTemplates)
      return

    if (passInIsChatMode === undefined ? isChatModel : passInIsChatMode) {
      draft.prompt_template = promptTemplates.chat_model.prompts
      return
    }

    draft.prompt_template = promptTemplates.completion_model.prompt
    setDefaultRolePrefix({
      user: promptTemplates.completion_model.conversation_histories_role.user_prefix,
      assistant: promptTemplates.completion_model.conversation_histories_role.assistant_prefix,
    })
  }, [isChatModel])

  useEffect(() => {
    const isReady = defaultConfig && Object.keys(defaultConfig).length > 0
    if (!isReady || inputs.prompt_template)
      return

    const nextInputs = produce(inputs, (draft) => {
      appendDefaultPromptConfig(draft, defaultConfig)
    })
    setInputs(nextInputs)
  }, [appendDefaultPromptConfig, defaultConfig, inputs, setInputs])

  return {
    inputRef,
    setInputs,
    appendDefaultPromptConfig,
  }
}

export default useLLMInputManager