From bbfe83c86b8c29ee758ba1adea17f5f46e6a68ed Mon Sep 17 00:00:00 2001
From: NFish <douxc512@gmail.com>
Date: Thu, 20 Feb 2025 14:31:56 +0800
Subject: [PATCH] Fix/upgrade btn show logic (#14072)

---
 .../workplace-selector/index.tsx              | 10 ---
 web/app/components/header/index.tsx           | 40 ++---------
 .../components/header/license-env/index.tsx   | 17 +++--
 .../components/header/plan-badge/index.tsx    | 70 +++++++++++++++++++
 4 files changed, 86 insertions(+), 51 deletions(-)
 create mode 100644 web/app/components/header/plan-badge/index.tsx

diff --git a/web/app/components/header/account-dropdown/workplace-selector/index.tsx b/web/app/components/header/account-dropdown/workplace-selector/index.tsx
index f8bde33d51..9649f2c70a 100644
--- a/web/app/components/header/account-dropdown/workplace-selector/index.tsx
+++ b/web/app/components/header/account-dropdown/workplace-selector/index.tsx
@@ -8,7 +8,6 @@ import { switchWorkspace } from '@/service/common'
 import { useWorkspacesContext } from '@/context/workspace-context'
 import { useProviderContext } from '@/context/provider-context'
 import { ToastContext } from '@/app/components/base/toast'
-import PremiumBadge from '@/app/components/base/premium-badge'
 
 const WorkplaceSelector = () => {
   const { t } = useTranslation()
@@ -72,15 +71,6 @@ const WorkplaceSelector = () => {
                       <div className='flex py-1 pl-3 pr-2 items-center gap-2 self-stretch hover:bg-state-base-hover rounded-lg' key={workspace.id} onClick={() => handleSwitchWorkspace(workspace.id)}>
                         <div className='flex items-center justify-center w-6 h-6 bg-[#EFF4FF] rounded-md text-xs font-medium text-primary-600'>{workspace.name[0].toLocaleUpperCase()}</div>
                         <div className='line-clamp-1 grow overflow-hidden text-text-secondary text-ellipsis system-md-regular cursor-pointer'>{workspace.name}</div>
-                        {
-                          <PremiumBadge size='s' color='gray' allowHover={false}>
-                            <div className='system-2xs-medium'>
-                              <span className='p-[2px]'>
-                                {plan.type === 'professional' ? 'PRO' : plan.type.toUpperCase()}
-                              </span>
-                            </div>
-                          </PremiumBadge>
-                        }
                       </div>
                     ))
                   }
diff --git a/web/app/components/header/index.tsx b/web/app/components/header/index.tsx
index d0807ff8ad..0e2df2e917 100644
--- a/web/app/components/header/index.tsx
+++ b/web/app/components/header/index.tsx
@@ -4,8 +4,6 @@ import Link from 'next/link'
 import { useBoolean } from 'ahooks'
 import { useSelectedLayoutSegment } from 'next/navigation'
 import { Bars3Icon } from '@heroicons/react/20/solid'
-import { SparklesSoft } from '@/app/components/base/icons/src/public/common'
-import PremiumBadge from '../base/premium-badge'
 import AccountDropdown from './account-dropdown'
 import AppNav from './app-nav'
 import DatasetNav from './dataset-nav'
@@ -13,7 +11,6 @@ import EnvNav from './env-nav'
 import PluginsNav from './plugins-nav'
 import ExploreNav from './explore-nav'
 import ToolsNav from './tools-nav'
-import LicenseNav from './license-env'
 import { WorkspaceProvider } from '@/context/workspace-context'
 import { useAppContext } from '@/context/app-context'
 import LogoSite from '@/app/components/base/logo/logo-site'
@@ -21,7 +18,9 @@ import WorkplaceSelector from '@/app/components/header/account-dropdown/workplac
 import useBreakpoints, { MediaType } from '@/hooks/use-breakpoints'
 import { useProviderContext } from '@/context/provider-context'
 import { useModalContext } from '@/context/modal-context'
-import { useTranslation } from 'react-i18next'
+import LicenseNav from './license-env'
+import PlanBadge from './plan-badge'
+import { Plan } from '../billing/type'
 
 const navClassName = `
   flex items-center relative mr-0 sm:mr-3 px-3 h-8 rounded-xl
@@ -31,15 +30,13 @@ const navClassName = `
 
 const Header = () => {
   const { isCurrentWorkspaceEditor, isCurrentWorkspaceDatasetOperator } = useAppContext()
-  const { t } = useTranslation()
-
   const selectedSegment = useSelectedLayoutSegment()
   const media = useBreakpoints()
   const isMobile = media === MediaType.mobile
   const [isShowNavMenu, { toggle, setFalse: hideNavMenu }] = useBoolean(false)
   const { enableBilling, plan } = useProviderContext()
   const { setShowPricingModal, setShowAccountSettingModal } = useModalContext()
-  const isFreePlan = plan.type === 'sandbox'
+  const isFreePlan = plan.type === Plan.sandbox
   const handlePlanClick = useCallback(() => {
     if (isFreePlan)
       setShowPricingModal()
@@ -71,18 +68,7 @@ const Header = () => {
               <WorkspaceProvider>
                 <WorkplaceSelector />
               </WorkspaceProvider>
-              {enableBilling && (
-                <div className='select-none'>
-                  <PremiumBadge color='blue' allowHover={true} onClick={handlePlanClick}>
-                    <SparklesSoft className='flex items-center py-[1px] pl-[3px] w-3.5 h-3.5 text-components-premium-badge-indigo-text-stop-0' />
-                    <div className='system-xs-medium'>
-                      <span className='p-1'>
-                        {t('billing.upgradeBtn.encourageShort')}
-                      </span>
-                    </div>
-                  </PremiumBadge>
-                </div>
-              )}
+              {enableBilling ? <PlanBadge allowHover sandboxAsUpgrade plan={plan.type} onClick={handlePlanClick} /> : <LicenseNav />}
             </div>
           </div>
         }
@@ -93,20 +79,7 @@ const Header = () => {
             <LogoSite />
           </Link>
           <div className='font-light text-divider-deep'>/</div>
-          {
-            enableBilling && (
-              <div className='select-none'>
-                <PremiumBadge color='blue' allowHover={true} onClick={handlePlanClick}>
-                  <SparklesSoft className='flex items-center py-[1px] pl-[3px] w-3.5 h-3.5 text-components-premium-badge-indigo-text-stop-0' />
-                  <div className='system-xs-medium'>
-                    <span className='p-1'>
-                      {t('billing.upgradeBtn.encourageShort')}
-                    </span>
-                  </div>
-                </PremiumBadge>
-              </div>
-            )
-          }
+          {enableBilling ? <PlanBadge allowHover sandboxAsUpgrade plan={plan.type} onClick={handlePlanClick} /> : <LicenseNav />}
         </div >
       )}
       {
@@ -120,7 +93,6 @@ const Header = () => {
         )
       }
       <div className='flex items-center shrink-0'>
-        <LicenseNav />
         <EnvNav />
         <div className='mr-3'>
           <PluginsNav />
diff --git a/web/app/components/header/license-env/index.tsx b/web/app/components/header/license-env/index.tsx
index 800d86d2b8..ca396db37a 100644
--- a/web/app/components/header/license-env/index.tsx
+++ b/web/app/components/header/license-env/index.tsx
@@ -5,6 +5,8 @@ import { LicenseStatus } from '@/types/feature'
 import { useTranslation } from 'react-i18next'
 import { useContextSelector } from 'use-context-selector'
 import dayjs from 'dayjs'
+import PremiumBadge from '../../base/premium-badge'
+import { RiHourglass2Fill } from '@remixicon/react'
 
 const LicenseNav = () => {
   const { t } = useTranslation()
@@ -13,15 +15,16 @@ const LicenseNav = () => {
   if (systemFeatures.license?.status === LicenseStatus.EXPIRING) {
     const expiredAt = systemFeatures.license?.expired_at
     const count = dayjs(expiredAt).diff(dayjs(), 'days')
-    return <div className='px-2 py-1 mr-4 rounded-full bg-util-colors-orange-orange-50 border-util-colors-orange-orange-100 system-xs-medium text-util-colors-orange-orange-600'>
-      {count <= 1 && <span>{t('common.license.expiring', { count })}</span>}
-      {count > 1 && <span>{t('common.license.expiring_plural', { count })}</span>}
-    </div>
+    return <PremiumBadge color='orange' className='select-none'>
+      <RiHourglass2Fill className='flex items-center pl-0.5 size-3 text-components-premium-badge-indigo-text-stop-0' />
+      {count <= 1 && <span className='system-xs-medium px-0.5'>{t('common.license.expiring', { count })}</span>}
+      {count > 1 && <span className='system-xs-medium px-0.5'>{t('common.license.expiring_plural', { count })}</span>}
+    </PremiumBadge>
   }
   if (systemFeatures.license.status === LicenseStatus.ACTIVE) {
-    return <div className='px-2 py-1 mr-4 rounded-md bg-util-colors-indigo-indigo-50 border-util-colors-indigo-indigo-100 system-xs-medium text-util-colors-indigo-indigo-600'>
-      Enterprise
-    </div>
+    return <PremiumBadge color="indigo" className='select-none'>
+      <span className='system-xs-medium px-1'>Enterprise</span>
+    </PremiumBadge>
   }
   return null
 }
diff --git a/web/app/components/header/plan-badge/index.tsx b/web/app/components/header/plan-badge/index.tsx
new file mode 100644
index 0000000000..5068e2da95
--- /dev/null
+++ b/web/app/components/header/plan-badge/index.tsx
@@ -0,0 +1,70 @@
+import { useProviderContext } from '@/context/provider-context'
+import classNames from '@/utils/classnames'
+import type { FC } from 'react'
+import { useTranslation } from 'react-i18next'
+import { SparklesSoft } from '../../base/icons/src/public/common'
+import PremiumBadge from '../../base/premium-badge'
+import { Plan } from '../../billing/type'
+
+type PlanBadgeProps = {
+  plan: Plan
+  size?: 's' | 'm'
+  allowHover?: boolean
+  sandboxAsUpgrade?: boolean
+  onClick?: () => void
+}
+
+const PlanBadge: FC<PlanBadgeProps> = ({ plan, allowHover, size = 'm', sandboxAsUpgrade = false, onClick }) => {
+  const { isFetchedPlan } = useProviderContext()
+  const { t } = useTranslation()
+
+  if (!isFetchedPlan) return null
+  if (plan === Plan.sandbox && sandboxAsUpgrade) {
+    return <div className='select-none'>
+      <PremiumBadge color='blue' allowHover={allowHover} onClick={onClick}>
+        <SparklesSoft className='flex items-center py-[1px] pl-[3px] w-3.5 h-3.5 text-components-premium-badge-indigo-text-stop-0' />
+        <div className='system-xs-medium'>
+          <span className='p-1'>
+            {t('billing.upgradeBtn.encourageShort')}
+          </span>
+        </div>
+      </PremiumBadge>
+    </div>
+  }
+  if (plan === Plan.sandbox) {
+    return <div className='select-none'>
+      <PremiumBadge size={size} color='gray' allowHover={allowHover} onClick={onClick}>
+        <div className={classNames(size === 's' ? 'system-2xs-medium-uppercase' : 'system-xs-medium-uppercase')}>
+          <span className='p-1'>
+            {plan}
+          </span>
+        </div>
+      </PremiumBadge>
+    </div>
+  }
+  if (plan === Plan.professional) {
+    return <div className='select-none'>
+      <PremiumBadge size={size} color='blue' allowHover={allowHover} onClick={onClick}>
+        <div className={classNames(size === 's' ? 'system-2xs-medium-uppercase' : 'system-xs-medium-uppercase')}>
+          <span className='p-1'>
+            pro
+          </span>
+        </div>
+      </PremiumBadge>
+    </div>
+  }
+  if (plan === Plan.team) {
+    return <div className='select-none'>
+      <PremiumBadge size={size} color='indigo' allowHover={allowHover} onClick={onClick}>
+        <div className={classNames(size === 's' ? 'system-2xs-medium-uppercase' : 'system-xs-medium-uppercase')}>
+          <span className='p-1'>
+            {plan}
+          </span>
+        </div>
+      </PremiumBadge>
+    </div>
+  }
+  return null
+}
+
+export default PlanBadge