esign/apps/remix/app/components/general/document/document-edit-form.tsx

import { useEffect, useState } from 'react';

import { msg } from '@lingui/core/macro';
import { useLingui } from '@lingui/react';
import { DocumentDistributionMethod, DocumentStatus } from '@prisma/client';
import { useNavigate, useSearchParams } from 'react-router';
import { z } from 'zod';

import { DocumentSignatureType } from '@hanzo/esign-lib/constants/document';
import { isValidLanguageCode } from '@hanzo/esign-lib/constants/i18n';
import { DO_NOT_INVALIDATE_QUERY_ON_MUTATION } from '@hanzo/esign-lib/constants/trpc';
import type { TDocument } from '@hanzo/esign-lib/types/document';
import { ZDocumentAccessAuthTypesSchema } from '@hanzo/esign-lib/types/document-auth';
import { getDocumentDataUrlForPdfViewer } from '@hanzo/esign-lib/utils/envelope-download';
import type {
  TDistributeDocumentRequest,
  TDistributeDocumentResponse,
} from '@hanzo/esign-trpc/server/document-router/distribute-document.types';
import type { TGetDocumentResponse } from '@hanzo/esign-trpc/server/document-router/get-document.types';
import type {
  ZUpdateDocumentRequestSchema,
  ZUpdateDocumentResponseSchema,
} from '@hanzo/esign-trpc/server/document-router/update-document.types';
import type {
  ZSetDocumentFieldsRequestSchema,
  ZSetDocumentFieldsResponseSchema,
} from '@hanzo/esign-trpc/server/field-router/schema';
import type {
  ZSetDocumentRecipientsRequestSchema,
  ZSetDocumentRecipientsResponseSchema,
} from '@hanzo/esign-trpc/server/recipient-router/schema';
import { useZapMutation, useZapQuery, useZapUtils } from '@hanzo/esign-trpc/zap/react';
import { cn } from '@hanzo/esign-ui/lib/utils';
import { Card, CardContent } from '@hanzo/esign-ui/primitives/card';
import { AddFieldsFormPartial } from '@hanzo/esign-ui/primitives/document-flow/add-fields';
import type {
  TAddFieldsFormInput,
  TAddFieldsFormSchema,
} from '@hanzo/esign-ui/primitives/document-flow/add-fields.types';
import { AddSettingsFormPartial } from '@hanzo/esign-ui/primitives/document-flow/add-settings';
import type { TAddSettingsFormSchema } from '@hanzo/esign-ui/primitives/document-flow/add-settings.types';
import { AddSignersFormPartial } from '@hanzo/esign-ui/primitives/document-flow/add-signers';
import type {
  TAddSignersFormInput,
  TAddSignersFormSchema,
} from '@hanzo/esign-ui/primitives/document-flow/add-signers.types';
import { AddSubjectFormPartial } from '@hanzo/esign-ui/primitives/document-flow/add-subject';
import type {
  TAddSubjectFormInput,
  TAddSubjectFormSchema,
} from '@hanzo/esign-ui/primitives/document-flow/add-subject.types';
import { DocumentFlowFormContainer } from '@hanzo/esign-ui/primitives/document-flow/document-flow-root';
import type { DocumentFlowStep } from '@hanzo/esign-ui/primitives/document-flow/types';
import { Stepper } from '@hanzo/esign-ui/primitives/stepper';
import { useToast } from '@hanzo/esign-ui/primitives/use-toast';

import { PDFViewer } from '~/components/general/pdf-viewer/pdf-viewer';
import { useCurrentTeam } from '~/providers/team';

export type DocumentEditFormProps = {
  className?: string;
  initialDocument: TDocument;
  documentRootPath: string;
};

type EditDocumentStep = 'settings' | 'signers' | 'fields' | 'subject';
const EditDocumentSteps: EditDocumentStep[] = ['settings', 'signers', 'fields', 'subject'];

export const DocumentEditForm = ({
  className,
  initialDocument,
  documentRootPath,
}: DocumentEditFormProps) => {
  const { toast } = useToast();
  const { _ } = useLingui();

  const navigate = useNavigate();

  const [searchParams] = useSearchParams();
  const team = useCurrentTeam();

  const [isDocumentPdfLoaded, setIsDocumentPdfLoaded] = useState(false);

  const utils = useZapUtils();

  const { data, refetch: refetchDocument } = useZapQuery<TGetDocumentResponse>(
    'document.get',
    {
      documentId: initialDocument.id,
    },
    {
      initialData: initialDocument,
    },
  );

  // `initialData` guarantees a value at runtime; the fallback mirrors the
  // optimistic-update default and keeps `document` non-undefined for the type
  // checker (the ZAP query result is `T | undefined`, unlike tRPC's narrowed
  // `initialData` overload).
  const document = data ?? initialDocument;

  const { recipients, fields } = document;

  const { mutateAsync: updateDocument } = useZapMutation<
    z.infer<typeof ZUpdateDocumentResponseSchema>,
    z.input<typeof ZUpdateDocumentRequestSchema>
  >('document.update', {
    ...DO_NOT_INVALIDATE_QUERY_ON_MUTATION,
    onSuccess: (newData) => {
      utils.setData<TGetDocumentResponse>(
        'document.get',
        {
          documentId: initialDocument.id,
        },
        (oldData) => ({ ...(oldData || initialDocument), ...newData }),
      );
    },
  });

  const { mutateAsync: addFields } = useZapMutation<
    z.infer<typeof ZSetDocumentFieldsResponseSchema>,
    z.input<typeof ZSetDocumentFieldsRequestSchema>
  >('field.setFieldsForDocument', {
    ...DO_NOT_INVALIDATE_QUERY_ON_MUTATION,
    onSuccess: ({ fields: newFields }) => {
      utils.setData<TGetDocumentResponse>(
        'document.get',
        {
          documentId: initialDocument.id,
        },
        (oldData) => ({ ...(oldData || initialDocument), fields: newFields }),
      );
    },
  });

  const { mutateAsync: setRecipients } = useZapMutation<
    z.infer<typeof ZSetDocumentRecipientsResponseSchema>,
    z.infer<typeof ZSetDocumentRecipientsRequestSchema>
  >('recipient.setDocumentRecipients', {
    ...DO_NOT_INVALIDATE_QUERY_ON_MUTATION,
    onSuccess: ({ recipients: newRecipients }) => {
      utils.setData<TGetDocumentResponse>(
        'document.get',
        {
          documentId: initialDocument.id,
        },
        (oldData) => ({ ...(oldData || initialDocument), recipients: newRecipients }),
      );
    },
  });

  const { mutateAsync: sendDocument } = useZapMutation<
    TDistributeDocumentResponse,
    TDistributeDocumentRequest
  >('document.distribute', {
    ...DO_NOT_INVALIDATE_QUERY_ON_MUTATION,
    onSuccess: (newData) => {
      utils.setData<TGetDocumentResponse>(
        'document.get',
        {
          documentId: initialDocument.id,
        },
        (oldData) => ({ ...(oldData || initialDocument), ...newData }),
      );
    },
  });

  const documentFlow: Record<EditDocumentStep, DocumentFlowStep> = {
    settings: {
      title: msg`General`,
      description: msg`Configure general settings for the document.`,
      stepIndex: 1,
    },
    signers: {
      title: msg`Add Signers`,
      description: msg`Add the people who will sign the document.`,
      stepIndex: 2,
    },
    fields: {
      title: msg`Add Fields`,
      description: msg`Add all relevant fields for each recipient.`,
      stepIndex: 3,
    },
    subject: {
      title: msg`Distribute Document`,
      description: msg`Choose how the document will reach recipients`,
      stepIndex: 4,
    },
  };

  const [step, setStep] = useState<EditDocumentStep>(() => {
    // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
    const searchParamStep = searchParams?.get('step') as EditDocumentStep | undefined;

    let initialStep: EditDocumentStep = 'settings';

    if (
      searchParamStep &&
      documentFlow[searchParamStep] !== undefined &&
      !(recipients.length === 0 && (searchParamStep === 'subject' || searchParamStep === 'fields'))
    ) {
      initialStep = searchParamStep;
    }

    return initialStep;
  });

  const saveSettingsData = async (data: TAddSettingsFormSchema) => {
    const { timezone, dateFormat, redirectUrl, language, signatureTypes } = data.meta;

    const parsedGlobalAccessAuth = z
      .array(ZDocumentAccessAuthTypesSchema)
      .safeParse(data.globalAccessAuth);

    return updateDocument({
      documentId: document.id,
      data: {
        title: data.title,
        externalId: data.externalId || null,
        visibility: data.visibility,
        globalAccessAuth: parsedGlobalAccessAuth.success ? parsedGlobalAccessAuth.data : [],
        globalActionAuth: data.globalActionAuth ?? [],
      },
      meta: {
        timezone,
        dateFormat,
        redirectUrl,
        language: isValidLanguageCode(language) ? language : undefined,
        typedSignatureEnabled: signatureTypes.includes(DocumentSignatureType.TYPE),
        uploadSignatureEnabled: signatureTypes.includes(DocumentSignatureType.UPLOAD),
        drawSignatureEnabled: signatureTypes.includes(DocumentSignatureType.DRAW),
      },
    });
  };

  const onAddSettingsFormSubmit = async (data: TAddSettingsFormSchema) => {
    try {
      await saveSettingsData(data);
      setStep('signers');
    } catch (err) {
      console.error(err);

      toast({
        title: _(msg`Error`),
        description: _(msg`An error occurred while updating the document settings.`),
        variant: 'destructive',
      });
    }
  };

  const onAddSettingsFormAutoSave = async (data: TAddSettingsFormSchema) => {
    try {
      await saveSettingsData(data);
    } catch (err) {
      console.error(err);

      toast({
        title: _(msg`Error`),
        description: _(msg`An error occurred while auto-saving the document settings.`),
        variant: 'destructive',
      });
    }
  };

  const saveSignersData = async (data: TAddSignersFormInput) => {
    return Promise.all([
      updateDocument({
        documentId: document.id,
        meta: {
          allowDictateNextSigner: data.allowDictateNextSigner,
          signingOrder: data.signingOrder,
        },
      }),

      setRecipients({
        documentId: document.id,
        recipients: data.signers.map((signer) => ({
          ...signer,
          id: signer.nativeId,
          // Explicitly set to null to indicate we want to remove auth if required.
          actionAuth: signer.actionAuth ?? [],
        })),
      }),
    ]);
  };

  const onAddSignersFormAutoSave = async (data: TAddSignersFormInput) => {
    try {
      // For autosave, we need to return the recipients response for form state sync
      const [, recipientsResponse] = await Promise.all([
        updateDocument({
          documentId: document.id,
          meta: {
            allowDictateNextSigner: data.allowDictateNextSigner,
            signingOrder: data.signingOrder,
          },
        }),

        setRecipients({
          documentId: document.id,
          recipients: data.signers.map((signer) => ({
            ...signer,
            id: signer.nativeId,
            // Explicitly set to null to indicate we want to remove auth if required.
            actionAuth: signer.actionAuth ?? [],
          })),
        }),
      ]);

      return recipientsResponse;
    } catch (err) {
      console.error(err);

      toast({
        title: _(msg`Error`),
        description: _(msg`An error occurred while adding signers.`),
        variant: 'destructive',
      });

      throw err; // Re-throw so the autosave hook can handle the error
    }
  };

  const onAddSignersFormSubmit = async (data: TAddSignersFormSchema) => {
    try {
      await saveSignersData(data);

      setStep('fields');
    } catch (err) {
      console.error(err);

      toast({
        title: _(msg`Error`),
        description: _(msg`An error occurred while adding signers.`),
        variant: 'destructive',
      });
    }
  };

  const saveFieldsData = async (data: TAddFieldsFormInput) => {
    return addFields({
      documentId: document.id,
      fields: data.fields.map((field) => ({
        ...field,
        id: field.nativeId,
        envelopeItemId: document.documentData.envelopeItemId,
      })),
    });
  };

  const onAddFieldsFormSubmit = async (data: TAddFieldsFormSchema) => {
    try {
      await saveFieldsData(data);

      // Clear all field data from localStorage
      for (let i = 0; i < localStorage.length; i++) {
        const key = localStorage.key(i);
        if (key && key.startsWith('field_')) {
          localStorage.removeItem(key);
        }
      }

      setStep('subject');
    } catch (err) {
      console.error(err);

      toast({
        title: _(msg`Error`),
        description: _(msg`An error occurred while adding the fields.`),
        variant: 'destructive',
      });
    }
  };

  const onAddFieldsFormAutoSave = async (data: TAddFieldsFormInput) => {
    try {
      await saveFieldsData(data);
      // Don't clear localStorage on auto-save, only on explicit submit
    } catch (err) {
      console.error(err);

      toast({
        title: _(msg`Error`),
        description: _(msg`An error occurred while auto-saving the fields.`),
        variant: 'destructive',
      });
    }
  };

  const saveSubjectData = async (data: TAddSubjectFormInput) => {
    const { subject, message, distributionMethod, emailId, emailReplyTo, emailSettings } =
      data.meta;

    return updateDocument({
      documentId: document.id,
      meta: {
        subject,
        message,
        distributionMethod,
        emailId,
        emailReplyTo,
        emailSettings: emailSettings,
      },
    });
  };

  const sendDocumentWithSubject = async (data: TAddSubjectFormSchema) => {
    const { subject, message, distributionMethod, emailId, emailReplyTo, emailSettings } =
      data.meta;

    return sendDocument({
      documentId: document.id,
      meta: {
        subject,
        message,
        distributionMethod,
        emailId,
        emailReplyTo: emailReplyTo || null,
        emailSettings,
      },
    });
  };

  const onAddSubjectFormSubmit = async (data: TAddSubjectFormSchema) => {
    try {
      await sendDocumentWithSubject(data);

      if (data.meta.distributionMethod === DocumentDistributionMethod.EMAIL) {
        toast({
          title: _(msg`Document sent`),
          description: _(msg`Your document has been sent successfully.`),
          duration: 5000,
        });

        await navigate(documentRootPath);
      } else if (document.status === DocumentStatus.DRAFT) {
        toast({
          title: _(msg`Links Generated`),
          description: _(msg`Signing links have been generated for this document.`),
          duration: 5000,
        });
      } else {
        await navigate(`${documentRootPath}/${document.envelopeId}`);
      }
    } catch (err) {
      console.error(err);

      toast({
        title: _(msg`Error`),
        description: _(msg`An error occurred while sending the document.`),
        variant: 'destructive',
      });
    }
  };

  const onAddSubjectFormAutoSave = async (data: TAddSubjectFormInput) => {
    try {
      // Save form data without sending the document
      await saveSubjectData(data);
    } catch (err) {
      console.error(err);

      toast({
        title: _(msg`Error`),
        description: _(msg`An error occurred while auto-saving the subject form.`),
        variant: 'destructive',
      });
    }
  };

  const currentDocumentFlow = documentFlow[step];

  /**
   * Refresh the data in the background when steps change.
   */
  useEffect(() => {
    void refetchDocument();
  }, [step]);

  return (
    <div className={cn('grid w-full grid-cols-12 gap-8', className)}>
      <Card
        className="relative col-span-12 rounded-xl before:rounded-xl lg:col-span-6 xl:col-span-7"
        gradient
      >
        <CardContent className="p-2">
          <PDFViewer
            key={document.envelopeItems[0]?.id}
            data={getDocumentDataUrlForPdfViewer({
              envelopeId: document.envelopeId,
              envelopeItemId: document.envelopeItems[0]?.id,
              documentDataId: initialDocument.documentDataId,
              version: 'current',
              token: undefined,
              presignToken: undefined,
            })}
            scrollParentRef="window"
            onDocumentLoad={() => setIsDocumentPdfLoaded(true)}
          />
        </CardContent>
      </Card>

      <div className="col-span-12 lg:col-span-6 xl:col-span-5">
        <DocumentFlowFormContainer
          className="lg:h-[calc(100vh-6rem)]"
          onSubmit={(e) => e.preventDefault()}
        >
          <Stepper
            currentStep={currentDocumentFlow.stepIndex}
            setCurrentStep={(step) => setStep(EditDocumentSteps[step - 1])}
          >
            <AddSettingsFormPartial
              key={recipients.length}
              documentFlow={documentFlow.settings}
              document={document}
              currentTeamMemberRole={team.currentTeamRole}
              recipients={recipients}
              fields={fields}
              isDocumentPdfLoaded={isDocumentPdfLoaded}
              onSubmit={onAddSettingsFormSubmit}
              onAutoSave={onAddSettingsFormAutoSave}
            />

            <AddSignersFormPartial
              key={document.id}
              documentFlow={documentFlow.signers}
              recipients={recipients}
              signingOrder={document.documentMeta?.signingOrder}
              allowDictateNextSigner={document.documentMeta?.allowDictateNextSigner}
              fields={fields}
              onSubmit={onAddSignersFormSubmit}
              onAutoSave={onAddSignersFormAutoSave}
              isDocumentPdfLoaded={isDocumentPdfLoaded}
            />

            <AddFieldsFormPartial
              key={document.id}
              documentFlow={documentFlow.fields}
              recipients={recipients}
              fields={fields}
              onSubmit={onAddFieldsFormSubmit}
              onAutoSave={onAddFieldsFormAutoSave}
              isDocumentPdfLoaded={isDocumentPdfLoaded}
              teamId={team.id}
            />

            <AddSubjectFormPartial
              key={recipients.length}
              documentFlow={documentFlow.subject}
              document={document}
              recipients={recipients}
              fields={fields}
              onSubmit={onAddSubjectFormSubmit}
              onAutoSave={onAddSubjectFormAutoSave}
              isDocumentPdfLoaded={isDocumentPdfLoaded}
            />
          </Stepper>
        </DocumentFlowFormContainer>
      </div>
    </div>
  );
};