"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);

          const res = await fetch("/api/bgremove", { 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);
      const res = await fetch("/api/bgremove", { 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 */}
        <div
          ref={frameRef}
          className="app-frame checkerboard relative w-full rounded-2xl border border-zinc-800/80 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: sits next to Run All on mobile; keeps position naturally on larger screens */}
          <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 collapse under on mobile */}
          <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>
  );
}

