jan/web/screens/Thread/ThreadCenterPanel/ChatInput/index.tsx

/* eslint-disable @typescript-eslint/no-explicit-any */
import { useContext, useEffect, useMemo, useRef, useState } from 'react'

import { InferenceEngine } from '@janhq/core'

import {
  TextArea,
  Button,
  Tooltip,
  useClickOutside,
  badgeVariants,
  Badge,
  Modal,
  Switch,
} from '@janhq/joi'
import { listen } from '@tauri-apps/api/event'
import { useAtom, useAtomValue } from 'jotai'
import {
  FileTextIcon,
  ImageIcon,
  StopCircle,
  PaperclipIcon,
  SettingsIcon,
  ChevronUpIcon,
  Settings2Icon,
  WrenchIcon,
} from 'lucide-react'

import { twMerge } from 'tailwind-merge'

import ModelDropdown from '@/containers/ModelDropdown'
import { currentPromptAtom, fileUploadAtom } from '@/containers/Providers/Jotai'

import { useActiveModel } from '@/hooks/useActiveModel'

import { useGetEngines } from '@/hooks/useEngineManagement'
import useSendChatMessage from '@/hooks/useSendChatMessage'

import { uploader } from '@/utils/file'
import { isLocalEngine } from '@/utils/modelEngine'

import { ChatContext } from '../../ThreadCenterPanel'
import { ChatContext } from '../../ThreadCenterPanel'
import FileUploadPreview from '../FileUploadPreview'
import ImageUploadPreview from '../ImageUploadPreview'

import RichTextEditor from './RichTextEditor'

import { showRightPanelAtom } from '@/helpers/atoms/App.atom'
import { experimentalFeatureEnabledAtom } from '@/helpers/atoms/AppConfig.atom'
import { activeAssistantAtom } from '@/helpers/atoms/Assistant.atom'
import { selectedModelAtom } from '@/helpers/atoms/Model.atom'
import { spellCheckAtom } from '@/helpers/atoms/Setting.atom'
import {
  activeSettingInputBoxAtom,
  activeThreadAtom,
  disabledThreadToolsAtom,
  getActiveThreadIdAtom,
  isBlockingSendAtom,
} from '@/helpers/atoms/Thread.atom'
import { activeTabThreadRightPanelAtom } from '@/helpers/atoms/ThreadRightPanel.atom'
import { ModelTool } from '@/types/model'

const ChatInput = () => {
  const activeThread = useAtomValue(activeThreadAtom)
  const { stateModel } = useActiveModel()
  const spellCheck = useAtomValue(spellCheckAtom)

  const [currentPrompt, setCurrentPrompt] = useAtom(currentPromptAtom)
  const [activeSettingInputBox, setActiveSettingInputBox] = useAtom(
    activeSettingInputBoxAtom
  )
  const { showApprovalModal } = useContext(ChatContext)
  const { sendChatMessage } = useSendChatMessage(showApprovalModal)
  const selectedModel = useAtomValue(selectedModelAtom)
  const [disabledTools, setDisableTools] = useAtom(disabledThreadToolsAtom)

  const activeThreadId = useAtomValue(getActiveThreadIdAtom)
  const [fileUpload, setFileUpload] = useAtom(fileUploadAtom)
  const [showAttachmentMenus, setShowAttachmentMenus] = useState(false)
  const textareaRef = useRef<HTMLTextAreaElement>(null)
  const fileInputRef = useRef<HTMLInputElement>(null)
  const imageInputRef = useRef<HTMLInputElement>(null)
  const { engines } = useGetEngines()
  const experimentalFeature = useAtomValue(experimentalFeatureEnabledAtom)
  const isBlockingSend = useAtomValue(isBlockingSendAtom)
  const activeAssistant = useAtomValue(activeAssistantAtom)
  const { stopInference } = useActiveModel()
  const [tools, setTools] = useState<any>([])
  const [showToolsModal, setShowToolsModal] = useState(false)

  const upload = uploader()
  const [activeTabThreadRightPanel, setActiveTabThreadRightPanel] = useAtom(
    activeTabThreadRightPanelAtom
  )

  const availableTools = useMemo(() => {
    return tools.filter((tool: ModelTool) => !disabledTools.includes(tool.name))
  }, [tools, disabledTools])

  const refAttachmentMenus = useClickOutside(() =>
    setShowAttachmentMenus(false)
  )
  const [showRightPanel, setShowRightPanel] = useAtom(showRightPanelAtom)

  useEffect(() => {
    if (textareaRef.current) {
      textareaRef.current.focus()
    }
  }, [activeThreadId])

  useEffect(() => {
    if (!selectedModel && !activeSettingInputBox) {
      setActiveSettingInputBox(true)
    }
  }, [activeSettingInputBox, selectedModel, setActiveSettingInputBox])

  useEffect(() => {
    window.core?.api?.getTools().then((data: ModelTool[]) => {
      setTools(data)
    })

    listen('mcp-update', (event) => {
      window.core?.api?.getTools().then((data: ModelTool[]) => {
        setTools(data)
      })
    })
  }, [])

  const onStopInferenceClick = async () => {
    stopInference()
  }

  const isModelSupportRagAndTools = isLocalEngine(
    engines,
    selectedModel?.engine as InferenceEngine
  )

  /**
   * Handles the change event of the extension file input element by setting the file name state.
   * Its to be used to display the extension file name of the selected file.
   * @param event - The change event object.
   */
  const handleFileChange = (event: React.ChangeEvent<HTMLInputElement>) => {
    const file = event.target.files?.[0]
    if (!file) return
    upload.addFile(file)
    upload.upload().then((data) => {
      setFileUpload({
        file: file,
        type: 'pdf',
        id: data?.successful?.[0]?.response?.body?.id,
        name: data?.successful?.[0]?.response?.body?.filename,
      })
    })
  }

  const handleImageChange = (event: React.ChangeEvent<HTMLInputElement>) => {
    const file = event.target.files?.[0]
    if (!file) return
    setFileUpload({ file: file, type: 'image' })
  }

  const renderPreview = (fileUpload: any) => {
    if (fileUpload) {
      if (fileUpload.type === 'image') {
        return <ImageUploadPreview file={fileUpload.file} />
      } else {
        return <FileUploadPreview />
      }
    }
  }

  return (
    <div className="relative p-4 pb-2">
      {renderPreview(fileUpload)}
      <div className="relative flex w-full flex-col">
        <RichTextEditor
          className={twMerge(
            'relative mb-1 max-h-[400px] resize-none rounded-lg border border-[hsla(var(--app-border))] p-3 pr-20',
            'focus-within:outline-none focus-visible:outline-0 focus-visible:ring-1 focus-visible:ring-[hsla(var(--primary-bg))] focus-visible:ring-offset-0',
            'overflow-y-auto',
            fileUpload && 'rounded-t-none border-t-0',
            experimentalFeature && 'pl-10',
            activeSettingInputBox && 'pb-14 pr-16'
          )}
          spellCheck={spellCheck}
          style={{ height: activeSettingInputBox ? '98px' : '44px' }}
          placeholder="Ask me anything"
          disabled={stateModel.loading || !activeThread}
        />
        <TextArea
          className="sr-only"
          data-testid="txt-input-chat"
          onChange={(e) => setCurrentPrompt(e.target.value)}
        />
        {experimentalFeature && (
          <Tooltip
            trigger={
              <Button
                theme="icon"
                className="absolute left-3 top-2.5"
                onClick={(e) => {
                  if (
                    !!fileUpload ||
                    (activeAssistant?.tools &&
                      !activeAssistant?.tools[0]?.enabled &&
                      !activeAssistant?.model.settings?.mmproj)
                  ) {
                    e.stopPropagation()
                  } else {
                    setShowAttachmentMenus(!showAttachmentMenus)
                  }
                }}
              >
                <PaperclipIcon
                  size={18}
                  className="text-[hsla(var(--text-secondary))]"
                />
              </Button>
            }
            disabled={
              !isModelSupportRagAndTools ||
              (activeAssistant?.tools && activeAssistant?.tools[0]?.enabled)
            }
            content={
              <>
                {!!fileUpload ||
                  (activeAssistant?.tools &&
                    !activeAssistant?.tools[0]?.enabled &&
                    !activeAssistant?.model.settings?.mmproj && (
                      <>
                        {!!fileUpload && (
                          <span>
                            Currently, we only support 1 attachment at the same
                            time.
                          </span>
                        )}
                        {activeAssistant?.tools &&
                          activeAssistant?.tools[0]?.enabled === false &&
                          isModelSupportRagAndTools && (
                            <span>
                              Turn on Retrieval in Tools settings to use this
                              feature
                            </span>
                          )}
                        {!isModelSupportRagAndTools && (
                          <span>Not supported for this model</span>
                        )}
                      </>
                    ))}
              </>
            }
          />
        )}

        {showAttachmentMenus && (
          <div
            ref={refAttachmentMenus}
            className={twMerge(
              'absolute bottom-14 left-0 z-30 w-36 cursor-pointer rounded-lg border border-[hsla(var(--app-border))] bg-[hsla(var(--app-bg))] py-1 shadow-sm',
              activeSettingInputBox && 'bottom-28'
            )}
          >
            <ul>
              <li
                className={twMerge(
                  'text-[hsla(var(--text-secondary)] hover:bg-secondary flex w-full items-center space-x-2 px-4 py-2 hover:bg-[hsla(var(--dropdown-menu-hover-bg))]',
                  activeAssistant?.model.settings?.mmproj &&
                    isModelSupportRagAndTools
                    ? 'cursor-pointer'
                    : 'cursor-not-allowed opacity-50'
                )}
                onClick={() => {
                  if (activeAssistant?.model.settings?.mmproj) {
                    imageInputRef.current?.click()
                    setShowAttachmentMenus(false)
                  }
                }}
              >
                <ImageIcon size={16} />
                <span className="font-medium">Image</span>
              </li>
              <Tooltip
                side="bottom"
                trigger={
                  <li
                    className={twMerge(
                      'text-[hsla(var(--text-secondary)] hover:bg-secondary flex w-full cursor-pointer items-center space-x-2 px-4 py-2 hover:bg-[hsla(var(--dropdown-menu-hover-bg))]',
                      isModelSupportRagAndTools
                        ? 'cursor-pointer'
                        : 'cursor-not-allowed opacity-50'
                    )}
                    onClick={() => {
                      if (isModelSupportRagAndTools) {
                        fileInputRef.current?.click()
                        setShowAttachmentMenus(false)
                      }
                    }}
                  >
                    <FileTextIcon size={16} />
                    <span className="font-medium">Document</span>
                  </li>
                }
                content={
                  (!activeAssistant?.tools ||
                    !activeAssistant?.tools[0]?.enabled) && (
                    <span>
                      Turn on Retrieval in Assistant Settings to use this
                      feature.
                    </span>
                  )
                }
              />
            </ul>
          </div>
        )}

        <div className={twMerge('absolute right-3 top-1.5')}>
          <div className="flex items-center gap-x-4">
            {!activeSettingInputBox && (
              <div className="flex h-8 items-center">
                <Button
                  theme="icon"
                  onClick={() => {
                    setActiveSettingInputBox(!activeSettingInputBox)
                  }}
                >
                  <SettingsIcon
                    size={18}
                    className="text-[hsla(var(--text-secondary))]"
                  />
                </Button>
              </div>
            )}

            {!isBlockingSend ? (
              <>
                {currentPrompt.length !== 0 && (
                  <Button
                    disabled={
                      stateModel.loading ||
                      !activeThread ||
                      currentPrompt.trim().length === 0
                    }
                    className="h-8 w-8 rounded-lg p-0"
                    data-testid="btn-send-chat"
                    onClick={() => sendChatMessage(currentPrompt)}
                  >
                    <svg
                      width="16"
                      height="16"
                      viewBox="0 0 16 16"
                      fill="none"
                      xmlns="http://www.w3.org/2000/svg"
                      className="fill-white stroke-white"
                    >
                      <path
                        d="M3.93098 4.26171L3.93108 4.26168L12.9041 1.27032C12.9041 1.27031 12.9041 1.27031 12.9041 1.27031C13.7983 0.972243 14.3972 0.77445 14.8316 0.697178C15.0428 0.659595 15.1663 0.660546 15.2355 0.671861C15.2855 0.680033 15.296 0.690905 15.3015 0.696542C15.3018 0.696895 15.3022 0.697228 15.3025 0.697538C15.3028 0.697847 15.3031 0.698168 15.3035 0.698509C15.3091 0.703965 15.32 0.71449 15.3282 0.764538C15.3395 0.8338 15.3405 0.957246 15.3029 1.16844C15.2258 1.60268 15.0282 2.20131 14.7307 3.09505L11.7383 12.0689L11.7383 12.069C11.3184 13.3293 11.0242 14.2078 10.7465 14.7789C10.6083 15.063 10.4994 15.2158 10.4215 15.292C10.3948 15.3182 10.3774 15.3295 10.3698 15.3338C10.3622 15.3295 10.3449 15.3181 10.3184 15.2921C10.2404 15.2158 10.1314 15.0629 9.99319 14.7788C9.71539 14.2077 9.42091 13.3291 9.00105 12.069L9.00094 12.0687L8.34059 10.0903L12.6391 5.79172L12.6392 5.7918L12.6472 5.78348C12.9604 5.45927 13.1337 5.02503 13.1297 4.57431C13.1258 4.12358 12.945 3.69242 12.6263 3.3737C12.3076 3.05497 11.8764 2.87418 11.4257 2.87027C10.975 2.86635 10.5407 3.03962 10.2165 3.35276L10.2165 3.35268L10.2083 3.36086L5.9106 7.65853L3.93098 6.99895C2.67072 6.57904 1.79218 6.28485 1.22115 6.00715C0.937001 5.86898 0.784237 5.76011 0.707981 5.68215C0.681839 5.65542 0.670463 5.63807 0.666163 5.63051C0.670529 5.62288 0.681934 5.60558 0.707909 5.57904C0.784233 5.50103 0.937088 5.3921 1.22125 5.25386C1.79226 4.97606 2.67087 4.68157 3.93098 4.26171Z"
                        strokeWidth="1.33"
                      />
                    </svg>
                  </Button>
                )}
              </>
            ) : (
              <Button
                theme="destructive"
                onClick={onStopInferenceClick}
                className="h-8 w-8 rounded-lg p-0"
              >
                <StopCircle size={20} />
              </Button>
            )}
          </div>
        </div>

        <div
          className={twMerge(
            'absolute bottom-[5px] left-[1px] flex w-[calc(100%-10px)] items-center justify-between rounded-b-lg bg-[hsla(var(--center-panel-bg))] p-3 pr-0',
            !activeThread && 'bg-transparent',
            !activeSettingInputBox && 'hidden',
            stateModel.loading && 'bg-transparent'
          )}
        >
          <div className="flex items-center gap-x-2">
            <ModelDropdown chatInputMode />
            <button
              className={twMerge(
                badgeVariants({
                  theme: 'secondary',
                  variant:
                    activeTabThreadRightPanel === 'model' ? 'solid' : 'outline',
                }),
                'flex cursor-pointer items-center gap-x-1',
                activeTabThreadRightPanel === 'model' &&
                  'border border-transparent'
              )}
              onClick={() => {
                // TODO @faisal: should be refactor later and better experience beetwen tab and toggle button
                if (showRightPanel && activeTabThreadRightPanel !== 'model') {
                  setShowRightPanel(true)
                  setActiveTabThreadRightPanel('model')
                }
                if (showRightPanel && activeTabThreadRightPanel === 'model') {
                  setShowRightPanel(false)
                  setActiveTabThreadRightPanel(undefined)
                }
                if (activeTabThreadRightPanel === undefined) {
                  setShowRightPanel(true)
                  setActiveTabThreadRightPanel('model')
                }
                if (!showRightPanel && activeTabThreadRightPanel !== 'model') {
                  setShowRightPanel(true)
                  setActiveTabThreadRightPanel('model')
                }
              }}
            >
              <Settings2Icon
                size={16}
                className="flex-shrink-0 cursor-pointer text-[hsla(var(--text-secondary))]"
              />
            </button>
            {tools && tools.length > 0 && (
              <>
                <Badge
                  theme="secondary"
                  className={twMerge(
                    'flex cursor-pointer items-center gap-x-1'
                  )}
                  variant={'outline'}
                  onClick={() => setShowToolsModal(true)}
                >
                  <WrenchIcon
                    size={16}
                    className="flex-shrink-0 cursor-pointer text-[hsla(var(--text-secondary))]"
                  />
                  <span className="text-xs">{availableTools.length}</span>
                </Badge>

                <Modal
                  open={showToolsModal}
                  onOpenChange={setShowToolsModal}
                  title="Available MCP Tools"
                  content={
                    <div className="overflow-y-auto">
                      <div className="mb-2 py-2 text-sm text-[hsla(var(--text-secondary))]">
                        Jan can use tools provided by specialized servers using
                        Model Context Protocol.{' '}
                        <a
                          href="https://modelcontextprotocol.io/introduction"
                          target="_blank"
                          className="text-[hsla(var(--app-link))]"
                        >
                          Learn more about MCP
                        </a>
                      </div>
                      {tools.map((tool: any) => (
                        <div
                          key={tool.name}
                          className="flex items-center gap-x-3 px-4 py-3 hover:bg-[hsla(var(--dropdown-menu-hover-bg))]"
                        >
                          <WrenchIcon
                            size={16}
                            className="flex-shrink-0 text-[hsla(var(--text-secondary))]"
                          />
                          <div className="flex w-full flex-1 flex-row justify-between">
                            <div>
                              <div className="font-medium">{tool.name}</div>
                              <div className="text-sm text-[hsla(var(--text-secondary))]">
                                {tool.description}
                              </div>
                            </div>
                            <Switch
                              checked={!disabledTools.includes(tool.name)}
                              onChange={(e) => {
                                if (e.target.checked) {
                                  setDisableTools((prev) =>
                                    prev.filter((t) => t !== tool.name)
                                  )
                                } else {
                                  setDisableTools([...disabledTools, tool.name])
                                }
                              }}
                              className="flex-shrink-0"
                            />
                          </div>
                        </div>
                      ))}
                    </div>
                  }
                />
              </>
            )}
          </div>
          {selectedModel && (
            <Button
              theme="icon"
              onClick={() => setActiveSettingInputBox(false)}
            >
              <ChevronUpIcon
                size={16}
                className="cursor-pointer text-[hsla(var(--text-secondary))]"
              />
            </Button>
          )}
        </div>
      </div>

      <input
        type="file"
        className="hidden"
        ref={imageInputRef}
        value=""
        onChange={handleImageChange}
        accept="image/png, image/jpeg, image/jpg"
      />

      <input
        type="file"
        className="hidden"
        ref={fileInputRef}
        value=""
        onChange={handleFileChange}
        accept="application/pdf"
      />
    </div>
  )
}

export default ChatInput