edr-platform/apps/edr-freight-web/backoffice/src/pages/trainScheduling/ScheduleCrewPage.tsx

import { useEffect, useMemo, useState } from "react";
import { useParams } from "react-router-dom";
import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query";
import {
  ActionIcon,
  Alert,
  Badge,
  Button,
  Card,
  Group,
  Loader,
  Select,
  Stack,
  Stepper,
  Text,
  ThemeIcon,
} from "@mantine/core";
import {
  AlertTriangle,
  CheckCircle2,
  Plus,
  ShieldCheck,
  Train,
  Trash2,
  Users,
  Wrench,
} from "lucide-react";

import { PageContainer, PageHeader } from "@/components/page";
import { useToast } from "@/hooks/use-toast";
import { useAuth } from "@/auth/useAuth";
import { FREIGHT_PERMS, hasPermission } from "@/lib/permissions";
import {
  trainCrewService,
  trainCrewRoleLabel,
  type TrainCrewMember,
  type TrainCrewRole,
} from "@/services/trainCrew.service";
import {
  DUTY_ROLE_OPTIONS,
  trainCrewAssignmentService,
  type CorridorYard,
  type CrewDutyRole,
} from "@/services/trainCrewAssignment.service";

/**
 * One driver row being built. The leg (two yards) and the duty role are
 * properties of THIS run, not of the person.
 */
interface DriverRow {
  key: string;
  crewMemberId: string | null;
  fromYardId: string | null;
  toYardId: string | null;
  dutyRole: CrewDutyRole | null;
}

/**
 * A new row pre-filled with the schedule's own endpoints — the common case is
 * one driver over the whole route, and staff narrow it from there.
 */
const newDriverRow = (yards: CorridorYard[]): DriverRow => ({
  key: `driver-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`,
  crewMemberId: null,
  fromYardId: yards[0]?.id ?? null,
  toYardId: yards[yards.length - 1]?.id ?? null,
  dutyRole: null,
});

/**
 * Assign a train crew to one schedule — ITLMS Rolling Stock §1.1 and §1.2.
 *
 * Crew sizes are free-form: operations add as many drivers, police, technicians
 * or specialists as a given run needs, rather than filling the fixed pairing
 * cases of §2. What is still enforced is what makes a run coherent — every
 * driver carries a leg and duty role, one Primary per leg, Djibouti drivers
 * confined to Dire Dawa and eastward (§1.1), and the specialized crew the cargo
 * actually demands (§1.2).
 *
 * A partial crew always saves: §1.2 puts the hard gate at departure, so this
 * page and the dispatch guard call the same server-side validator.
 */
export default function ScheduleCrewPage() {
  const { scheduleId = "" } = useParams();
  const { toast } = useToast();
  const qc = useQueryClient();
  const { user } = useAuth();
  const canAssign = hasPermission(user, FREIGHT_PERMS.trainCrew.assign);

  const [step, setStep] = useState(0);
  const [drivers, setDrivers] = useState<DriverRow[]>([]);
  /** Support and specialist picks, keyed by role. */
  const [supportIds, setSupportIds] = useState<Record<string, Array<string | null>>>({});

  const { data: crew, isLoading } = useQuery({
    queryKey: ["schedule-crew", scheduleId],
    queryFn: async () => (await trainCrewAssignmentService.get(scheduleId)).data,
    enabled: Boolean(scheduleId),
  });

  const { data: roster = [] } = useQuery({
    queryKey: ["train-crew", "roster-all"],
    queryFn: async () => {
      const res = await trainCrewService.getAll({ limit: 200, status: "ACTIVE" });
      return res.data.data;
    },
  });

  // Seed from what is already saved, so reopening resumes rather than restarts.
  useEffect(() => {
    if (!crew) return;
    const driverRows: DriverRow[] = [];
    const support: Record<string, Array<string | null>> = {};
    for (const a of crew.assignments) {
      if (a.role === "TRAIN_DRIVER") {
        driverRows.push({
          key: a.id,
          crewMemberId: a.crewMemberId,
          fromYardId: a.fromYardId ?? null,
          toYardId: a.toYardId ?? null,
          dutyRole: a.dutyRole ?? null,
        });
      } else {
        support[a.role] = [...(support[a.role] ?? []), a.crewMemberId];
      }
    }
    setDrivers(driverRows);
    setSupportIds(support);
  }, [crew]);

  const corridorYards = crew?.corridorYards ?? [];

  const yardOptions = useMemo(
    () => corridorYards.map((y) => ({ value: y.id, label: y.label })),
    [corridorYards],
  );

  /**
   * §1.1 — a leg is open to a Djibouti driver only when both ends sit at or
   * beyond Dire Dawa. Position along the corridor answers this without naming
   * station pairs, so a handover anywhere east of Dire Dawa works.
   */
  const legOpenToDjibouti = (leg: {
    fromYardId: string | null;
    toYardId: string | null;
  }) => {
    const boundary = corridorYards.find((y) => /dire dawa/i.test(y.label));
    const from = corridorYards.find((y) => y.id === leg.fromYardId);
    const to = corridorYards.find((y) => y.id === leg.toYardId);
    // An unknown boundary or half-built leg is not a breach — the server-side
    // validator reports the incomplete leg on its own.
    if (!boundary || !from || !to) return true;
    return Math.min(from.displayOrder, to.displayOrder) >= boundary.displayOrder;
  };

  const byRole = useMemo(() => {
    const map = new Map<TrainCrewRole, TrainCrewMember[]>();
    for (const m of roster) {
      map.set(m.role, [...(map.get(m.role) ?? []), m]);
    }
    return map;
  }, [roster]);

  /** Everyone already picked — nobody may hold two seats on one run. */
  const takenIds = useMemo(() => {
    const ids = [
      ...drivers.map((d) => d.crewMemberId),
      ...Object.values(supportIds).flat(),
    ].filter(Boolean) as string[];
    return new Set(ids);
  }, [drivers, supportIds]);

  const memberOptions = (
    role: TrainCrewRole,
    currentValue: string | null,
    leg?: { fromYardId: string | null; toYardId: string | null },
  ) =>
    (byRole.get(role) ?? [])
      .filter((m) => {
        // §1.1 territorial boundary: a Djibouti driver never appears on a leg
        // they may not work. Enforced by making the invalid choice unavailable
        // rather than by rejecting it afterwards.
        if (leg && m.nationality === "DJIBOUTIAN" && !legOpenToDjibouti(leg)) {
          return false;
        }
        return m.id === currentValue || !takenIds.has(m.id);
      })
      .map((m) => ({
        value: m.id,
        label: `${m.firstName} ${m.lastName} · ${m.nationality === "ETHIOPIAN" ? "ET" : "DJ"}`,
      }));

  const setDriver = (key: string, patch: Partial<DriverRow>) =>
    setDrivers((prev) =>
      prev.map((row) => {
        if (row.key !== key) return row;
        const next = { ...row, ...patch };
        // Moving the leg can invalidate the person already chosen — clear
        // rather than silently persist a territorial breach.
        const legMoved = patch.fromYardId !== undefined || patch.toYardId !== undefined;
        if (legMoved && next.crewMemberId) {
          const member = roster.find((m) => m.id === next.crewMemberId);
          if (member?.nationality === "DJIBOUTIAN" && !legOpenToDjibouti(next)) {
            next.crewMemberId = null;
          }
        }
        return next;
      }),
    );

  const setSupportCount = (role: TrainCrewRole, count: number) =>
    setSupportIds((prev) => ({
      ...prev,
      [role]: Array.from({ length: count }, (_, i) => prev[role]?.[i] ?? null),
    }));

  const buildPayload = () => {
    const assignments: Array<{
      crewMemberId: string;
      role: TrainCrewRole;
      dutyRole?: CrewDutyRole;
      fromYardId?: string;
      toYardId?: string;
    }> = [];
    for (const row of drivers) {
      if (row.crewMemberId) {
        assignments.push({
          crewMemberId: row.crewMemberId,
          role: "TRAIN_DRIVER",
          ...(row.dutyRole ? { dutyRole: row.dutyRole } : {}),
          ...(row.fromYardId ? { fromYardId: row.fromYardId } : {}),
          ...(row.toYardId ? { toYardId: row.toYardId } : {}),
        });
      }
    }
    for (const [role, ids] of Object.entries(supportIds)) {
      for (const id of ids) {
        if (id) assignments.push({ crewMemberId: id, role: role as TrainCrewRole });
      }
    }
    return { assignments };
  };

  const saveMutation = useMutation({
    mutationFn: () => trainCrewAssignmentService.save(scheduleId, buildPayload()),
    onSuccess: (res) => {
      const validation = res.data;
      toast({
        title: validation.complete
          ? "Crew saved — composition complete"
          : "Crew saved (still incomplete)",
        description: validation.complete
          ? undefined
          : "The train cannot be dispatched until every rule passes.",
      });
      qc.invalidateQueries({ queryKey: ["schedule-crew", scheduleId] });
    },
    onError: (error: unknown) => {
      const message = (error as { response?: { data?: { message?: unknown } } })
        ?.response?.data?.message;
      toast({
        title: "Could not save crew",
        description: Array.isArray(message)
          ? message.join(", ")
          : typeof message === "string"
            ? message
            : "The request failed. Please try again.",
        variant: "destructive",
      });
    },
  });

  if (isLoading) {
    return (
      <PageContainer>
        <Group justify="center" py="xl">
          <Loader size="sm" />
        </Group>
      </PageContainer>
    );
  }

  const demand = crew?.demand;
  const specialized = crew?.requirements.specialized ?? [];
  const technicianRule = crew?.requirements.technician;
  const validation = crew?.validation;

  return (
    <PageContainer>
      <PageHeader
        title="Assign Train Crew"
        subtitle="Add as many drivers and crew as this run needs"
        backTo={`/dashboard/operations/train-scheduling-v2/${scheduleId}`}
        meta={
          validation ? (
            <Badge
              variant="light"
              color={validation.complete ? "green" : "orange"}
              leftSection={
                validation.complete ? <CheckCircle2 size={12} /> : <AlertTriangle size={12} />
              }
            >
              {validation.complete ? "Ready to dispatch" : "Incomplete"}
            </Badge>
          ) : null
        }
        action={
          canAssign ? (
            <Button
              onClick={() => saveMutation.mutate()}
              loading={saveMutation.isPending}
              color="edr-green"
            >
              Save Crew
            </Button>
          ) : null
        }
      />

      <Stepper active={step} onStepClick={setStep} mt="md" size="sm">
        <Stepper.Step label="Drivers" description="Any number">
          <Stack gap="md" mt="lg">
            <Text size="sm" c="dimmed">
              Add a row per driver and set the leg they work — any two yards on this
              schedule's route, so a handover at Feto or Meiso is as easy as one at
              Dire Dawa. Djibouti drivers are offered only on legs from Dire Dawa
              eastward.
            </Text>

            {drivers.length === 0 ? (
              <Alert color="gray">No drivers added yet.</Alert>
            ) : (
              drivers.map((row, index) => (
                <Card key={row.key} withBorder padding="md">
                  <Group justify="space-between" mb="sm">
                    <Group gap="sm">
                      <ThemeIcon size={28} radius="md" variant="light" color="edr-green">
                        <Train size={14} />
                      </ThemeIcon>
                      <Text fw={600} size="sm">
                        Driver {index + 1}
                      </Text>
                    </Group>
                    <ActionIcon
                      variant="subtle"
                      color="red"
                      aria-label="Remove driver"
                      onClick={() =>
                        setDrivers((prev) => prev.filter((d) => d.key !== row.key))
                      }
                    >
                      <Trash2 size={16} />
                    </ActionIcon>
                  </Group>
                  <Group grow align="flex-start" wrap="wrap">
                    <Select
                      label="From yard"
                      placeholder="Start of this leg"
                      searchable
                      data={yardOptions}
                      value={row.fromYardId}
                      onChange={(val) => setDriver(row.key, { fromYardId: val })}
                    />
                    <Select
                      label="To yard"
                      placeholder="End of this leg"
                      searchable
                      data={yardOptions}
                      value={row.toYardId}
                      onChange={(val) => setDriver(row.key, { toYardId: val })}
                    />
                    <Select
                      label="Duty role"
                      placeholder="Select a duty role"
                      data={DUTY_ROLE_OPTIONS}
                      value={row.dutyRole}
                      onChange={(val) =>
                        setDriver(row.key, { dutyRole: (val as CrewDutyRole) ?? null })
                      }
                    />
                    <Select
                      label="Driver"
                      placeholder="Select a driver"
                      searchable
                      clearable
                      data={memberOptions("TRAIN_DRIVER", row.crewMemberId, row)}
                      value={row.crewMemberId}
                      onChange={(val) => setDriver(row.key, { crewMemberId: val })}
                    />
                  </Group>
                </Card>
              ))
            )}

            <Button
              variant="light"
              leftSection={<Plus size={16} />}
              onClick={() => setDrivers((prev) => [...prev, newDriverRow(corridorYards)])}
            >
              Add Driver
            </Button>
          </Stack>
        </Stepper.Step>

        <Stepper.Step label="Support crew" description="Police, technical, cargo">
          <Stack gap="lg" mt="lg">
            <SupportSection
              role="FEDERAL_POLICE"
              icon={<ShieldCheck size={16} />}
              color="blue"
              title="Security detail"
              hint="Add as many federal police as this run needs"
              values={supportIds.FEDERAL_POLICE ?? []}
              onCount={(n) => setSupportCount("FEDERAL_POLICE", n)}
              onPick={(i, val) =>
                setSupportIds((prev) => ({
                  ...prev,
                  FEDERAL_POLICE: (prev.FEDERAL_POLICE ?? []).map((v, idx) =>
                    idx === i ? val : v,
                  ),
                }))
              }
              options={(value) => memberOptions("FEDERAL_POLICE", value)}
            />

            <SupportSection
              role="TECHNICIAN"
              icon={<Wrench size={16} />}
              color="orange"
              title="Technical maintenance crew"
              hint={technicianRule?.reason ?? "Optional technical maintenance crew"}
              alert={
                demand?.hasBadOrderWagon
                  ? "A defective wagon is attached, so at least one technician is mandatory."
                  : undefined
              }
              values={supportIds.TECHNICIAN ?? []}
              onCount={(n) => setSupportCount("TECHNICIAN", n)}
              onPick={(i, val) =>
                setSupportIds((prev) => ({
                  ...prev,
                  TECHNICIAN: (prev.TECHNICIAN ?? []).map((v, idx) => (idx === i ? val : v)),
                }))
              }
              options={(value) => memberOptions("TECHNICIAN", value)}
            />

            {specialized.length ? (
              specialized.map((rule) => (
                <SupportSection
                  key={rule.role}
                  role={rule.role}
                  icon={<Users size={16} />}
                  color="grape"
                  title={trainCrewRoleLabel(rule.role)}
                  hint={rule.reason}
                  values={supportIds[rule.role] ?? []}
                  onCount={(n) => setSupportCount(rule.role, n)}
                  onPick={(i, val) =>
                    setSupportIds((prev) => ({
                      ...prev,
                      [rule.role]: (prev[rule.role] ?? []).map((v, idx) =>
                        idx === i ? val : v,
                      ),
                    }))
                  }
                  options={(value) => memberOptions(rule.role, value)}
                />
              ))
            ) : (
              <Alert color="gray">
                No specialized cargo detected on this train — no reefer, HAZMAT, break-bulk
                or livestock crew is required.
              </Alert>
            )}
          </Stack>
        </Stepper.Step>

        <Stepper.Completed>
          <Stack gap="md" mt="lg">
            <Card withBorder padding="lg">
              <Text fw={600} mb="sm">
                Composition checklist
              </Text>
              {validation?.complete ? (
                <Group gap="xs">
                  <ThemeIcon size={22} radius="xl" color="green" variant="light">
                    <CheckCircle2 size={14} />
                  </ThemeIcon>
                  <Text size="sm">Every rule passes — this train may be dispatched.</Text>
                </Group>
              ) : (
                <Stack gap="xs">
                  {validation?.issues.map((issue) => (
                    <Group key={`${issue.code}-${issue.message}`} gap="xs" wrap="nowrap">
                      <ThemeIcon size={22} radius="xl" color="orange" variant="light">
                        <AlertTriangle size={14} />
                      </ThemeIcon>
                      <Text size="sm">{issue.message}</Text>
                    </Group>
                  ))}
                </Stack>
              )}
            </Card>
            {validation?.runType ? (
              <Text size="sm" c="dimmed">
                Derived run type:{" "}
                <Text span fw={600}>
                  {validation.runType === "LONG_RUN" ? "Long run" : "Short run"}
                </Text>
              </Text>
            ) : null}
          </Stack>
        </Stepper.Completed>
      </Stepper>

      <Group justify="space-between" mt="xl">
        <Button variant="light" disabled={step === 0} onClick={() => setStep((s) => s - 1)}>
          Back
        </Button>
        <Button variant="light" disabled={step > 1} onClick={() => setStep((s) => s + 1)}>
          Next
        </Button>
      </Group>
    </PageContainer>
  );
}

/** A crew block: add/remove rows freely, each naming one person. */
function SupportSection({
  role,
  icon,
  color,
  title,
  hint,
  alert,
  values,
  onCount,
  onPick,
  options,
}: {
  role: TrainCrewRole;
  icon: React.ReactNode;
  color: string;
  title: string;
  hint: string;
  alert?: string;
  values: Array<string | null>;
  onCount: (count: number) => void;
  onPick: (index: number, value: string | null) => void;
  options: (currentValue: string | null) => Array<{ value: string; label: string }>;
}) {
  return (
    <Card withBorder padding="lg">
      <Group gap="sm" mb="md">
        <ThemeIcon size={32} radius="md" variant="light" color={color}>
          {icon}
        </ThemeIcon>
        <div>
          <Text fw={600}>{title}</Text>
          <Text size="xs" c="dimmed">
            {hint}
          </Text>
        </div>
      </Group>

      {alert ? (
        <Alert color="orange" icon={<AlertTriangle size={16} />} mb="md">
          {alert}
        </Alert>
      ) : null}

      <Stack gap="sm">
        {values.map((value, index) => (
          <Group key={index} align="flex-end" wrap="nowrap">
            <Select
              label={`${trainCrewRoleLabel(role)} ${index + 1}`}
              placeholder="Select a crew member"
              searchable
              clearable
              data={options(value)}
              value={value}
              onChange={(val) => onPick(index, val)}
              style={{ flex: 1 }}
            />
            <ActionIcon
              variant="subtle"
              color="red"
              aria-label="Remove"
              onClick={() => {
                const next = values.filter((_, i) => i !== index);
                onCount(next.length);
                next.forEach((v, i) => onPick(i, v));
              }}
            >
              <Trash2 size={16} />
            </ActionIcon>
          </Group>
        ))}
        <Button
          variant="light"
          size="xs"
          leftSection={<Plus size={14} />}
          onClick={() => onCount(values.length + 1)}
          style={{ alignSelf: "flex-start" }}
        >
          Add {trainCrewRoleLabel(role)}
        </Button>
      </Stack>
    </Card>
  );
}