dify/web/app/components/custom/custom-web-app-brand/hooks/use-web-app-brand.ts

import type { ChangeEvent } from 'react'
import { toast } from '@langgenius/dify-ui/toast'
import { useSuspenseQuery } from '@tanstack/react-query'
import { useAtomValue, useSetAtom } from 'jotai'
import { useState } from 'react'
import { useTranslation } from 'react-i18next'
import { getImageUploadErrorMessage, imageUpload } from '@/app/components/base/image-uploader/utils'
import { Plan } from '@/app/components/billing/type'
import { workspacePermissionKeysAtom } from '@/context/permission-state'
import { useProviderContext } from '@/context/provider-context'
import { currentWorkspaceAtom, refreshCurrentWorkspaceAtom } from '@/context/workspace-state'
import { systemFeaturesQueryOptions } from '@/features/system-features/client'
import { updateCurrentWorkspace } from '@/service/common'
import { hasPermission } from '@/utils/permission'

const MAX_LOGO_FILE_SIZE = 5 * 1024 * 1024
const CUSTOM_CONFIG_URL = '/workspaces/custom-config'
const WEB_APP_LOGO_UPLOAD_URL = '/workspaces/custom-config/webapp-logo/upload'
const useWebAppBrand = () => {
  const { t } = useTranslation()
  const { plan, enableBilling } = useProviderContext()
  const currentWorkspace = useAtomValue(currentWorkspaceAtom)
  const mutateCurrentWorkspace = useSetAtom(refreshCurrentWorkspaceAtom)
  const workspacePermissionKeys = useAtomValue(workspacePermissionKeysAtom)
  const [fileId, setFileId] = useState('')
  const [imgKey, setImgKey] = useState(() => Date.now())
  const [uploadProgress, setUploadProgress] = useState(0)
  const { data: systemFeatures } = useSuspenseQuery(systemFeaturesQueryOptions())
  const isSandbox = enableBilling && plan.type === Plan.sandbox
  const uploading = uploadProgress > 0 && uploadProgress < 100
  const webappLogo = currentWorkspace.custom_config?.replace_webapp_logo || ''
  const webappBrandRemoved = currentWorkspace.custom_config?.remove_webapp_brand
  const canManageCustomBrand = hasPermission(workspacePermissionKeys, 'customization.manage')
  const uploadDisabled = isSandbox || webappBrandRemoved || !canManageCustomBrand
  const workspaceLogo = systemFeatures.branding.enabled
    ? systemFeatures.branding.workspace_logo
    : ''
  const persistWorkspaceBrand = async (body: Record<string, unknown>) => {
    await updateCurrentWorkspace({
      url: CUSTOM_CONFIG_URL,
      body,
    })
    mutateCurrentWorkspace()
  }
  const handleChange = (e: ChangeEvent<HTMLInputElement>) => {
    const file = e.target.files?.[0]
    if (!file) return
    if (file.size > MAX_LOGO_FILE_SIZE) {
      toast.error(t(($) => $['imageUploader.uploadFromComputerLimit'], { ns: 'common', size: 5 }))
      return
    }
    imageUpload(
      {
        file,
        onProgressCallback: setUploadProgress,
        onSuccessCallback: (res) => {
          setUploadProgress(100)
          setFileId(res.id)
        },
        onErrorCallback: (error) => {
          const errorMessage = getImageUploadErrorMessage(
            error,
            t(($) => $['imageUploader.uploadFromComputerUploadError'], { ns: 'common' }),
            t,
          )
          toast.error(errorMessage)
          setUploadProgress(-1)
        },
      },
      false,
      WEB_APP_LOGO_UPLOAD_URL,
    )
  }
  const handleApply = async () => {
    await persistWorkspaceBrand({
      remove_webapp_brand: webappBrandRemoved,
      replace_webapp_logo: fileId,
    })
    setFileId('')
    setImgKey(Date.now())
  }
  const handleRestore = async () => {
    await persistWorkspaceBrand({
      remove_webapp_brand: false,
      replace_webapp_logo: '',
    })
  }
  const handleSwitch = async (checked: boolean) => {
    await persistWorkspaceBrand({
      remove_webapp_brand: checked,
    })
  }
  const handleCancel = () => {
    setFileId('')
    setUploadProgress(0)
  }
  return {
    fileId,
    imgKey,
    uploadProgress,
    uploading,
    webappLogo,
    webappBrandRemoved,
    uploadDisabled,
    workspaceLogo,
    isSandbox,
    canManageCustomBrand,
    handleApply,
    handleCancel,
    handleChange,
    handleRestore,
    handleSwitch,
  }
}
export default useWebAppBrand