makearmy-app/components/utilities/BackgroundRemoverPanel.tsx

"use client";

import React, { useCallback, useEffect, useMemo, useRef, useState } from "react";
import { Loader2 } from "lucide-react";

// ---------- Preview + batch helpers ----------
const PREVIEW_MAX = 2048; // max long-edge for on-screen previews
const BATCH_SIZES = [2048, 1536, 1280, 1024, 864, 720]; // adaptive preview long-edges
const COOLDOWN_MS = 150; // tiny cooldown between requests to ease VRAM

async function makePreview(blob: Blob, maxEdge = PREVIEW_MAX): Promise<Blob> {
    const bitmap = await createImageBitmap(blob);
    const { width, height } = bitmap;
    const scale = Math.min(1, maxEdge / Math.max(width, height));
    const outW = Math.max(1, Math.round(width * scale));
    const outH = Math.max(1, Math.round(height * scale));
    const canvas = document.createElement("canvas");
    canvas.width = outW;
    canvas.height = outH;
    const ctx = canvas.getContext("2d")!;
    ctx.drawImage(bitmap, 0, 0, outW, outH);
    bitmap.close();
    const outBlob = await new Promise<Blob>((resolve) =>
    canvas.toBlob((b) => resolve(b!), "image/png")
    );
    return outBlob;
}

function revoke(url: string | null | undefined) {
    if (!url) return;
    try {
        URL.revokeObjectURL(url);
    } catch {}
}

// ---------- Methods ----------
type Canonical =
| "ormbg"
| "u2net"
| "basnet"
| "deeplab"
| "tracer"
| "u2net_human_seg"
| "isnet-general-use"
| "isnet-anime"
| "bria"
| "inspyrenet";

const METHODS: { key: Canonical; label: string }[] = [
    { key: "ormbg", label: "ORMBG" },
{ key: "u2net", label: "U2NET" },
{ key: "basnet", label: "BASNET" },
{ key: "deeplab", label: "DEEPLAB" },
{ key: "tracer", label: "TRACER-B7" },
{ key: "u2net_human_seg", label: "U2NET (Human)" },
{ key: "isnet-general-use", label: "ISNET (General)" },
{ key: "isnet-anime", label: "ISNET (Anime)" },
{ key: "bria", label: "BRIA RMBG1.4" },
{ key: "inspyrenet", label: "INSPYRENET" },
];

const DEFAULT_CONCURRENCY = 2;

type Status = "idle" | "pending" | "ok" | "error";

type ResultMap = {
    [K in Canonical]?: {
        fullBlob: Blob;
        previewUrl: string;
        bytes: number;
        ms: number;
    };
};

export default function BackgroundRemoverPage() {
    // ---------- State ----------
    const [file, setFile] = useState<File | null>(null);
    const [sourceUrl, setSourceUrl] = useState<string | null>(null);
    const [natural, setNatural] = useState<{ w: number; h: number } | null>(null);

    const [status, setStatus] = useState<Record<Canonical, Status>>(
        () => Object.fromEntries(METHODS.map((m) => [m.key, "idle"])) as Record<Canonical, Status>
    );
    const [results, setResults] = useState<ResultMap>({});
    const resultsRef = useRef<ResultMap>({});
    useEffect(() => {
        resultsRef.current = results;
    }, [results]);

    const [active, setActive] = useState<Canonical | null>(null);
    const [reveal, setReveal] = useState<number>(50);
    const [gpuSafe, setGpuSafe] = useState(true);

    const frameRef = useRef<HTMLDivElement | null>(null);
    const draggingRef = useRef(false);
    const batchBlobCache = useRef<Map<number, Blob>>(new Map());

    useEffect(() => {
        return () => {
            revoke(sourceUrl);
            Object.values(resultsRef.current).forEach((v) => revoke(v?.previewUrl));
        };
    }, [sourceUrl]);

    // ---------- Styles ----------
    const styles = (
        <style>{`
            html, body { width: 100%; overflow-x: hidden; }
            :root { font-size: 17px; }
            .checkerboard {
                background-size: 24px 24px;
                background-image:
                linear-gradient(45deg,#2a2a2a 25%,transparent 25%),
                    linear-gradient(-45deg,#2a2a2a 25%,transparent 25%),
                    linear-gradient(45deg,transparent 75%,#2a2a2a 75%),
                    linear-gradient(-45deg,transparent 75%,#2a2a2a 75%);
                    background-position: 0 0,0 12px,12px -12px,-12px 0;
            }
            .slider-handle { position: absolute; top: 0; bottom: 0; width: 0; left: calc(var(--reveal, 50) * 1%); }
            .slider-handle::before { content: ""; position: absolute; top: 0; bottom: 0; width: 2px; left: -1px; background: rgba(255,255,255,0.85); }
            .slider-thumb { position: absolute; top: 50%; transform: translate(-50%, -50%); width: 26px; height: 26px; border-radius: 9999px; background: rgba(24,24,27,0.9); border: 1px solid rgba(255,255,255,0.85); display: grid; place-items: center; cursor: ew-resize; }
            /* Mobile: keep the page from panning left/right while using the slider */
            .app-frame { touch-action: pan-y; overscroll-behavior-x: contain; }
            `}</style>
    );

    // ---------- File pick ----------
    const onPick = useCallback(
        async (f: File | null) => {
            revoke(sourceUrl);
            Object.values(resultsRef.current).forEach((v) => revoke(v?.previewUrl));
            batchBlobCache.current.clear();

            setFile(f);
            setResults({});
            setActive(null);
            setReveal(50);
            setStatus(Object.fromEntries(METHODS.map((m) => [m.key, "idle"])) as any);

            if (!f) {
                setSourceUrl(null);
                setNatural(null);
                return;
            }

            try {
                const bmp = await createImageBitmap(f);
                setNatural({ w: bmp.width, h: bmp.height });
                bmp.close();
            } catch {}

            const previewBlob = await makePreview(f, PREVIEW_MAX);
            const previewUrl = URL.createObjectURL(previewBlob);
            setSourceUrl(previewUrl);
        },
        [sourceUrl]
    );

    // Get or create a cached resized blob for batch preview runs
    const getBatchBlob = useCallback(
        async (longEdge: number): Promise<Blob> => {
            const cache = batchBlobCache.current;
            if (cache.has(longEdge)) return cache.get(longEdge)!;
            if (!file) throw new Error("No file selected");
            const b = await makePreview(file, longEdge);
            cache.set(longEdge, b);
            return b;
        },
        [file]
    );

    // ---------- Batch run (adaptive) ----------
    const startAll = useCallback(async () => {
        if (!file) return;

        setResults({});
        setStatus((prev) => {
            const next = { ...prev };
            METHODS.forEach((m) => (next[m.key] = "pending"));
            return next;
        });

        const runOne = async (key: Canonical) => {
            // When GPU-safe is on, try progressively smaller long-edge previews.
            const sizes = gpuSafe ? BATCH_SIZES : [Math.max(natural?.w || 0, natural?.h || 0) || 4096];

            let lastErr: string | null = null;
            const t0 = performance.now();

            for (const size of sizes) {
                try {
                    const blobToSend = gpuSafe ? await getBatchBlob(size) : file!;
                    const fd = new FormData();
                    fd.append("file", blobToSend);
                    fd.append("method", key);

                    // CHANGED: call local proxy instead of a hardcoded service
                    const res = await fetch("/api/bgbye/process", { method: "POST", body: fd });
                    if (!res.ok) {
                        const txt = await res.text().catch(() => "");
                        const retryable = /out of memory|onnxruntime|cuda|allocate|500/i.test(txt);
                        if (gpuSafe && retryable) {
                            lastErr = txt || `HTTP ${res.status}`;
                            continue; // try next smaller size
                        }
                        throw new Error(txt || `HTTP ${res.status}`);
                    }

                    const outBlob = await res.blob();
                    const ms = performance.now() - t0;
                    const previewBlob = await makePreview(outBlob);
                    const previewUrl = URL.createObjectURL(previewBlob);

                    setResults((r) => ({
                        ...r,
                        [key]: { fullBlob: outBlob, previewUrl, bytes: outBlob.size, ms },
                    }));
                    setStatus((s) => ({ ...s, [key]: "ok" }));
                    await new Promise((r) => setTimeout(r, COOLDOWN_MS)); // tiny cooldown
                    return;
                } catch (e: any) {
                    lastErr = e?.message || String(e);
                    if (!gpuSafe) break; // not in adaptive mode, bail after first failure
                }
            }

            setStatus((s) => ({ ...s, [key]: "error" }));
            // (Optional) console.debug("Last error for", key, lastErr);
        };

        const concurrency = gpuSafe ? 1 : DEFAULT_CONCURRENCY;
        const queue = [...METHODS.map((m) => m.key)];
        let inFlight: Promise<void>[] = [];

        const launch = () => {
            while (inFlight.length < concurrency && queue.length) {
                const key = queue.shift()!;
                const p = runOne(key).finally(() => {
                    inFlight = inFlight.filter((q) => q !== p);
                });
                inFlight.push(p);
            }
        };

        launch();
        while (inFlight.length) {
            await Promise.race(inFlight);
            launch();
        }

        setActive((prev) => {
            if (prev) return prev;
            for (const m of METHODS) if ((resultsRef.current as any)[m.key]) return m.key;
            return METHODS[0]?.key ?? null;
        });
    }, [file, gpuSafe, natural, getBatchBlob]);

    // ---------- Upload & slider handlers ----------
    const onDrop = (e: React.DragEvent) => {
        e.preventDefault();
        const f = e.dataTransfer.files?.[0];
        if (f) onPick(f);
    };
        const onSelect = (e: React.ChangeEvent<HTMLInputElement>) => {
            const f = e.target.files?.[0] ?? null;
            onPick(f);
        };

        const aspect = useMemo(() => (!natural ? 16 / 9 : natural.w / natural.h || 16 / 9), [natural]);

        const updateByClientX = useCallback((clientX: number) => {
            const el = frameRef.current;
            if (!el) return;
            const rect = el.getBoundingClientRect();
            const pct = ((clientX - rect.left) / rect.width) * 100;
            setReveal(Math.min(100, Math.max(0, pct)));
        }, []);
        const onMouseDown = (e: React.MouseEvent) => {
            draggingRef.current = true;
            updateByClientX(e.clientX);
        };
        const onMouseMove = (e: React.MouseEvent) => {
            if (!draggingRef.current) return;
            updateByClientX(e.clientX);
        };
        const onMouseUp = () => (draggingRef.current = false);
        const onTouchStart = (e: React.TouchEvent) => {
            draggingRef.current = true;
            updateByClientX(e.touches[0].clientX);
        };
        const onTouchMove = (e: React.TouchEvent) => {
            if (!draggingRef.current) return;
            e.preventDefault(); // keep page from horizontal panning while sliding
            updateByClientX(e.touches[0].clientX);
        };
        const onTouchEnd = () => (draggingRef.current = false);

        // ---------- Active result & actions ----------
        const activeResult = active ? results[active] : undefined;
        const canDownload = Boolean(active && activeResult?.fullBlob);

        const download = () => {
            if (!active || !activeResult) return;
            const a = document.createElement("a");
            const base = file?.name?.replace(/\.[^.]+$/, "") || "image";
            const fullUrl = URL.createObjectURL(activeResult.fullBlob);
            a.href = fullUrl;
            setTimeout(() => revoke(fullUrl), 5000);
            a.download = `${base}_${active}.png`;
            document.body.appendChild(a);
            a.click();
            a.remove();
        };

        // Re-run the selected method on the ORIGINAL file for full-resolution output
        const renderFullRes = useCallback(async () => {
            if (!file || !active) return;
            setStatus((s) => ({ ...s, [active]: "pending" }));
            const t0 = performance.now();
            try {
                const fd = new FormData();
                fd.append("file", file);
                fd.append("method", active);
                // CHANGED: call local proxy instead of a hardcoded service
                const res = await fetch("/api/bgbye/process", { method: "POST", body: fd });
                if (!res.ok) throw new Error(await res.text());
                const outBlob = await res.blob();
                const ms = performance.now() - t0;
                const previewBlob = await makePreview(outBlob);
                const previewUrl = URL.createObjectURL(previewBlob);
                const prev = resultsRef.current[active];
                if (prev) revoke(prev.previewUrl);
                setResults((r) => ({ ...r, [active]: { fullBlob: outBlob, previewUrl, bytes: outBlob.size, ms } }));
                setStatus((s) => ({ ...s, [active]: "ok" }));
            } catch {
                setStatus((s) => ({ ...s, [active]: "error" }));
            }
        }, [file, active]);

        const doneCount = useMemo(() => METHODS.filter((m) => status[m.key] === "ok").length, [status]);
        const pendingCount = useMemo(() => METHODS.filter((m) => status[m.key] === "pending").length, [status]);

        function StatusDot({ s }: { s: Status }) {
            const cls =
            s === "ok"
            ? "bg-emerald-500"
            : s === "pending"
            ? "bg-amber-400 animate-pulse"
            : s === "error"
            ? "bg-rose-500"
            : "bg-zinc-600";
            return <span className={`inline-block w-2 h-2 rounded-full ${cls}`} />;
        }

        // ---------- Render ----------
        return (
            <div className="p-6 text-zinc-100 overflow-x-hidden">
            {styles}

            <div className="mx-auto w-full max-w-[1200px] px-4">
            {/* Header row: title left, back button right */}
            <div className="mb-4 flex items-center justify-between gap-3">
            <h1 className="text-2xl font-semibold">Background Remover</h1>
            <a
            href="https://makearmy.io"
            className="px-3 py-1 rounded-md border border-zinc-700 hover:bg-zinc-800/60 text-sm"
            >
            Back to main
            </a>
            </div>

            {/* Source filename */}
            <div className="text-zinc-400 mb-3">
            <span className="text-zinc-300">Source:</span>{" "}
            {file?.name ?? <span className="italic">— none —</span>}
            </div>

            {/* Preview frame (border removed per your note) */}
            <div
            ref={frameRef}
            className="app-frame checkerboard relative w-full rounded-2xl shadow-inner"
            style={{ aspectRatio: `${aspect}`, maxWidth: "1200px", maxHeight: "80vh", marginInline: "auto" }}
            onDragOver={(e) => e.preventDefault()}
            onDrop={onDrop}
            onMouseDown={onMouseDown}
            onMouseMove={onMouseMove}
            onMouseLeave={onMouseUp}
            onMouseUp={onMouseUp}
            onTouchStart={onTouchStart}
            onTouchMove={onTouchMove}
            onTouchEnd={onTouchEnd}
            >
            {/* Drop hint */}
            {!sourceUrl && (
                <label className="absolute inset-0 grid place-items-center cursor-pointer">
                <input type="file" accept="image/*" className="hidden" onChange={onSelect} />
                <div className="text-zinc-400 border-2 border-dashed border-zinc-600/70 rounded-xl px-6 py-10">
                <div className="text-center">
                <div className="mb-1">Drop an image here</div>
                <div className="text-zinc-500">or click to select a file</div>
                </div>
                </div>
                </label>
            )}

            {/* Before/After */}
            {sourceUrl && (
                <>
                {/* LEFT (BEFORE) */}
                <img
                src={sourceUrl}
                loading="lazy"
                decoding="async"
                className="absolute inset-0 w-full h-full object-contain select-none"
                alt="Source"
                style={{ clipPath: `inset(0 0 0 ${reveal}%)` }}
                draggable={false}
                />

                {/* RIGHT (AFTER) */}
                {activeResult ? (
                    <img
                    src={activeResult.previewUrl}
                    loading="lazy"
                    decoding="async"
                    className="absolute inset-0 w-full h-full object-contain select-none pointer-events-none"
                    alt="Result"
                    style={{ clipPath: `inset(0 ${100 - reveal}% 0 0)` }}
                    draggable={false}
                    />
                ) : status[active as Canonical] === "pending" ? (
                    <div className="absolute inset-0 grid place-items-center">
                    <Loader2 className="animate-spin" />
                    </div>
                ) : null}

                {/* Divider & Thumb */}
                <div
                className="slider-handle"
                style={{ "--reveal": `${Math.min(100, Math.max(0, reveal))}` } as React.CSSProperties}
                >
                <div className="slider-thumb">
                <div className="w-1.5 h-4 bg-white/80 rounded" />
                </div>
                </div>
                </>
            )}
            </div>

            {/* Method options – EVEN GRID under the preview */}
            <div className="mt-4 grid grid-cols-2 sm:grid-cols-3 md:grid-cols-4 lg:grid-cols-5 gap-2">
            {METHODS.map(({ key, label }) => (
                <button
                key={key}
                className={`w-full justify-center px-3 py-2 rounded-md border flex items-center gap-2 ${
                    active === key ? "border-blue-400 bg-blue-500/20" : "border-zinc-700 hover:bg-zinc-800/60"
                }`}
                onClick={() => setActive(key)}
                disabled={!file}
                title={!file ? "Select a file first" : label}
                >
                <StatusDot s={status[key]} />
                <span className="truncate">{label}</span>
                </button>
            ))}
            </div>

            {/* Actions & global status */}
            <div className="mt-4 flex items-center gap-3 flex-wrap">
            <button
            onClick={startAll}
            className="px-3 py-1 rounded-md border border-zinc-700 hover:bg-zinc-800/60 flex items-center gap-2 order-0"
            disabled={!file || pendingCount > 0}
            title={!file ? "Select a file first" : pendingCount > 0 ? "Running…" : "Run all methods"}
            >
            {pendingCount > 0 && <Loader2 className="animate-spin w-4 h-4" />}{" "}
            {pendingCount > 0 ? `Running… ${doneCount}/${METHODS.length}` : "Run all methods"}
            </button>

            {/* GPU-safe toggle */}
            <label className="flex items-center gap-2 text-sm text-zinc-300 cursor-pointer select-none order-1">
            <input type="checkbox" checked={gpuSafe} onChange={(e) => setGpuSafe(e.target.checked)} /> GPU-safe mode
            </label>

            <div className="text-zinc-400 text-sm order-2">
            {file ? (
                pendingCount > 0 ? (
                    <span>Processing… {doneCount}/{METHODS.length} finished</span>
                ) : doneCount > 0 ? (
                    <span>Done: {doneCount} methods succeeded</span>
                ) : (
                    <span>Ready. Click Run all methods</span>
                )
            ) : (
                <span>Drop an image to begin</span>
            )}
            </div>

            {/* Right-side controls */}
            <div className="sm:ml-auto flex items-center gap-3 w-full sm:w-auto order-3">
            <input
            type="range"
            min={0}
            max={100}
            value={reveal}
            onChange={(e) => setReveal(parseInt(e.target.value, 10))}
            className="w-full sm:w-56"
            title="Slide to compare before/after"
            />

            <button
            onClick={renderFullRes}
            disabled={!file || !active}
            className={`px-3 py-1 rounded-md border ${
                file && active
                ? "border-sky-600 bg-sky-600/20 hover:bg-sky-600/30"
                : "border-zinc-700 text-zinc-400 cursor-not-allowed"
            }`}
            title={!file ? "Select a file first" : !active ? "Choose a method" : "Render selected method at full resolution"}
            >
            Full-res render
            </button>

            <button
            onClick={download}
            disabled={!canDownload}
            className={`px-3 py-1 rounded-md border ${
                canDownload
                ? "border-emerald-600 bg-emerald-600/20 hover:bg-emerald-600/30"
                : "border-zinc-700 text-zinc-400 cursor-not-allowed"
            }`}
            >
            Download
            </button>
            </div>
            </div>
            </div>
            </div>
        );
}