frigate/web/src/components/overlay/ClassificationSelectionDialog.tsx

import {
  Drawer,
  DrawerClose,
  DrawerContent,
  DrawerDescription,
  DrawerHeader,
  DrawerTitle,
  DrawerTrigger,
} from "@/components/ui/drawer";
import {
  DropdownMenu,
  DropdownMenuContent,
  DropdownMenuItem,
  DropdownMenuLabel,
  DropdownMenuTrigger,
} from "@/components/ui/dropdown-menu";
import {
  Tooltip,
  TooltipContent,
  TooltipTrigger,
} from "@/components/ui/tooltip";
import { isDesktop, isMobile } from "react-device-detect";
import { useTranslation } from "react-i18next";
import { cn } from "@/lib/utils";
import React, { ReactNode, useCallback, useMemo, useState } from "react";
import TextEntryDialog from "./dialog/TextEntryDialog";
import { Button } from "../ui/button";
import axios from "axios";
import { toast } from "sonner";
import { Separator } from "../ui/separator";

type ClassificationSelectionDialogProps = {
  className?: string;
  classes: string[];
  modelName: string;
  image: string;
  onRefresh?: () => void;
  onCategorize?: (category: string) => void;
  excludeCategory?: string;
  dialogLabel?: string;
  tooltipLabel?: string;
  children: ReactNode;
};
export default function ClassificationSelectionDialog({
  className,
  classes,
  modelName,
  image,
  onRefresh,
  onCategorize,
  excludeCategory,
  dialogLabel,
  tooltipLabel,
  children,
}: ClassificationSelectionDialogProps) {
  const { t } = useTranslation(["views/classificationModel"]);

  const onCategorizeImage = useCallback(
    (category: string) => {
      if (onCategorize) {
        onCategorize(category);
        return;
      }

      axios
        .post(`/classification/${modelName}/dataset/categorize`, {
          category,
          training_file: image,
        })
        .then((resp) => {
          if (resp.status == 200) {
            toast.success(t("toast.success.categorizedImage"), {
              position: "top-center",
            });
            onRefresh?.();
          }
        })
        .catch((error) => {
          const errorMessage =
            error.response?.data?.message ||
            error.response?.data?.detail ||
            "Unknown error";
          toast.error(t("toast.error.categorizeFailed", { errorMessage }), {
            position: "top-center",
          });
        });
    },
    [modelName, image, onRefresh, onCategorize, t],
  );

  const filteredClasses = useMemo(
    () =>
      excludeCategory ? classes.filter((c) => c !== excludeCategory) : classes,
    [classes, excludeCategory],
  );

  const isChildButton = useMemo(
    () => React.isValidElement(children) && children.type === Button,
    [children],
  );

  // control
  const [newClass, setNewClass] = useState(false);

  // Non-modal Radix DropdownMenu doesn't propagate wheel events to nested
  // scroll containers, so attach a non-passive listener that scrolls manually.
  const scrollContainerRef = useCallback((el: HTMLDivElement | null) => {
    if (!el || !isDesktop) return;
    const handleWheel = (e: WheelEvent) => {
      if (el.scrollHeight <= el.clientHeight) return;
      e.preventDefault();
      el.scrollTop += e.deltaY;
    };
    el.addEventListener("wheel", handleWheel, { passive: false });
    return () => el.removeEventListener("wheel", handleWheel);
  }, []);

  // components
  const Selector = isDesktop ? DropdownMenu : Drawer;
  const SelectorTrigger = isDesktop ? DropdownMenuTrigger : DrawerTrigger;
  const SelectorContent = isDesktop ? DropdownMenuContent : DrawerContent;
  const SelectorItem = isDesktop
    ? DropdownMenuItem
    : (props: React.HTMLAttributes<HTMLDivElement>) => (
        <DrawerClose asChild>
          <div {...props} className={cn(props.className, "my-2")} />
        </DrawerClose>
      );

  // keep modal false on desktop to prevent dismissable layer pointer events
  // issue with dialog auto-close
  return (
    <div className={className ?? "flex"}>
      <TextEntryDialog
        open={newClass}
        setOpen={setNewClass}
        title={t("createCategory.new")}
        onSave={(newCat) => onCategorizeImage(newCat)}
      />
      <Selector {...(isDesktop ? { modal: false } : {})}>
        <Tooltip>
          <TooltipTrigger asChild={isChildButton}>
            <SelectorTrigger asChild>{children}</SelectorTrigger>
          </TooltipTrigger>
          <TooltipContent>
            {tooltipLabel ?? t("categorizeImage")}
          </TooltipContent>
        </Tooltip>
        <SelectorContent
          ref={scrollContainerRef}
          className={cn(
            isDesktop && "scrollbar-container max-h-[40dvh] overflow-y-auto",
            isMobile && "mx-1 gap-2 rounded-t-2xl px-4",
          )}
          onCloseAutoFocus={(e) => e.preventDefault()}
        >
          {isMobile && (
            <DrawerHeader className="sr-only">
              <DrawerTitle>Details</DrawerTitle>
              <DrawerDescription>Details</DrawerDescription>
            </DrawerHeader>
          )}
          <DropdownMenuLabel>
            {dialogLabel ?? t("categorizeImageAs")}
          </DropdownMenuLabel>
          <div className={cn("flex flex-col", isMobile && "gap-2 pb-4")}>
            {filteredClasses
              .sort((a, b) => {
                if (a === "none") return 1;
                if (b === "none") return -1;
                return a.localeCompare(b);
              })
              .map((category) => (
                <SelectorItem
                  key={category}
                  className="flex cursor-pointer gap-2 smart-capitalize"
                  onClick={() => onCategorizeImage(category)}
                >
                  {category === "none"
                    ? t("details.none")
                    : category.replaceAll("_", " ")}
                </SelectorItem>
              ))}
            <Separator />
            <SelectorItem
              className="flex cursor-pointer gap-2 smart-capitalize"
              onClick={() => setNewClass(true)}
            >
              {t("createCategory.new")}
            </SelectorItem>
          </div>
        </SelectorContent>
      </Selector>
    </div>
  );
}