Pien-Studio/apps/web/hooks/use-face-blur-workflow.ts

import React from "react";
import {
  getLayerRuntimeSource,
  type FaceBlurMethod,
  type Layer,
  type LayerEffect,
} from "@pien-studio/types";
import type { FaceDetectionOverlay } from "./use-face-detection";

type FaceBlurPreview = {
  layerId: string;
  effects: LayerEffect[];
};

type FaceSelectionState = {
  key: string;
  indices: number[];
};

type UseFaceBlurWorkflowOptions = {
  selectedLayer: Layer | null;
  faceDetectionsLayerId: string | null;
  faceDetections: FaceDetectionOverlay[];
  setLayerEffect: (layerId: string, effect: LayerEffect) => void;
  removeLayerEffect: (layerId: string, kind: LayerEffect["kind"]) => void;
};

export function useFaceBlurWorkflow(options: UseFaceBlurWorkflowOptions) {
  const {
    selectedLayer,
    faceDetectionsLayerId,
    faceDetections,
    setLayerEffect,
    removeLayerEffect,
  } = options;
  const [blurMethod, setBlurMethod] =
    React.useState<FaceBlurMethod>("gaussian");
  const [blurAmount, setBlurAmount] = React.useState(14);
  const [censorColor, setCensorColor] = React.useState("#111111");
  const [faceSelection, setFaceSelection] =
    React.useState<FaceSelectionState | null>(null);
  const hasDetectableSelection = Boolean(
    selectedLayer &&
    selectedLayer.type === "raster" &&
    faceDetectionsLayerId === selectedLayer.id,
  );

  const faceBlurEffect = React.useMemo(
    () =>
      selectedLayer?.effects.find(
        (e): e is Extract<LayerEffect, { kind: "face-blur" }> =>
          e.kind === "face-blur",
      ) ?? null,
    [selectedLayer?.effects],
  );

  const selectionKey = `${selectedLayer?.id ?? "none"}:${faceDetectionsLayerId ?? "none"}:${faceDetections.length}:${faceBlurEffect ? "effect" : "empty"}`;

  const defaultSelectedFaceIndices = React.useMemo(() => {
    if (!hasDetectableSelection || faceBlurEffect) return [];
    return faceDetections.map((_, index) => index);
  }, [faceDetections, faceBlurEffect, hasDetectableSelection]);

  const selectedFaceIndices =
    faceSelection?.key === selectionKey
      ? faceSelection.indices
      : defaultSelectedFaceIndices;

  const buildBlurRegions = React.useCallback(
    (indices: number[]) => {
      if (!selectedLayer || selectedLayer.type !== "raster") return [];
      if (
        faceDetectionsLayerId !== selectedLayer.id ||
        faceDetections.length === 0
      )
        return [];
      const indexSet = new Set(indices);
      return faceDetections
        .filter((_, index) => indexSet.has(index))
        .map((face) => {
          const baseWidth = Math.max(1, selectedLayer.width);
          const baseHeight = Math.max(1, selectedLayer.height);
          const scaleX = face.sourceWidth / baseWidth;
          const scaleY = face.sourceHeight / baseHeight;
          return {
            x: Math.max(0, Math.floor(face.x * scaleX)),
            y: Math.max(0, Math.floor(face.y * scaleY)),
            width: Math.max(1, Math.floor(face.width * scaleX)),
            height: Math.max(1, Math.floor(face.height * scaleY)),
            sourceWidth: face.sourceWidth,
            sourceHeight: face.sourceHeight,
            censorColor,
          };
        });
    },
    [censorColor, faceDetections, faceDetectionsLayerId, selectedLayer],
  );

  const toggleFaceIndex = React.useCallback(
    (index: number) => {
      setFaceSelection((prev) => {
        const current =
          prev?.key === selectionKey
            ? prev.indices
            : defaultSelectedFaceIndices;
        const indices = current.includes(index)
          ? current.filter((item) => item !== index)
          : [...current, index];
        return { key: selectionKey, indices };
      });
    },
    [defaultSelectedFaceIndices, selectionKey],
  );

  const clearBlur = React.useCallback(() => {
    if (!selectedLayer) return;
    removeLayerEffect(selectedLayer.id, "face-blur");
    setFaceSelection({ key: selectionKey, indices: [] });
  }, [removeLayerEffect, selectedLayer, selectionKey]);

  const blurFaces = React.useCallback(
    (indices: number[]) => {
      if (
        !selectedLayer ||
        selectedLayer.type !== "raster" ||
        !getLayerRuntimeSource(selectedLayer)
      )
        return;
      if (
        faceDetectionsLayerId !== selectedLayer.id ||
        faceDetections.length === 0
      )
        return;
      const regions = buildBlurRegions(indices);
      setLayerEffect(selectedLayer.id, {
        kind: "face-blur",
        enabled: true,
        method: blurMethod,
        amount: blurAmount,
        regions,
        censorColor,
      });
      setFaceSelection({ key: selectionKey, indices: [] });
    },
    [
      blurAmount,
      blurMethod,
      buildBlurRegions,
      censorColor,
      faceDetections.length,
      faceDetectionsLayerId,
      selectedLayer,
      selectionKey,
      setLayerEffect,
    ],
  );

  const faceBlurPreview = React.useMemo<FaceBlurPreview | null>(() => {
    if (
      !hasDetectableSelection ||
      !selectedLayer ||
      selectedLayer.type !== "raster"
    ) {
      return null;
    }

    if (selectedFaceIndices.length === 0) {
      return null;
    }

    return {
      layerId: selectedLayer.id,
      effects: [
        {
          kind: "face-blur",
          enabled: true,
          method: blurMethod,
          amount: blurAmount,
          regions: buildBlurRegions(selectedFaceIndices),
          censorColor,
        },
      ],
    };
  }, [
    blurAmount,
    blurMethod,
    buildBlurRegions,
    censorColor,
    hasDetectableSelection,
    selectedFaceIndices,
    selectedLayer,
  ]);

  return {
    blurMethod,
    setBlurMethod,
    blurAmount,
    setBlurAmount,
    censorColor,
    setCensorColor,
    selectedFaceIndices,
    faceBlurPreview,
    toggleFaceIndex,
    clearBlur,
    blurFaces,
  };
}