From 576dfcc2d4f2c6721b79d2e9b5a44bb566a7ecf9 Mon Sep 17 00:00:00 2001
From: Hardikk Kamboj <64458111+khardikk@users.noreply.github.com>
Date: Wed, 13 Aug 2025 01:53:44 +0530
Subject: [PATCH] Feat: label creation flow fix (#1908)

Co-authored-by: Adam <13007539+MrgSub@users.noreply.github.com>
---
 .../components/context/thread-context.tsx     | 63 ++++++++++++++++++-
 apps/mail/components/labels/label-dialog.tsx  | 14 +++--
 apps/mail/components/ui/nav-main.tsx          | 24 ++++---
 3 files changed, 86 insertions(+), 15 deletions(-)

diff --git a/apps/mail/components/context/thread-context.tsx b/apps/mail/components/context/thread-context.tsx
index 76a70c9d6..2785a5cc0 100644
--- a/apps/mail/components/context/thread-context.tsx
+++ b/apps/mail/components/context/thread-context.tsx
@@ -21,15 +21,17 @@ import {
   Star,
   StarOff,
   Tag,
+  Plus,
   Trash,
 } from 'lucide-react';
 import { useOptimisticThreadState } from '@/components/mail/optimistic-thread-state';
+import { LabelDialog } from '@/components/labels/label-dialog';
 import { useOptimisticActions } from '@/hooks/use-optimistic-actions';
 import { ExclamationCircle, Mail, Clock } from '../icons/icons';
 import { SnoozeDialog } from '@/components/mail/snooze-dialog';
 import { type ThreadDestination } from '@/lib/thread-actions';
 import { useThread, useThreads } from '@/hooks/use-threads';
-import { useMemo, type ReactNode, useState } from 'react';
+import { useMemo, type ReactNode, useState, useCallback } from 'react';
 import { useTRPC } from '@/providers/query-provider';
 import { useMutation } from '@tanstack/react-query';
 import { useLabels } from '@/hooks/use-labels';
@@ -40,6 +42,7 @@ import { m } from '@/paraglide/messages';
 import { useParams } from 'react-router';
 import { useQueryState } from 'nuqs';
 import { toast } from 'sonner';
+import type { Label as LabelType } from '@/types';
 
 interface EmailAction {
   id: string;
@@ -61,7 +64,7 @@ interface EmailContextMenuProps {
   refreshCallback?: () => void;
 }
 
-const LabelsList = ({ threadId, bulkSelected }: { threadId: string; bulkSelected: string[] }) => {
+const LabelsList = ({ threadId, bulkSelected, onCreateLabel }: { threadId: string; bulkSelected: string[]; onCreateLabel: () => void }) => {
   const { userLabels: labels } = useLabels();
   const { optimisticToggleLabel } = useOptimisticActions();
   const targetThreadIds = bulkSelected.length > 0 ? bulkSelected : [threadId];
@@ -90,6 +93,19 @@ const LabelsList = ({ threadId, bulkSelected }: { threadId: string; bulkSelected
     optimisticToggleLabel(targetThreadIds, labelId, !hasLabel);
   };
 
+  // If no labels exist, show create label button
+  if (!labels || labels.length === 0) {
+    return (
+      <ContextMenuItem 
+        onClick={onCreateLabel}
+        className="font-normal"
+      >
+        <Plus className="mr-2 h-4 w-4 opacity-60" />
+        {m['common.mail.createNewLabel']()}
+      </ContextMenuItem>
+    );
+  }
+
   return (
     <>
       {labels
@@ -148,6 +164,7 @@ export function ThreadContextMenu({
   const [, setActiveReplyId] = useQueryState('activeReplyId');
   const optimisticState = useOptimisticThreadState(threadId);
   const trpc = useTRPC();
+  const { refetch: refetchLabels } = useLabels();
   const {
     optimisticMoveThreadsTo,
     optimisticToggleStar,
@@ -159,6 +176,7 @@ export function ThreadContextMenu({
     optimisticUnsnooze,
   } = useOptimisticActions();
   const { mutateAsync: deleteThread } = useMutation(trpc.mail.delete.mutationOptions());
+  const { mutateAsync: createLabel } = useMutation(trpc.labels.create.mutationOptions());
 
   const { isUnread, isStarred, isImportant } = useMemo(() => {
     const unread = threadData?.hasUnread ?? false;
@@ -454,6 +472,11 @@ export function ThreadContextMenu({
   }, [isSpam, isBin, isArchiveFolder, isInbox, isSent, handleMove, handleDelete]);
 
   const [snoozeOpen, setSnoozeOpen] = useState(false);
+  const [createLabelOpen, setCreateLabelOpen] = useState(false);
+
+  const handleOpenCreateLabel = useCallback(() => {
+    setCreateLabelOpen(true);
+  }, []);
 
   const handleSnoozeConfirm = (wakeAt: Date) => {
     const targets = mail.bulkSelected.length ? mail.bulkSelected : [threadId];
@@ -461,6 +484,35 @@ export function ThreadContextMenu({
     setSnoozeOpen(false);
   };
 
+  const handleCreateLabel = async (data: LabelType) => {
+    const labelData = {
+      name: data.name,
+      color: {
+        backgroundColor: data.color?.backgroundColor || '#202020',
+        textColor: data.color?.textColor || '#FFFFFF'
+      }
+    };
+    
+    try {
+      const promise = createLabel(labelData).then(async (result) => {
+        await refetchLabels();
+        return result;
+      });
+      
+      toast.promise(promise, {
+        loading: m['common.labels.savingLabel'](),
+        success: m['common.labels.saveLabelSuccess'](),
+        error: m['common.labels.failedToSavingLabel'](),
+      });
+      
+      await promise;
+    } catch (error) {
+      console.error('Failed to create label:', error);
+    } finally {
+      setCreateLabelOpen(false);
+    }
+  };
+
   const otherActions: EmailAction[] = useMemo(
     () => [
       {
@@ -520,6 +572,11 @@ export function ThreadContextMenu({
 
   return (
     <>
+      <LabelDialog
+        open={createLabelOpen}
+        onOpenChange={setCreateLabelOpen}
+        onSubmit={handleCreateLabel}
+      />
       <ContextMenu>
         <ContextMenuTrigger disabled={isLoading || isFetching} className="w-full">
           {children}
@@ -538,7 +595,7 @@ export function ThreadContextMenu({
               {m['common.mail.labels']()}
             </ContextMenuSubTrigger>
             <ContextMenuSubContent className="dark:bg-panelDark max-h-[520px] w-48 overflow-y-auto bg-white">
-              <LabelsList threadId={threadId} bulkSelected={mail.bulkSelected} />
+              <LabelsList threadId={threadId} bulkSelected={mail.bulkSelected} onCreateLabel={handleOpenCreateLabel} />
             </ContextMenuSubContent>
           </ContextMenuSub>
 
diff --git a/apps/mail/components/labels/label-dialog.tsx b/apps/mail/components/labels/label-dialog.tsx
index d26c3e5dd..da6aa1ff4 100644
--- a/apps/mail/components/labels/label-dialog.tsx
+++ b/apps/mail/components/labels/label-dialog.tsx
@@ -1,6 +1,7 @@
 import {
   Dialog,
   DialogContent,
+  DialogDescription,
   DialogHeader,
   DialogTitle,
   DialogTrigger,
@@ -64,12 +65,12 @@ export function LabelDialog({
       if (editingLabel) {
         form.reset({
           name: editingLabel.name,
-          color: editingLabel.color || { backgroundColor: '#E2E2E2', textColor: '#000000' },
+          color: editingLabel.color || { backgroundColor: '#202020', textColor: '#FFFFFF' },
         });
       } else {
         form.reset({
           name: '',
-          color: { backgroundColor: '#E2E2E2', textColor: '#000000' },
+          color: { backgroundColor: '#202020', textColor: '#FFFFFF' },
         });
       }
     }
@@ -85,7 +86,7 @@ export function LabelDialog({
     setDialogOpen(false);
     form.reset({
       name: '',
-      color: { backgroundColor: '#E2E2E2', textColor: '#000000' },
+      color: { backgroundColor: '#202020', textColor: '#FFFFFF' },
     });
   };
 
@@ -97,6 +98,11 @@ export function LabelDialog({
           <DialogTitle>
             {editingLabel ? m['common.labels.editLabel']() : m['common.mail.createNewLabel']()}
           </DialogTitle>
+          <DialogDescription>
+            {editingLabel
+              ? 'Modify the label name and color to update this label.'
+              : 'Create a new label to organize your emails. Choose a name and color for easy identification.'}
+          </DialogDescription>
         </DialogHeader>
         <Form {...form}>
           <form
@@ -126,7 +132,7 @@ export function LabelDialog({
               <div className="space-y-2">
                 <Label>{m['common.labels.color']()}</Label>
                 <div className="w-full">
-                  <div className="flex flex-wrap gap-2">
+                  <div className="flex flex-wrap gap-2 mt-2">
                     {LABEL_COLORS.map((color) => (
                       <button
                         key={color.backgroundColor}
diff --git a/apps/mail/components/ui/nav-main.tsx b/apps/mail/components/ui/nav-main.tsx
index 392a3bbe7..254e83996 100644
--- a/apps/mail/components/ui/nav-main.tsx
+++ b/apps/mail/components/ui/nav-main.tsx
@@ -163,14 +163,22 @@ export function NavMain({ items }: NavMainProps) {
   );
 
   const onSubmit = async (data: LabelType) => {
-    toast.promise(createLabel(data), {
-      loading: 'Creating label...',
-      success: 'Label created successfully',
-      error: 'Failed to create label',
-      finally: () => {
-        refetch();
-      },
-    });
+    try {
+      const promise = createLabel(data).then(async (result) => {
+        await refetch();
+        return result;
+      });
+      
+      toast.promise(promise, {
+        loading: 'Creating label...',
+        success: 'Label created successfully',
+        error: 'Failed to create label',
+      });
+      
+      await promise;
+    } catch (error) {
+      console.error('Failed to create label:', error);
+    }
   };
 
   return (