import type {NodeProps} from '@xyflow/react'
import {Tag} from 'antd'
import AmisNode, {inputsFormColumns, outputsFormColumns} from './AmisNode.tsx'
import React from 'react'

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

const LlmNode = (props: NodeProps) => AmisNode({
  nodeProps: props,
  type: 'normal',
  defaultNodeName: '大模型',
  defaultNodeDescription: '使用大模型对话',
  extraNodeDescription: nodeData => {
    const model = nodeData?.model as string | undefined
    return model
      ? <div className="mt-2 flex justify-between">
        <span>大模型</span>
        <Tag className="m-0" color="blue">{modelMap[model]}</Tag>
      </div>
      : <></>
  },
  columnSchema: (props) => [
    ...inputsFormColumns(props),
    {
      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, {text: {type: 'string'}}),
  ],
})

export default React.memo(LlmNode)