Pien-Studio/apps/web/hooks/use-face-detection.ts

import React from "react";
import type { EditorToolId } from "../store/editor-store";
import {
  buildFacePreviews,
  loadImageFromUri,
  toFaceDetectionOverlays,
} from "../lib/face-detection-utils";

export type FaceDetectionOverlay = {
  x: number;
  y: number;
  width: number;
  height: number;
  sourceWidth: number;
  sourceHeight: number;
  label: string;
  gender?: string;
  genderScore?: number;
};

export type FacePreview = {
  id: string;
  src: string;
};

type SelectedImageLayer = {
  id: string;
  sourceUri: string;
  width?: number;
  height?: number;
};

type UseFaceDetectionOptions = {
  tool: EditorToolId;
  selectedLayerId: string | null;
  selectedImageLayer: SelectedImageLayer | null;
  activeLayerStillSelected: (layerId: string) => boolean;
};

export function useFaceDetection(options: UseFaceDetectionOptions) {
  const {
    tool,
    selectedLayerId,
    selectedImageLayer,
    activeLayerStillSelected,
  } = options;
  const selectedImageLayerId = selectedImageLayer?.id ?? null;
  const selectedImageSourceUri = selectedImageLayer?.sourceUri ?? null;
  const selectedImageWidth = selectedImageLayer?.width;
  const selectedImageHeight = selectedImageLayer?.height;
  const [faceDetections, setFaceDetections] = React.useState<
    FaceDetectionOverlay[]
  >([]);
  const [faceDetectionsLayerId, setFaceDetectionsLayerId] = React.useState<
    string | null
  >(null);
  const [faceStatus, setFaceStatus] = React.useState<
    "idle" | "detecting" | "unsupported"
  >("idle");
  const [facePreviews, setFacePreviews] = React.useState<FacePreview[]>([]);

  const resetFaceState = React.useCallback(
    (status: "idle" | "detecting" | "unsupported" = "idle") => {
      setFaceDetections((prev) => (prev.length === 0 ? prev : []));
      setFaceDetectionsLayerId((prev) => (prev === null ? prev : null));
      setFaceStatus((prev) => (prev === status ? prev : status));
    },
    [],
  );

  React.useEffect(() => {
    let canceled = false;

    async function detectFaces() {
      const layerId = selectedImageLayerId;
      if (tool !== "face" || !layerId || !selectedImageSourceUri) {
        resetFaceState("idle");
        return;
      }

      try {
        resetFaceState("detecting");
        const { detectFaceBoxes } = await import("../lib/face-ml");
        const result = await detectFaceBoxes(selectedImageSourceUri);
        if (canceled) return;
        if (!activeLayerStillSelected(layerId) || tool !== "face") return;
        const overlays = toFaceDetectionOverlays(
          result.faces,
          result,
          selectedImageWidth,
          selectedImageHeight,
        );
        setFaceDetections(overlays);
        setFaceDetectionsLayerId(layerId);
        setFaceStatus((prev) => (prev === "idle" ? prev : "idle"));
      } catch (err) {
        console.error("[face-detection] detectFaces error:", err);
        if (!canceled) {
          resetFaceState("unsupported");
        }
      }
    }

    detectFaces();
    return () => {
      canceled = true;
    };
  }, [
    activeLayerStillSelected,
    resetFaceState,
    selectedImageHeight,
    selectedImageLayerId,
    selectedImageSourceUri,
    selectedImageWidth,
    tool,
  ]);

  React.useEffect(() => {
    if (tool !== "face") {
      setFaceDetections([]);
      setFaceDetectionsLayerId(null);
      setFacePreviews([]);
      return;
    }
    setFaceDetections([]);
    setFaceDetectionsLayerId(null);
    setFacePreviews([]);
  }, [tool, selectedLayerId]);

  React.useEffect(() => {
    let canceled = false;

    async function generateFacePreviews() {
      if (
        tool !== "face" ||
        !selectedImageSourceUri ||
        faceDetections.length === 0
      ) {
        setFacePreviews([]);
        return;
      }

      const image = await loadImageFromUri(selectedImageSourceUri);
      if (!image || canceled) {
        setFacePreviews([]);
        return;
      }

      const previews = buildFacePreviews(
        image,
        faceDetections,
        selectedImageWidth,
        selectedImageHeight,
      );

      if (!canceled) setFacePreviews(previews);
    }

    generateFacePreviews();
    return () => {
      canceled = true;
    };
  }, [
    faceDetections,
    selectedImageHeight,
    selectedImageSourceUri,
    selectedImageWidth,
    tool,
  ]);

  return { faceDetections, faceDetectionsLayerId, facePreviews, faceStatus };
}