Assist_Design/packages/validation/src/zod-form.ts

/**
 * Framework-agnostic Zod form utilities for React environments.
 * Provides predictable error and touched state handling.
 */

import { useCallback, useMemo, useState } from "react";
import type { FormEvent } from "react";
import { ZodError, type ZodIssue, type ZodSchema } from "zod";

export type FormErrors<T> = Record<string, string | undefined>;
export type FormTouched<T> = Record<string, boolean | undefined>;

export interface ZodFormOptions<T> {
  schema: ZodSchema<T>;
  initialValues: T;
  onSubmit?: (data: T) => Promise<unknown> | unknown;
}

export interface UseZodFormReturn<T extends Record<string, unknown>> {
  values: T;
  errors: FormErrors<T>;
  touched: FormTouched<T>;
  submitError: string | null;
  isSubmitting: boolean;
  isValid: boolean;
  setValue: <K extends keyof T>(field: K, value: T[K]) => void;
  setTouched: <K extends keyof T>(field: K, touched: boolean) => void;
  setTouchedField: <K extends keyof T>(field: K, touched?: boolean) => void;
  validate: () => boolean;
  validateField: <K extends keyof T>(field: K) => boolean;
  handleSubmit: (event?: FormEvent) => Promise<void>;
  reset: () => void;
}

function issuesToErrors<T>(issues: ZodIssue[]): FormErrors<T> {
  const nextErrors: FormErrors<T> = {};

  issues.forEach(issue => {
    const [first, ...rest] = issue.path;
    const key = issue.path.join(".");

    if (typeof first === "string" && nextErrors[first] === undefined) {
      nextErrors[first] = issue.message;
    }

    if (key) {
      nextErrors[key] = issue.message;

      if (rest.length > 0) {
        const topLevelKey = String(first);
        if (nextErrors[topLevelKey] === undefined) {
          nextErrors[topLevelKey] = issue.message;
        }
      }
    } else if (nextErrors._form === undefined) {
      nextErrors._form = issue.message;
    }
  });

  return nextErrors;
}

export function useZodForm<T extends Record<string, unknown>>({
  schema,
  initialValues,
  onSubmit,
}: ZodFormOptions<T>): UseZodFormReturn<T> {
  const [values, setValues] = useState<T>(initialValues);
  const [errors, setErrors] = useState<FormErrors<T>>({});
  const [touched, setTouchedState] = useState<FormTouched<T>>({});
  const [submitError, setSubmitError] = useState<string | null>(null);
  const [isSubmitting, setIsSubmitting] = useState(false);

  const clearFieldError = useCallback((field: keyof T) => {
    const fieldKey = String(field);
    setErrors(prev => {
      const prefix = `${fieldKey}.`;
      const hasDirectError = prev[fieldKey] !== undefined;
      const hasNestedError = Object.keys(prev).some(key => key.startsWith(prefix));

      if (!hasDirectError && !hasNestedError) {
        return prev;
      }

      const next: FormErrors<T> = { ...prev };
      delete next[fieldKey];
      Object.keys(next).forEach(key => {
        if (key.startsWith(prefix)) {
          delete next[key];
        }
      });
      return next;
    });
  }, []);

  const validate = useCallback(() => {
    try {
      schema.parse(values);
      setErrors({});
      return true;
    } catch (error) {
      if (error instanceof ZodError) {
        setErrors(issuesToErrors<T>(error.issues));
      }
      return false;
    }
  }, [schema, values]);

  const validateField = useCallback(<K extends keyof T>(field: K) => {
    const result = schema.safeParse(values);

    if (result.success) {
      clearFieldError(field);
      setErrors(prev => {
        if (prev._form === undefined) {
          return prev;
        }
        const next: FormErrors<T> = { ...prev };
        delete next._form;
        return next;
      });
      return true;
    }

    const fieldKey = String(field);
    const relatedIssues = result.error.issues.filter(issue => issue.path[0] === field);

    setErrors(prev => {
      const next: FormErrors<T> = { ...prev };

      if (relatedIssues.length > 0) {
        const message = relatedIssues[0]?.message ?? "";
        next[fieldKey] = message;
        relatedIssues.forEach(issue => {
          const nestedKey = issue.path.join(".");
          if (nestedKey) {
            next[nestedKey] = issue.message;
          }
        });
      } else {
        delete next[fieldKey];
      }

      const formLevelIssue = result.error.issues.find(issue => issue.path.length === 0);
      if (formLevelIssue) {
        next._form = formLevelIssue.message;
      } else if (relatedIssues.length === 0) {
        delete next._form;
      }

      return next;
    });

    return relatedIssues.length === 0;
  }, [schema, values, clearFieldError]);

  const setValue = useCallback(<K extends keyof T>(field: K, value: T[K]) => {
    setValues(prev => ({ ...prev, [field]: value }));
    clearFieldError(field);
  }, [clearFieldError]);

  const setTouched = useCallback(<K extends keyof T>(field: K, value: boolean) => {
    setTouchedState(prev => ({ ...prev, [String(field)]: value }));
  }, []);

  const setTouchedField = useCallback(<K extends keyof T>(field: K, value: boolean = true) => {
    setTouched(field, value);
    void validateField(field);
  }, [setTouched, validateField]);

  const handleSubmit = useCallback(
    async (event?: FormEvent) => {
      event?.preventDefault();

      if (!onSubmit) {
        return;
      }

      const valid = validate();
      if (!valid) {
        return;
      }

      setIsSubmitting(true);
      setSubmitError(null);
      setErrors(prev => {
        if (prev._form === undefined) {
          return prev;
        }
        const next: FormErrors<T> = { ...prev };
        delete next._form;
        return next;
      });

      try {
        await onSubmit(values);
      } catch (error) {
        const message = error instanceof Error ? error.message : String(error);
        setSubmitError(message);
        setErrors(prev => ({ ...prev, _form: message }));
        console.error("Zod form submission error", error);
        throw error;
      } finally {
        setIsSubmitting(false);
      }
    },
    [validate, onSubmit, values]
  );

  const reset = useCallback(() => {
    setValues(initialValues);
    setErrors({});
    setTouchedState({});
    setSubmitError(null);
    setIsSubmitting(false);
  }, [initialValues]);

  const isValid = useMemo(() => Object.values(errors).every(error => !error), [errors]);

  return {
    values,
    errors,
    touched,
    submitError,
    isSubmitting,
    isValid,
    setValue,
    setTouched,
    setTouchedField,
    validate,
    validateField,
    handleSubmit,
    reset,
  };
}