feat(warehouses): pick export train when allocating empty return to wagon

Advancing an empty container return to WAGON_ALLOCATED opened no dialog,
so the wagon_allocation_reference column stayed null and nobody knew which
departure carried the empty. Now the action opens a picker of scheduled
EXPORT trains (DRAFT/SCHEDULED) and records the selection.

apps/edr-freight-web/backoffice/src/pages/warehouses/ContainerReturnsPage.tsx

@@ -37,6 +37,7 @@ import type {
   EmptyContainerReturn,
   EmptyContainerReturnStatus,
 } from "@/types/importOperations";
+import type { TrainScheduleListItem } from "@/types/trainScheduling";
 import { formatDateTime } from "@/lib/format";
 
 type ReturnType = "all" | "edr" | "customer";
@@ -100,6 +101,7 @@ export default function ContainerReturnsPage() {
   const [standaloneModalOpen, setStandaloneModalOpen] = useState(false);
   const [activeKey, setActiveKey] = useState<string | null>(null);
   const [historyRow, setHistoryRow] = useState<any | null>(null);
+  const [allocateRow, setAllocateRow] = useState<EmptyContainerReturn | null>(null);
 
   const { data: unloadedQueue = [], isLoading: queueLoading } = useQuery({
     queryKey: ["import-unloaded-queue"],
@@ -306,15 +308,25 @@ export default function ContainerReturnsPage() {
   });
 
   const advanceStatusMutation = useMutation({
-    mutationFn: (id: string) => {
+    mutationFn: ({
+      id,
+      wagonAllocationReference,
+    }: {
+      id: string;
+      wagonAllocationReference?: string;
+    }) => {
       const current = returnedContainers.find((r: any) => r.id === id);
       const nextIndex = RETURN_STATUS_ORDER.indexOf(current?.status ?? "RETURNED") + 1;
       const status = RETURN_STATUS_ORDER[nextIndex] ?? "COMPLETED";
-      return importOperationsService.updateEmptyReturnStatus(id, { status });
+      return importOperationsService.updateEmptyReturnStatus(id, {
+        status,
+        wagonAllocationReference,
+      });
     },
     onSuccess: () => {
       toast({ title: "Return status updated" });
       qc.invalidateQueries({ queryKey: ["empty-container-returns"] });
+      setAllocateRow(null);
     },
     onError: (error: any) => {
       toast({
@@ -381,9 +393,16 @@ export default function ContainerReturnsPage() {
       id: "status",
       header: "Status",
       cell: ({ row }) => (
-        <Badge size="sm">
-          {RETURN_STATUS_LABEL[row.original.status] ?? row.original.status}
-        </Badge>
+        <Stack gap={2}>
+          <Badge size="sm">
+            {RETURN_STATUS_LABEL[row.original.status] ?? row.original.status}
+          </Badge>
+          {row.original.wagonAllocationReference && (
+            <Text size="xs" c="dimmed">
+              Train {row.original.wagonAllocationReference}
+            </Text>
+          )}
+        </Stack>
       ),
     },
     {
@@ -406,8 +425,15 @@ export default function ContainerReturnsPage() {
               <Button
                 size="xs"
                 variant="light"
-                loading={advanceStatusMutation.isPending && advanceStatusMutation.variables === ret.id}
-                onClick={() => advanceStatusMutation.mutate(ret.id)}
+                loading={
+                  advanceStatusMutation.isPending &&
+                  advanceStatusMutation.variables?.id === ret.id
+                }
+                onClick={() =>
+                  nextStatus === "WAGON_ALLOCATED"
+                    ? setAllocateRow(ret)
+                    : advanceStatusMutation.mutate({ id: ret.id })
+                }
               >
                 Advance to {RETURN_STATUS_LABEL[nextStatus]}
               </Button>
@@ -629,6 +655,19 @@ export default function ContainerReturnsPage() {
         loading={createReturnsMutation.isPending}
       />
 
+      <ExportTrainAllocationModal
+        row={allocateRow}
+        onClose={() => setAllocateRow(null)}
+        onSubmit={(reference) =>
+          allocateRow &&
+          advanceStatusMutation.mutate({
+            id: allocateRow.id,
+            wagonAllocationReference: reference,
+          })
+        }
+        loading={advanceStatusMutation.isPending}
+      />
+
       <Modal opened={!!historyRow} onClose={() => setHistoryRow(null)} title="Status History" size="sm">
         {historyRow && (
           <Stack gap="sm">
@@ -649,6 +688,123 @@ export default function ContainerReturnsPage() {
   );
 }
 
+/**
+ * Empties ride an EXPORT departure back to Djibouti, so wagon allocation picks
+ * from the export schedules that have not left yet (DRAFT/SCHEDULED). The pick
+ * is recorded as the return's `wagonAllocationReference`.
+ */
+function ExportTrainAllocationModal({
+  row,
+  onClose,
+  onSubmit,
+  loading,
+}: {
+  row: EmptyContainerReturn | null;
+  onClose: () => void;
+  onSubmit: (reference: string) => void;
+  loading: boolean;
+}) {
+  const [scheduleId, setScheduleId] = useState<string | null>(null);
+
+  const schedulesQuery = useQuery(
+    api.trainScheduling.scheduleList.queryOptions({
+      input: { filters: { pageSize: 100, sortBy: "scheduledDepartureDate", sortOrder: "ASC" } },
+      enabled: Boolean(row),
+    }),
+  );
+
+  const exportTrains = useMemo(
+    () =>
+      ((schedulesQuery.data?.items ?? []) as TrainScheduleListItem[]).filter(
+        (s) => s.direction === "EXPORT" && (s.status === "DRAFT" || s.status === "SCHEDULED"),
+      ),
+    [schedulesQuery.data],
+  );
+
+  const referenceOf = (train: TrainScheduleListItem) =>
+    train.trainNumber || train.reference || train.id;
+
+  return (
+    <Modal opened={!!row} onClose={onClose} title="Allocate to Export Train" size="lg">
+      {row && (
+        <Stack gap="md">
+          <Text fw={600}>{row.containerNumber}</Text>
+
+          {schedulesQuery.isLoading ? (
+            <Group justify="center" py="md">
+              <Loader size="sm" />
+            </Group>
+          ) : exportTrains.length === 0 ? (
+            <Alert color="gray">No export train is scheduled — create one in Train Scheduling.</Alert>
+          ) : (
+            <Table highlightOnHover>
+              <Table.Thead>
+                <Table.Tr>
+                  <Table.Th w={40} />
+                  <Table.Th>Train</Table.Th>
+                  <Table.Th>Departure</Table.Th>
+                  <Table.Th>Route</Table.Th>
+                  <Table.Th>Wagons</Table.Th>
+                  <Table.Th>Status</Table.Th>
+                </Table.Tr>
+              </Table.Thead>
+              <Table.Tbody>
+                {exportTrains.map((train) => (
+                  <Table.Tr
+                    key={train.id}
+                    onClick={() => setScheduleId(train.id)}
+                    style={{ cursor: "pointer" }}
+                  >
+                    <Table.Td>
+                      <Checkbox
+                        checked={scheduleId === train.id}
+                        onChange={() => setScheduleId(train.id)}
+                      />
+                    </Table.Td>
+                    <Table.Td>
+                      <Text fw={600} size="sm">
+                        {referenceOf(train)}
+                      </Text>
+                    </Table.Td>
+                    <Table.Td>
+                      {train.scheduleDate ? new Date(train.scheduleDate).toLocaleDateString() : "—"}
+                    </Table.Td>
+                    <Table.Td>
+                      {train.origin ?? "—"} → {train.destination ?? "—"}
+                    </Table.Td>
+                    <Table.Td>
+                      {train.wagonsUsed ?? 0}/{train.wagonCount}
+                    </Table.Td>
+                    <Table.Td>
+                      <Badge size="sm">{train.status}</Badge>
+                    </Table.Td>
+                  </Table.Tr>
+                ))}
+              </Table.Tbody>
+            </Table>
+          )}
+
+          <Group justify="flex-end" gap="sm">
+            <Button variant="default" onClick={onClose} disabled={loading}>
+              Cancel
+            </Button>
+            <Button
+              disabled={!scheduleId}
+              loading={loading}
+              onClick={() => {
+                const train = exportTrains.find((t) => t.id === scheduleId);
+                if (train) onSubmit(referenceOf(train));
+              }}
+            >
+              Allocate
+            </Button>
+          </Group>
+        </Stack>
+      )}
+    </Modal>
+  );
+}
+
 interface ContainerReturnModalProps {
   opened: boolean;
   onClose: () => void;