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],
);
const [newClass, setNewClass] = useState(false);
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);
}, []);
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>
);
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 &&
"max-h-[40dvh] gap-2 overflow-y-auto overflow-x-hidden 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>
);
}