tsms/src/components/CreateProject/MilestoneSection.tsx

import { CreateProjectInputs, PaymentInputs } from "@/app/api/projects/actions";
import { TaskGroup } from "@/app/api/tasks";
import { Add, Check, Close, Delete } from "@mui/icons-material";
import {
  Stack,
  Typography,
  Grid,
  FormControl,
  Box,
  Button,
} from "@mui/material";
import {
  GridColDef,
  GridActionsCellItem,
  GridToolbarContainer,
  GridRowModesModel,
  GridRowModes,
  FooterPropsOverrides,
  GridRowId,
  GridEventListener,
  useGridApiRef,
  GridRowModel,
} from "@mui/x-data-grid";
import { LocalizationProvider, DatePicker } from "@mui/x-date-pickers";
import { AdapterDayjs } from "@mui/x-date-pickers/AdapterDayjs";
import dayjs from "dayjs";
import "dayjs/locale/zh-hk";
import React, { useCallback, useEffect, useMemo, useState } from "react";
import { useFormContext } from "react-hook-form";
import { useTranslation } from "react-i18next";
import StyledDataGrid from "../StyledDataGrid";
import { moneyFormatter } from "@/app/utils/formatUtil";
import isDate from "lodash/isDate";

interface Props {
  taskGroupId: TaskGroup["id"];
}

declare module "@mui/x-data-grid" {
  interface FooterPropsOverrides {
    onAdd: () => void;
  }
}

type PaymentRow = Partial<PaymentInputs & { _isNew: boolean; _error: string }>;

const MilestoneSection: React.FC<Props> = ({ taskGroupId }) => {
  const { t } = useTranslation();
  const { getValues, setValue } = useFormContext<CreateProjectInputs>();
  const [payments, setPayments] = useState<PaymentRow[]>(
    getValues("milestones")[taskGroupId]?.payments || [],
  );
  const [rowModesModel, setRowModesModel] = useState<GridRowModesModel>({});
  // Change the payments table depending on the TaskGroupId
  useEffect(() => {
    setPayments(getValues("milestones")[taskGroupId]?.payments || []);
    setRowModesModel({});
  }, [getValues, taskGroupId]);

  const apiRef = useGridApiRef();
  const addRow = useCallback(() => {
    const id = Date.now();
    setPayments((p) => [...p, { id, _isNew: true }]);
    setRowModesModel((model) => ({
      ...model,
      [id]: { mode: GridRowModes.Edit, fieldToFocus: "description" },
    }));
  }, []);

  const validateRow = useCallback(
    (id: GridRowId) => {
      const row = apiRef.current.getRowWithUpdatedValues(id, "") as PaymentRow;
      let error: keyof PaymentInputs | "" = "";
      if (!row.description) {
        error = "description";
      } else if (!(isDate(row.date) && row.date.getTime())) {
        error = "date";
      } else if (!(row.amount && row.amount >= 0)) {
        error = "amount";
      }

      apiRef.current.updateRows([{ id, _error: error }]);
      return !error;
    },
    [apiRef],
  );

  const handleCancel = useCallback(
    (id: GridRowId) => () => {
      setRowModesModel((model) => ({
        ...model,
        [id]: { mode: GridRowModes.View, ignoreModifications: true },
      }));
      const editedRow = payments.find((payment) => payment.id === id);
      if (editedRow?._isNew) {
        setPayments((ps) => ps.filter((p) => p.id !== id));
      }
    },
    [payments],
  );

  const handleDelete = useCallback(
    (id: GridRowId) => () => {
      setPayments((ps) => ps.filter((p) => p.id !== id));
    },
    [],
  );

  const handleSave = useCallback(
    (id: GridRowId) => () => {
      if (validateRow(id)) {
        setRowModesModel((model) => ({
          ...model,
          [id]: { mode: GridRowModes.View },
        }));
      }
    },
    [validateRow],
  );

  const handleEditStop = useCallback<GridEventListener<"rowEditStop">>(
    (params, event) => {
      if (!validateRow(params.id)) {
        event.defaultMuiPrevented = true;
      }
    },
    [validateRow],
  );

  const processRowUpdate = useCallback((newRow: GridRowModel) => {
    const updatedRow = { ...newRow, _isNew: false };
    setPayments((ps) => ps.map((p) => (p.id === newRow.id ? updatedRow : p)));
    return updatedRow;
  }, []);

  const columns = useMemo<GridColDef[]>(
    () => [
      {
        type: "actions",
        field: "actions",
        headerName: t("Actions"),
        getActions: ({ id }) => {
          if (rowModesModel[id]?.mode === GridRowModes.Edit) {
            return [
              <GridActionsCellItem
                key="accpet-action"
                icon={<Check />}
                label={t("Save")}
                onClick={handleSave(id)}
              />,
              <GridActionsCellItem
                key="cancel-action"
                icon={<Close />}
                label={t("Cancel")}
                onClick={handleCancel(id)}
              />,
            ];
          }

          return [
            <GridActionsCellItem
              key="delete-action"
              icon={<Delete />}
              label={t("Remove")}
              onClick={handleDelete(id)}
            />,
          ];
        },
      },
      {
        field: "description",
        headerName: t("Payment Milestone Description"),
        width: 300,
        editable: true,
      },
      {
        field: "date",
        headerName: t("Payment Milestone Date"),
        width: 200,
        type: "date",
        editable: true,
        valueGetter(params) {
          return new Date(params.value);
        },
      },
      {
        field: "amount",
        headerName: t("Payment Milestone Amount"),
        width: 300,
        editable: true,
        type: "number",
        valueFormatter(params) {
          return moneyFormatter.format(params.value);
        },
      },
    ],
    [handleCancel, handleDelete, handleSave, rowModesModel, t],
  );

  useEffect(() => {
    const milestones = getValues("milestones");
    setValue("milestones", {
      ...milestones,
      [taskGroupId]: {
        ...milestones[taskGroupId],
        payments: payments
          .filter((p) => !p._isNew && !p._error)
          .map((p) => ({
            description: p.description!,
            id: p.id!,
            amount: p.amount!,
            date: dayjs(p.date!).toISOString(),
          })),
      },
    });
  }, [getValues, payments, setValue, taskGroupId]);

  return (
    <Stack gap={1}>
      <Typography variant="overline" display="block" marginBlockEnd={1}>
        {t("Task Stage Milestones")}
      </Typography>
      <LocalizationProvider dateAdapter={AdapterDayjs} adapterLocale="zh-hk">
        <Grid container spacing={2} columns={{ xs: 6, sm: 12 }}>
          <Grid item xs>
            <FormControl fullWidth>
              <DatePicker
                label={t("Stage Start Date")}
                value={dayjs(getValues("milestones")[taskGroupId]?.startDate)}
                onChange={(date) => {
                  if (!date) return;
                  const milestones = getValues("milestones");
                  setValue("milestones", {
                    ...milestones,
                    [taskGroupId]: {
                      ...milestones[taskGroupId],
                      startDate: date.toISOString(),
                    },
                  });
                }}
              />
            </FormControl>
          </Grid>
          <Grid item xs>
            <FormControl fullWidth>
              <DatePicker
                label={t("Stage End Date")}
                value={dayjs(getValues("milestones")[taskGroupId]?.endDate)}
                onChange={(date) => {
                  if (!date) return;
                  const milestones = getValues("milestones");
                  setValue("milestones", {
                    ...milestones,
                    [taskGroupId]: {
                      ...milestones[taskGroupId],
                      endDate: date.toISOString(),
                    },
                  });
                }}
              />
            </FormControl>
          </Grid>
        </Grid>
      </LocalizationProvider>
      <Box
        sx={(theme) => ({
          marginBlockStart: 1,
          marginInline: -3,
          borderBottom: `1px solid ${theme.palette.divider}`,
        })}
      >
        <StyledDataGrid
          apiRef={apiRef}
          autoHeight
          sx={{
            "--DataGrid-overlayHeight": "100px",
            ".MuiDataGrid-row .MuiDataGrid-cell.hasError": {
              border: "1px solid",
              borderColor: "error.main",
            },
          }}
          disableColumnMenu
          editMode="row"
          rows={payments}
          rowModesModel={rowModesModel}
          onRowModesModelChange={setRowModesModel}
          onRowEditStop={handleEditStop}
          processRowUpdate={processRowUpdate}
          columns={columns}
          getCellClassName={(params) => {
            return params.row._error === params.field ? "hasError" : "";
          }}
          slots={{
            footer: FooterToolbar,
            noRowsOverlay: NoRowsOverlay,
          }}
          slotProps={{
            footer: { onAdd: addRow },
          }}
        />
      </Box>
    </Stack>
  );
};

const NoRowsOverlay: React.FC = () => {
  const { t } = useTranslation();
  return (
    <Box
      display="flex"
      justifyContent="center"
      alignItems="center"
      height="100%"
    >
      <Typography variant="caption">
        {t("Add some payment milestones!")}
      </Typography>
    </Box>
  );
};

const FooterToolbar: React.FC<FooterPropsOverrides> = ({ onAdd }) => {
  const { t } = useTranslation();
  return (
    <GridToolbarContainer sx={{ p: 2 }}>
      <Button
        variant="outlined"
        startIcon={<Add />}
        onClick={onAdd}
        size="small"
      >
        {t("Add Payment Milestone")}
      </Button>
    </GridToolbarContainer>
  );
};

export default MilestoneSection;