Update beamtime assignment logic for pucks and samples

Simplified and unified beamtime assignment handling for pucks and samples in the backend. Enhanced the frontend to display detailed assignment state, including shift, date, and beamline, for both pucks and dewars. This ensures consistent and accurate state management across the application.

backend/app/routers/dewar.py

@@ -595,14 +595,23 @@ async def assign_beamtime_to_dewar(
     db: Session = Depends(get_db),
 ):
     dewar = db.query(DewarModel).filter(DewarModel.id == dewar_id).first()
-    if not dewar:  # <- Move check earlier!
+    if not dewar:
         raise HTTPException(status_code=404, detail="Dewar not found")
-    if beamtime_id == 0:
+
-        dewar.beamtime_id = None
+    dewar.beamtime_id = None if beamtime_id == 0 else beamtime_id
-    else:
-        dewar.beamtime_id = beamtime_id
     db.commit()
     db.refresh(dewar)
+    for puck in dewar.pucks:
+        puck.beamtime_id = None if beamtime_id == 0 else beamtime_id
+        for sample in puck.samples:
+            has_sample_event = (
+                db.query(SampleEvent).filter(SampleEvent.sample_id == sample.id).count()
+                > 0
+            )
+            if not has_sample_event:
+                sample.beamtime_id = None if beamtime_id == 0 else beamtime_id
+
+    db.commit()
     return {"status": "success", "dewar_id": dewar.id, "beamtime_id": beamtime_id}
 
 

backend/app/routers/puck.py

@@ -20,6 +20,7 @@ from app.models import (
     Sample as SampleModel,
     LogisticsEvent as LogisticsEventModel,
     Dewar as DewarModel,
+    SampleEvent,
 )
 from app.dependencies import get_db
 import logging
@@ -669,10 +670,16 @@ async def assign_beamtime_to_puck(
     puck = db.query(PuckModel).filter(PuckModel.id == puck_id).first()
     if not puck:
         raise HTTPException(status_code=404, detail="Puck not found")
-    if beamtime_id == 0:
+
-        puck.beamtime_id = None
+    puck.beamtime_id = None if beamtime_id == 0 else beamtime_id
-    else:
-        puck.beamtime_id = beamtime_id
     db.commit()
     db.refresh(puck)
+    # Update samples
+    for sample in puck.samples:
+        has_sample_event = (
+            db.query(SampleEvent).filter(SampleEvent.sample_id == sample.id).count() > 0
+        )
+        if not has_sample_event:
+            sample.beamtime_id = None if beamtime_id == 0 else beamtime_id
+    db.commit()
     return {"status": "success", "puck_id": puck.id, "beamtime_id": beamtime_id}

frontend/src/components/Calendar.tsx

@@ -268,12 +268,11 @@ const Calendar: React.FC = () => {
             return;
         }
 
-        // Find out if it's already assigned to this shift!
+        // Get current association state
         const prev = eventAssociations[selectedEventId] || { dewars: [], pucks: [] };
         const isAssigned = prev.pucks.includes(puckId);
 
         if (!isAssigned) {
-            // Assign it immediately
             assignPuckToBeamtime(Number(puckId), Number(beamtimeId))
                 .then(() => {
                     setEventAssociations(prevAssoc => {
@@ -291,7 +290,6 @@ const Calendar: React.FC = () => {
                     console.error("Failed to assign puck to beamtime", e);
                 });
         } else {
-            // Unassign (patch to None) immediately
             unassignPuckFromBeamtime(Number(puckId))
                 .then(() => {
                     setEventAssociations(prevAssoc => {
@@ -387,10 +385,8 @@ const Calendar: React.FC = () => {
                     <h4>Select Dewars</h4>
                     <ul>
                         {shipments.map(dewar => {
-                            // Are *all* pucks assigned to this event? (Assigned at Dewar level)
+                            const thisEvent = eventAssociations[selectedEventId] || { dewars: [], pucks: [] };
-                            const thisEvent = eventAssociations[selectedEventId] || {dewars: [], pucks: []};
                             const dewarAssigned = thisEvent.dewars.includes(dewar.id);
-                            const puckAssignments = thisEvent.pucks;
 
                             return (
                                 <li key={dewar.id}>
@@ -405,8 +401,20 @@ const Calendar: React.FC = () => {
                                     {!dewarAssigned && dewar.pucks && dewar.pucks.length > 0 && (
                                         <ul>
                                             {(dewar.pucks || []).map(puck => {
-                                                const isAssigned =
+                                                // Find eventId for this puck, if assigned
-                                                    !!eventAssociations[selectedEventId]?.pucks.includes(puck.id);
+                                                const associatedEventId = Object.keys(eventAssociations).find(eid =>
+                                                    eventAssociations[eid]?.pucks.includes(puck.id)
+                                                );
+                                                const associatedEvent = associatedEventId
+                                                    ? events.find(ev => ev.id === associatedEventId)
+                                                    : null;
+
+                                                const associatedShift = associatedEvent?.beamtime_shift;
+                                                const associatedDate = associatedEvent?.start;
+                                                const associatedBeamline = associatedEvent?.beamline;
+
+                                                const currentShift = eventDetails?.beamtime_shift;
+                                                const isAssignedToCurrentShift = associatedShift && currentShift && associatedShift === currentShift;
 
                                                 return (
                                                     <div key={puck.id} style={{ display: 'flex', alignItems: 'center' }}>
@@ -415,32 +423,67 @@ const Calendar: React.FC = () => {
                                                             type="button"
                                                             style={{
                                                                 marginLeft: 8,
-                                                                background: isAssigned ? '#4CAF50' : '#e0e0e0',
+                                                                background: isAssignedToCurrentShift ? '#4CAF50' : (associatedShift ? '#B3E5B3' : '#e0e0e0'),
-                                                                color: isAssigned ? 'white' : 'black',
+                                                                color: isAssignedToCurrentShift ? 'white' : 'black',
-                                                                border: isAssigned ? '1px solid #388e3c' : '1px solid #bdbdbd',
+                                                                border: isAssignedToCurrentShift ? '1px solid #388e3c' : '1px solid #bdbdbd',
                                                                 borderRadius: 4,
                                                                 padding: '4px 10px',
                                                                 cursor: 'pointer',
                                                                 transition: 'background 0.2s',
                                                             }}
-                                                            onClick={() => handlePuckAssignment(dewar.id, puck.id)}
+                                                            onClick={() => handlePuckAssignment(puck.id)}
                                                         >
-                                                            {isAssigned ? puck.puck_name : puck.puck_name}
+                                                            {puck.puck_name}
                                                         </button>
-                                                        {/* Optionally show assignment info, as before */}
+                                                        {associatedEvent && (
-                                                        {findAssociatedEventForPuck(puck.id) && (
+                                                            <span
-                                                            <span style={{marginLeft: 8, color: 'green'}}>← Assigned to shift: {findAssociatedEventForPuck(puck.id)}</span>
+                                                                style={{
+                                                                    marginLeft: 8,
+                                                                    color: isAssignedToCurrentShift ? 'green' : '#388e3c',
+                                                                    fontWeight: isAssignedToCurrentShift ? 700 : 400
+                                                                }}
+                                                            >
+                      ← Assigned to shift: {associatedShift}
+                                                                {associatedDate && (
+                                                                    <> on {new Date(associatedDate).toLocaleDateString()}</>
+                                                                )}
+                                                                {associatedBeamline && (
+                                                                    <> ({associatedBeamline})</>
+                                                                )}
+                    </span>
                                                         )}
                                                     </div>
                                                 );
                                             })}
-
                                         </ul>
                                     )}
-                                    {/* Show associated shift, if any */}
+                                    {/* Show associated shift, date, and beamline for Dewar, if any */}
-                                    {findAssociatedEventForDewar(dewar.id) && (
+                                    {(() => {
-                                        <span style={{marginLeft: 8, color: 'green'}}>← Assigned to shift: {findAssociatedEventForDewar(dewar.id)}</span>
+                                        const associatedEventId = Object.keys(eventAssociations).find(eid =>
-                                    )}
+                                            eventAssociations[eid]?.dewars.includes(dewar.id)
+                                        );
+                                        const associatedEvent = associatedEventId
+                                            ? events.find(ev => ev.id === associatedEventId)
+                                            : null;
+                                        const associatedShift = associatedEvent?.beamtime_shift;
+                                        const associatedDate = associatedEvent?.start;
+                                        const associatedBeamline = associatedEvent?.beamline;
+
+                                        const currentShift = eventDetails?.beamtime_shift;
+                                        const isAssignedToCurrentShift = associatedShift && currentShift && associatedShift === currentShift;
+
+                                        return associatedEvent && (
+                                            <span style={{ marginLeft: 8, color: isAssignedToCurrentShift ? 'green' : '#388e3c', fontWeight: isAssignedToCurrentShift ? 700 : 400 }}>
+              ← Assigned to shift: {associatedShift}
+                                                {associatedDate && (
+                                                    <> on {new Date(associatedDate).toLocaleDateString()}</>
+                                                )}
+                                                {associatedBeamline && (
+                                                    <> ({associatedBeamline})</>
+                                                )}
+            </span>
+                                        );
+                                    })()}
                                 </li>
                             );
                         })}