import type {NodeProps} from '@xyflow/react'
import {Tag} from 'antd'
import React, {useCallback, useEffect, useMemo} from 'react'
import {useContextStore} from '../store/ContextStore.ts'
import {useDataStore} from '../store/DataStore.ts'
import {useFlowStore} from '../store/FlowStore.ts'
import AmisNode, {inputsFormColumns, nodeClassName, NormalNodeHandler, outputsFormColumns} from './AmisNode.tsx'

const modelMap: Record<string, string> = {
  qwen3: 'Qwen3',
  deepseek: 'Deepseek',
}

const LlmNode = (props: NodeProps) => {
  const {getNodes, getEdges} = useFlowStore()
  const {getData, mergeDataById, getDataById} = useDataStore()
  const {getInputSchema} = useContextStore()

  const nodeData = getDataById(props.id)

  useEffect(() => {
    mergeDataById(
      props.id,
      {
        outputs: {
          text: {
            type: 'text',
          },
        },
      },
    )
  }, [props.id])

  const columnsSchema = useCallback(() => [
    ...inputsFormColumns(props.id, getInputSchema(), getNodes(), getEdges(), getData()),
    {
      type: 'divider',
    },
    {
      type: 'select',
      name: 'model',
      label: '大模型',
      required: true,
      selectFirst: true,
      options: Object.keys(modelMap).map(key => ({label: modelMap[key], value: key})),
    },
    {
      type: 'textarea',
      name: 'systemPrompt',
      label: '系统提示词',
      required: true,
    },
    {
      type: 'divider',
    },
    ...outputsFormColumns(false, true),
  ], [props.id])

  const extraNodeDescription = useMemo(() => {
    return nodeData?.model
      ? <div className="mt-2 flex justify-between">
        <span>模型名称</span>
        <Tag className="m-0" color="blue">{modelMap[nodeData.model]}</Tag>
      </div>
      : <></>
  }, [nodeData])

  return (
    <AmisNode
      className={nodeClassName('llm')}
      nodeProps={props}
      extraNodeDescription={extraNodeDescription}
      columnSchema={columnsSchema}
      handler={<NormalNodeHandler/>}
    />
  )
}

export default React.memo(LlmNode)