improvement: treatment plan turned into a wizard. shift+click control added to FDI tooth chart for cunnected prosthesises.

frontend/src/components/ui/treatment/TreatmentWorkspace.tsx

@@ -5,6 +5,7 @@ import { useLocale, useTranslations } from 'next-intl';
 import { useRouter } from '@/i18n/navigation';
 import { Button } from '@/components/ui/shared/Button';
 import { Checkbox } from '@/components/ui/shared/Checkbox';
+import { WizardStepper } from '@/components/ui/shared/WizardStepper';
 import { PatientSearchCombobox } from '@/components/ui/patient/PatientSearchCombobox';
 import {
   TreatmentLabCasesPanel,
@@ -39,6 +40,14 @@ import {
   isEmptyDraftDetail,
   isLabDependentDetailMissingTeeth,
 } from '@/components/treatment/treatmentDetailRules';
+import {
+  applyShiftRange,
+  connectedTeethSet,
+  deriveTeethFromGroups,
+  groupsFromFlatTeeth,
+  pruneToothProsthesisForGroups,
+  toggleToothInGroups,
+} from '@/components/treatment/toothSelectionGroups';
 import type { LabDispatchAttentionItem } from '@/components/treatment/labDispatchAttention';
 import { collectLabDispatchAttention } from '@/components/treatment/labDispatchAttention';
 import { canEditTreatment, canViewTreatment, canAccessDashboardRoute } from '@/components/shared/permissions';
@@ -68,6 +77,9 @@ import type {
 import type { PatientLabCaseSummary } from '@/types/lab-case-activity';
 
 type WorkspaceMode = 'live' | 'historical';
+type EntryStep = 'teeth' | 'content' | 'lab';
+
+const ENTRY_STEPS: EntryStep[] = ['teeth', 'content', 'lab'];
 
 function isTreatmentDayHistorical(treatmentAt: string, todayStart: Date): boolean {
   return compareLocalDayStart(new Date(treatmentAt), todayStart) < 0;
@@ -98,6 +110,7 @@ function labCaseDraftsToPast(
             clientId: linkedDetail.clientId,
             treatmentType: linkedDetail.treatmentType,
             teeth: linkedDetail.teeth,
+            toothSelectionGroups: linkedDetail.toothSelectionGroups,
           }
         : null,
       sends: lc.sends ?? [],
@@ -154,6 +167,7 @@ function newDetail(defaultTreatmentType?: string): TreatmentDetailDraft {
         : `detail-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`,
     treatmentType: defaultTreatmentType ?? '',
     teeth: [],
+    toothSelectionGroups: [],
     comment: '',
     attachmentMetas: [],
     sendToOrganizationIds: [],
@@ -190,11 +204,14 @@ function mapAppointment(record: AppointmentRecord): TreatmentAppointment {
 }
 
 function mapDetailFromApi(d: PastTreatmentCase): TreatmentDetailDraft {
+  const teeth = d.teeth;
+  const toothSelectionGroups = groupsFromFlatTeeth(teeth, d.toothSelectionGroups ?? null);
   return {
     clientId: d.clientId,
     id: d.id,
     treatmentType: d.treatmentType,
-    teeth: d.teeth,
+    teeth,
+    toothSelectionGroups,
     comment: d.notes ?? '',
     attachmentMetas: d.attachmentMetas ?? [],
     labCaseId: d.labCaseId ?? null,
@@ -215,6 +232,7 @@ function mapLabCaseDraftFromApi(lc: PastLabCase): LabCaseDraft {
       detailClientId: lc.detail?.clientId ?? tp.treatmentDetailId,
       tooth: tp.tooth,
       prosthesisTypeCode: tp.prosthesisTypeCode,
+      selectionGroupId: tp.selectionGroupId ?? '',
     })),
     attachmentIds: (lc.attachments ?? []).map((a) => a.id),
     sentAt: lc.sentAt ?? null,
@@ -231,6 +249,7 @@ function serializeDetails(details: TreatmentDetailDraft[]) {
       id: d.id,
       treatmentType: d.treatmentType,
       teeth: d.teeth,
+      toothSelectionGroups: d.toothSelectionGroups,
       comment: d.comment,
       attachmentMetas: d.attachmentMetas,
     })),
@@ -261,6 +280,7 @@ function detailsToPreviewTreatment(
       clientId: d.clientId,
       treatmentType: d.treatmentType,
       teeth: d.teeth,
+      toothSelectionGroups: d.toothSelectionGroups,
       notes: d.comment || null,
       attachmentMetas: d.attachmentMetas,
       labCaseId: d.labCaseId ?? null,
@@ -380,6 +400,7 @@ export function TreatmentWorkspace({
   const [organizationSearch, setOrganizationSearch] = useState('');
   const [recentOrganizationIds, setRecentOrganizationIds] = useState<string[]>([]);
   const [showWholeTreatmentPlan, setShowWholeTreatmentPlan] = useState(false);
+  const [entryStep, setEntryStep] = useState<EntryStep>('teeth');
 
   const isDetailLocked = useCallback(
     (detail: TreatmentDetailDraft) =>
@@ -478,6 +499,20 @@ export function TreatmentWorkspace({
     [details, labDependentCodes],
   );
 
+  /** Lab wizard step only for prosthesis (lab-dependent) treatment types on the active detail. */
+  const showLabWizardStep = useMemo(
+    () => Boolean(activeDetail && labDependentCodes.has(activeDetail.treatmentType)),
+    [activeDetail, labDependentCodes],
+  );
+
+  const visibleEntrySteps = useMemo(
+    () =>
+      showLabWizardStep
+        ? ENTRY_STEPS
+        : (ENTRY_STEPS.filter((step) => step !== 'lab') as EntryStep[]),
+    [showLabWizardStep],
+  );
+
   const showLabShipmentBlocked = useMemo(
     () =>
       Boolean(
@@ -546,6 +581,11 @@ export function TreatmentWorkspace({
   }, []);
 
   const selectedTeethSet = useMemo(() => new Set(activeDetail?.teeth ?? []), [activeDetail?.teeth]);
+  const connectedSelectedTeeth = useMemo(
+    () => connectedTeethSet(activeDetail?.toothSelectionGroups ?? []),
+    [activeDetail?.toothSelectionGroups],
+  );
+  const rangeAnchorRef = useRef<FdiToothId | null>(null);
 
   const wholePlanTeethSet = useMemo(() => {
     const set = new Set<FdiToothId>();
@@ -574,8 +614,17 @@ export function TreatmentWorkspace({
   // Reset whole-plan overview when switching details.
   useEffect(() => {
     setShowWholeTreatmentPlan(false);
+    rangeAnchorRef.current = null;
+    setEntryStep('teeth');
   }, [activeDetailId]);
 
+  // Leave Lab step if the active detail is no longer prosthesis / lab-dependent.
+  useEffect(() => {
+    if (entryStep === 'lab' && !showLabWizardStep) {
+      setEntryStep('content');
+    }
+  }, [entryStep, showLabWizardStep]);
+
   // Sync active lab shipment when the selected treatment detail changes.
   useEffect(() => {
     const match = labCaseDrafts.find((lc) => lc.detailClientId === activeDetailId);
@@ -864,14 +913,17 @@ export function TreatmentWorkspace({
       }
 
       const response = await treatmentsApi.saveDraft(selectedAppointment.id, {
-        details: currentDetails.map(({ clientId, id, treatmentType, teeth, comment, attachmentMetas }) => ({
-          clientId,
-          id,
-          treatmentType,
-          teeth,
-          comment,
-          attachmentIds: attachmentMetas.map((a) => a.id),
-        })),
+        details: currentDetails.map(
+          ({ clientId, id, treatmentType, teeth, toothSelectionGroups, comment, attachmentMetas }) => ({
+            clientId,
+            id,
+            treatmentType,
+            teeth,
+            toothSelectionGroups,
+            comment,
+            attachmentIds: attachmentMetas.map((a) => a.id),
+          }),
+        ),
       });
       const mapped = response.data.details.map(mapDetailFromApi);
       setDetails(mapped);
@@ -1081,6 +1133,7 @@ export function TreatmentWorkspace({
           setActiveLabCaseId(linked.clientId);
         }
         if (options?.scrollToLabPanel !== false) {
+          setEntryStep('lab');
           requestAnimationFrame(() => {
             scrollWithinMainScrollContainer(labPanelRef.current);
           });
@@ -1140,6 +1193,7 @@ export function TreatmentWorkspace({
         if (linked) {
           setActiveLabCaseId(linked.clientId);
         }
+        setEntryStep('lab');
         requestAnimationFrame(() => {
           scrollWithinMainScrollContainer(labPanelRef.current);
         });
@@ -1306,11 +1360,18 @@ export function TreatmentWorkspace({
                   treatmentDetailId: detailId,
                   tooth: tp.tooth,
                   prosthesisTypeCode: tp.prosthesisTypeCode,
+                  selectionGroupId: tp.selectionGroupId ?? '',
                 };
               })
               .filter(
-                (row): row is { treatmentDetailId: string; tooth: string; prosthesisTypeCode: string } =>
-                  row !== null,
+                (
+                  row,
+                ): row is {
+                  treatmentDetailId: string;
+                  tooth: string;
+                  prosthesisTypeCode: string;
+                  selectionGroupId: string;
+                } => row !== null,
               ),
             attachmentIds: lc.attachmentIds,
             dueDate: lc.dueDate ?? null,
@@ -1379,35 +1440,42 @@ export function TreatmentWorkspace({
     ],
   );
 
-  const handleRemoveActiveDetail = useCallback(() => {
-    if (!canEditTreatmentForDay) return;
-    const idx = details.findIndex((d) => d.clientId === activeDetailId);
-    if (idx < 0) return;
-    const active = details[idx];
-    if (!active || isDetailLocked(active) || details.length <= 1) return;
-    if (!window.confirm(t('confirmRemoveDetail'))) return;
+  const handleRemoveDetail = useCallback(
+    (detailClientId: string) => {
+      if (!canEditTreatmentForDay) return;
+      const idx = details.findIndex((d) => d.clientId === detailClientId);
+      if (idx < 0) return;
+      const target = details[idx];
+      if (!target || isDetailLocked(target) || details.length <= 1) return;
+      if (!window.confirm(t('confirmRemoveDetail'))) return;
 
-    const removedId = active.clientId;
-    const nextDetails = details.filter((d) => d.clientId !== removedId);
-    const nextActive =
-      nextDetails[Math.min(idx, nextDetails.length - 1)]?.clientId ?? nextDetails[0]?.clientId;
-    setDetails(nextDetails);
-    if (nextActive) setActiveDetailId(nextActive);
+      const nextDetails = details.filter((d) => d.clientId !== detailClientId);
+      const nextActive =
+        activeDetailId === detailClientId
+          ? (nextDetails[Math.min(idx, nextDetails.length - 1)]?.clientId ??
+            nextDetails[0]?.clientId)
+          : activeDetailId;
+      setDetails(nextDetails);
+      if (nextActive) setActiveDetailId(nextActive);
 
-    if (labCaseDrafts.some((lc) => lc.detailClientId === removedId)) {
-      handleLabCasesChange(
-        withoutEmptyLabCaseDrafts(labCaseDrafts.filter((lc) => lc.detailClientId !== removedId)),
-      );
-    }
-  }, [
-    activeDetailId,
-    canEditTreatmentForDay,
-    details,
-    handleLabCasesChange,
-    isDetailLocked,
-    labCaseDrafts,
-    t,
-  ]);
+      if (labCaseDrafts.some((lc) => lc.detailClientId === detailClientId)) {
+        handleLabCasesChange(
+          withoutEmptyLabCaseDrafts(
+            labCaseDrafts.filter((lc) => lc.detailClientId !== detailClientId),
+          ),
+        );
+      }
+    },
+    [
+      activeDetailId,
+      canEditTreatmentForDay,
+      details,
+      handleLabCasesChange,
+      isDetailLocked,
+      labCaseDrafts,
+      t,
+    ],
+  );
 
   const handleAddLabCase = useCallback(async () => {
     if (!canEditTreatmentForDay || !selectedAppointment) return;
@@ -1764,35 +1832,6 @@ export function TreatmentWorkspace({
         </div>
 
         <div className="space-y-3 min-w-0 w-full">
-          <FdiToothChart
-            selected={chartSelectedTeeth}
-            toothColors={chartToothColors}
-            readOnly={showWholeTreatmentPlan}
-            headerControl={
-              details.length > 1 ? (
-                <Checkbox
-                  checked={showWholeTreatmentPlan}
-                  onChange={setShowWholeTreatmentPlan}
-                  label={t('toothChartWholePlan')}
-                  className="text-[11px] [&_span:last-child]:text-[11px] [&_span:last-child]:text-text-muted"
-                />
-              ) : undefined
-            }
-            onToggle={(fdi) => {
-              if (!canEditTreatmentForDay || isDetailLocked(activeDetail) || showWholeTreatmentPlan) return;
-              setDetails((prev) =>
-                prev.map((d) => {
-                  if (d.clientId !== activeDetailId) return d;
-                  const set = new Set(d.teeth);
-                  if (set.has(fdi)) set.delete(fdi);
-                  else set.add(fdi);
-                  return { ...d, teeth: [...set].sort() as FdiToothId[] };
-                }),
-              );
-            }}
-            disabled={!canEditTreatmentForDay || isDetailLocked(activeDetail)}
-          />
-
           <TreatmentDetailsEditor
             details={details}
             activeDetailId={activeDetailId}
@@ -1811,60 +1850,223 @@ export function TreatmentWorkspace({
               );
               setDetails((prev) => [...prev, next]);
               setActiveDetailId(next.clientId);
+              setEntryStep('teeth');
             }}
-            onRemoveDetail={handleRemoveActiveDetail}
+            onRemoveDetail={handleRemoveDetail}
             onUploadFiles={(files) => void uploadForDetail(activeDetailId, files ?? [])}
+            showChrome
+            showFields={false}
           />
 
-          <div ref={labPanelRef}>
-          {showLabShipmentBlocked ? <LabShipmentBlockedNotice /> : null}
-          {showLabDispatchPanel ? (
-          <LabCasesDispatchPanel
-            details={details}
-            activeDetailId={activeDetailId}
-            labCases={labCaseDrafts}
-            labDependentCodes={labDependentCodes}
-            treatmentCatalog={treatmentCatalog}
-            labCaseSummary={activeLabCaseSummary}
-            locale={locale}
-            onLabCaseSummaryChange={handleLabCaseSummaryChange}
-            onLabCaseMarkedRead={handleLabCaseMarkedRead}
-            onLabCaseActivityChange={() => {
-              if (historyPatientId) {
-                void refreshPatientLabCases(historyPatientId, { silent: true });
-              }
-            }}
-            activeLabCaseId={activeLabCaseId}
-            onLabCasesChange={handleLabCasesChange}
-            disabled={!canEditTreatmentForDay}
-            canEdit={canEdit}
-            orgs={orgs}
-            organizationSearch={organizationSearch}
-            onOrganizationSearchChange={setOrganizationSearch}
-            recentOrganizationIds={recentOrganizationIds}
-            onRecentOrganizationPick={(orgId) => {
-              setLabCaseDrafts((prev) => {
-                const targetId =
-                  activeLabCaseId ??
-                  prev.find((lc) => !lc.sentAt && lc.detailClientId === activeDetailId)
-                    ?.clientId;
-                if (!targetId) return prev;
-                return prev.map((lc) =>
-                  lc.clientId === targetId && !lc.sentAt
-                    ? { ...lc, destinationOrganizationId: orgId }
-                    : lc,
-                );
-              });
-            }}
-            sendBusyId={sendBusyId}
-            onAddLabCase={() => void handleAddLabCase()}
-            onSendLabCase={(lc, comment) => void handleSendLabCase(lc, comment)}
-            onCommentError={showError}
-            canInviteLab={canAccessOrganizations}
-            onInviteLab={() => router.push('/organizations?action=invite-lab')}
-          />
-          ) : null}
+          <div className="surface-card px-3 py-2 flex flex-col gap-2 sm:flex-row sm:items-center sm:justify-between sm:gap-3">
+            <WizardStepper
+              className="min-w-0 flex-1"
+              aria-label={t('entryWizardLabel')}
+              steps={visibleEntrySteps.map((step) => ({
+                id: step,
+                label:
+                  step === 'teeth'
+                    ? t('entryStepTeeth')
+                    : step === 'content'
+                      ? t('entryStepContent')
+                      : t('entryStepLab'),
+              }))}
+              currentStepId={entryStep}
+              onStepChange={(stepId) => setEntryStep(stepId as EntryStep)}
+            />
+            <div className="flex flex-wrap gap-2 sm:justify-end shrink-0">
+              <Button
+                type="button"
+                variant="ghost"
+                disabled={entryStep === 'teeth'}
+                onClick={() => {
+                  const idx = visibleEntrySteps.indexOf(entryStep);
+                  if (idx > 0) setEntryStep(visibleEntrySteps[idx - 1]);
+                }}
+              >
+                {t('entryStepBack')}
+              </Button>
+              <Button
+                type="button"
+                variant="primary"
+                disabled={
+                  entryStep === 'lab' || (entryStep === 'content' && !showLabWizardStep)
+                }
+                onClick={() => {
+                  if (entryStep === 'teeth') setEntryStep('content');
+                  else if (entryStep === 'content' && showLabWizardStep) setEntryStep('lab');
+                }}
+              >
+                {t('entryStepNext')}
+              </Button>
+            </div>
           </div>
+
+          {entryStep === 'teeth' ? (
+            <FdiToothChart
+              selected={chartSelectedTeeth}
+              connectedTeeth={showWholeTreatmentPlan ? undefined : connectedSelectedTeeth}
+              toothColors={chartToothColors}
+              readOnly={showWholeTreatmentPlan}
+              headerControl={
+                details.length > 1 ? (
+                  <Checkbox
+                    checked={showWholeTreatmentPlan}
+                    onChange={setShowWholeTreatmentPlan}
+                    label={t('toothChartWholePlan')}
+                    className="text-[11px] [&_span:last-child]:text-[11px] [&_span:last-child]:text-text-muted"
+                  />
+                ) : undefined
+              }
+              onToggle={(fdi, event) => {
+                if (
+                  !canEditTreatmentForDay ||
+                  isDetailLocked(activeDetail) ||
+                  showWholeTreatmentPlan ||
+                  !activeDetailId
+                ) {
+                  return;
+                }
+
+                const detail = details.find((d) => d.clientId === activeDetailId);
+                if (!detail) return;
+                const currentGroups =
+                  detail.toothSelectionGroups.length > 0
+                    ? detail.toothSelectionGroups
+                    : groupsFromFlatTeeth(detail.teeth);
+
+                let nextGroups: ReturnType<typeof toggleToothInGroups> | null = null;
+
+                if (event.shiftKey) {
+                  const anchor = rangeAnchorRef.current;
+                  // Need a prior click as range start; shift alone on one tooth does nothing.
+                  if (!anchor || anchor === fdi) {
+                    rangeAnchorRef.current = fdi;
+                    return;
+                  }
+                  nextGroups = applyShiftRange(currentGroups, anchor, fdi);
+                  rangeAnchorRef.current = fdi;
+                  if (!nextGroups) return;
+                } else {
+                  nextGroups = toggleToothInGroups(currentGroups, fdi);
+                  rangeAnchorRef.current = fdi;
+                }
+
+                setDetails((prev) =>
+                  prev.map((d) =>
+                    d.clientId !== activeDetailId
+                      ? d
+                      : {
+                          ...d,
+                          toothSelectionGroups: nextGroups!,
+                          teeth: deriveTeethFromGroups(nextGroups!),
+                        },
+                  ),
+                );
+
+                setLabCaseDrafts((prev) =>
+                  prev.map((lc) => {
+                    if (lc.detailClientId !== activeDetailId) return lc;
+                    return {
+                      ...lc,
+                      toothProsthesis: pruneToothProsthesisForGroups(
+                        lc.toothProsthesis,
+                        activeDetailId,
+                        nextGroups!,
+                      ),
+                    };
+                  }),
+                );
+              }}
+              disabled={!canEditTreatmentForDay || isDetailLocked(activeDetail)}
+            />
+          ) : null}
+
+          {entryStep === 'content' ? (
+            <TreatmentDetailsEditor
+              details={details}
+              activeDetailId={activeDetailId}
+              onActiveDetailChange={setActiveDetailId}
+              onDetailsChange={setDetails}
+              isDetailLocked={isDetailLocked}
+              labDependentCodes={labDependentCodes}
+              treatmentCatalog={treatmentDropdownCatalog}
+              disabled={!canEditTreatmentForDay}
+              canEdit={canEdit}
+              saveStatus={saveStatus}
+              uploadBusy={uploadBusyDetailId === activeDetailId}
+              onAddDetail={() => {
+                const next = newDetail(
+                  defaultTreatmentTypeForAppointment(
+                    selectedAppointment?.purpose,
+                    treatmentCatalog,
+                  ),
+                );
+                setDetails((prev) => [...prev, next]);
+                setActiveDetailId(next.clientId);
+                setEntryStep('teeth');
+              }}
+              onUploadFiles={(files) => void uploadForDetail(activeDetailId, files ?? [])}
+              showChrome={false}
+              showFields
+            />
+          ) : null}
+
+          {entryStep === 'lab' ? (
+            <div ref={labPanelRef}>
+              {showLabShipmentBlocked ? <LabShipmentBlockedNotice /> : null}
+              {showLabDispatchPanel ? (
+                <LabCasesDispatchPanel
+                  details={details}
+                  activeDetailId={activeDetailId}
+                  labCases={labCaseDrafts}
+                  labDependentCodes={labDependentCodes}
+                  treatmentCatalog={treatmentCatalog}
+                  labCaseSummary={activeLabCaseSummary}
+                  locale={locale}
+                  onLabCaseSummaryChange={handleLabCaseSummaryChange}
+                  onLabCaseMarkedRead={handleLabCaseMarkedRead}
+                  onLabCaseActivityChange={() => {
+                    if (historyPatientId) {
+                      void refreshPatientLabCases(historyPatientId, { silent: true });
+                    }
+                  }}
+                  activeLabCaseId={activeLabCaseId}
+                  onLabCasesChange={handleLabCasesChange}
+                  disabled={!canEditTreatmentForDay}
+                  canEdit={canEdit}
+                  orgs={orgs}
+                  organizationSearch={organizationSearch}
+                  onOrganizationSearchChange={setOrganizationSearch}
+                  recentOrganizationIds={recentOrganizationIds}
+                  onRecentOrganizationPick={(orgId) => {
+                    setLabCaseDrafts((prev) => {
+                      const targetId =
+                        activeLabCaseId ??
+                        prev.find((lc) => !lc.sentAt && lc.detailClientId === activeDetailId)
+                          ?.clientId;
+                      if (!targetId) return prev;
+                      return prev.map((lc) =>
+                        lc.clientId === targetId && !lc.sentAt
+                          ? { ...lc, destinationOrganizationId: orgId }
+                          : lc,
+                      );
+                    });
+                  }}
+                  sendBusyId={sendBusyId}
+                  onAddLabCase={() => void handleAddLabCase()}
+                  onSendLabCase={(lc, comment) => void handleSendLabCase(lc, comment)}
+                  onCommentError={showError}
+                  canInviteLab={canAccessOrganizations}
+                  onInviteLab={() => router.push('/organizations?action=invite-lab')}
+                />
+              ) : (
+                <p className="text-sm text-text-muted surface-card p-4">
+                  {t('entryStepLabUnavailable')}
+                </p>
+              )}
+            </div>
+          ) : null}
         </div>
       </div>
     </div>