src/routes/index.tsx

import { createFileRoute } from "@tanstack/react-router";
import { useState, useMemo } from "react";
import { useMutation } from "@tanstack/react-query";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import {
  Select,
  SelectContent,
  SelectItem,
  SelectTrigger,
  SelectValue,
} from "@/components/ui/select";
import { Checkbox } from "@/components/ui/checkbox";
import { Badge } from "@/components/ui/badge";
import { Plus, ChevronLeft, ChevronRight, Download, Search, Trash2, Loader2 } from "lucide-react";
import { useSites } from "@/lib/grabber/useSites";
import { searchPosts } from "@/lib/grabber/search.functions";
import type { NormalizedPost } from "@/lib/grabber/sites";
import { AddSiteDialog } from "@/components/grabber/AddSiteDialog";
import { Lightbox } from "@/components/grabber/Lightbox";

export const Route = createFileRoute("/")({
  head: () => ({
    meta: [
      { title: "Imageboard Grabber" },
      {
        name: "description",
        content: "Search booru-style imageboards, browse thumbnails, and batch-download as ZIP.",
      },
      { property: "og:title", content: "Imageboard Grabber" },
      {
        property: "og:description",
        content: "Search booru-style imageboards, browse thumbnails, and batch-download as ZIP.",
      },
    ],
  }),
  component: Grabber,
});

function Grabber() {
  const { sites, custom, addSite, removeSite, hydrated } = useSites();
  const [siteId, setSiteId] = useState("e621");
  const [tags, setTags] = useState("");
  const [limit, setLimit] = useState(40);
  const [page, setPage] = useState(1);
  const [selected, setSelected] = useState<Set<string>>(new Set());
  const [addOpen, setAddOpen] = useState(false);
  const [lightbox, setLightbox] = useState<NormalizedPost | null>(null);
  const [downloading, setDownloading] = useState(false);
  const [excludeTags, setExcludeTags] = useState("");
  const [hfToken, setHfToken] = useState("");
  const [datasetName, setDatasetName] = useState("");
  const [uploading, setUploading] = useState(false);
  const [uploadStatus, setUploadStatus] = useState("");

  const site = useMemo(() => sites.find((s) => s.id === siteId) ?? sites[0], [sites, siteId]);

  const searchMut = useMutation({
    mutationFn: async (vars: { page: number }) => {
      return await searchPosts({ data: { site, tags, page: vars.page, limit } });
    },
  });

  const posts = searchMut.data?.posts ?? [];
  const err = searchMut.data?.error;

  const runSearch = (p: number) => {
    setPage(p);
    setSelected(new Set());
    searchMut.mutate({ page: p });
  };

  const toggle = (id: string) => {
    setSelected((prev) => {
      const n = new Set(prev);
      if (n.has(id)) n.delete(id);
      else n.add(id);
      return n;
    });
  };

  const selectAll = () => {
    if (selected.size === posts.length) setSelected(new Set());
    else setSelected(new Set(posts.map((p) => p.id)));
  };

  const downloadZip = async () => {
    const chosen = posts.filter((p) => selected.has(p.id));
    if (!chosen.length) return;
    setDownloading(true);
    try {
      const res = await fetch("/api/download-zip", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({
          siteName: site.id,
          posts: chosen,
          excludeTags: excludeTags
            .split(/\s+|,/)
            .map((t) => t.trim().toLowerCase())
            .filter(Boolean),
        }),
      });
      if (!res.ok) {
        const t = await res.text();
        alert(`Download failed: ${t}`);
        return;
      }
      const blob = await res.blob();
      const url = URL.createObjectURL(blob);
      const a = document.createElement("a");
      a.href = url;
      a.download = `grabber-${site.id}-${Date.now()}.zip`;
      document.body.appendChild(a);
      a.click();
      a.remove();
      URL.revokeObjectURL(url);
    } finally {
      setDownloading(false);
    }
  };

  const uploadToDataset = async () => {
    const chosen = posts.filter((p) => selected.has(p.id));
    if (!chosen.length) {
      alert("Please select some images to upload.");
      return;
    }
    if (!hfToken.trim() || !datasetName.trim()) {
      alert("Hugging Face write token and dataset name are required.");
      return;
    }
    setUploading(true);
    setUploadStatus("Uploading to Hugging Face...");
    try {
      const res = await fetch("/api/upload-to-dataset", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({
          siteName: site.id,
          posts: chosen,
          excludeTags: excludeTags
            .split(/\s+|,/)
            .map((t) => t.trim().toLowerCase())
            .filter(Boolean),
          hfToken: hfToken.trim(),
          datasetName: datasetName.trim(),
        }),
      });
      const data = await res.json();
      if (!res.ok) {
        setUploadStatus(`Upload failed: ${data.error || "Unknown error"}`);
      } else {
        setUploadStatus(`Success! Uploaded ZIP as ${data.filename}`);
      }
    } catch (err: unknown) {
      const errMsg = err instanceof Error ? err.message : "Network error";
      setUploadStatus(`Upload failed: ${errMsg}`);
    } finally {
      setUploading(false);
    }
  };

  return (
    <div className="min-h-screen bg-background text-foreground">
      <header className="border-b sticky top-0 bg-background/95 backdrop-blur z-10">
        <div className="max-w-7xl mx-auto px-4 py-3 flex items-center gap-3 flex-wrap">
          <h1 className="text-lg font-bold">Grabber</h1>
          <div className="flex-1" />
          <Select value={siteId} onValueChange={setSiteId}>
            <SelectTrigger className="w-40">
              <SelectValue />
            </SelectTrigger>
            <SelectContent>
              {sites.map((s) => (
                <SelectItem key={s.id} value={s.id}>
                  {s.name}
                </SelectItem>
              ))}
            </SelectContent>
          </Select>
          {hydrated && custom.some((c) => c.id === siteId) && (
            <Button
              variant="ghost"
              size="icon"
              onClick={() => {
                removeSite(siteId);
                setSiteId("e621");
              }}
              title="Remove site"
            >
              <Trash2 className="w-4 h-4" />
            </Button>
          )}
          <Button variant="outline" size="sm" onClick={() => setAddOpen(true)}>
            <Plus className="w-4 h-4 mr-1" /> Add site
          </Button>
        </div>
        <div className="max-w-7xl mx-auto px-4 pb-3 flex gap-2 flex-wrap items-center">
          <Input
            placeholder="tags (e.g. rating:safe fluffy)"
            value={tags}
            onChange={(e) => setTags(e.target.value)}
            onKeyDown={(e) => e.key === "Enter" && runSearch(1)}
            className="flex-1 min-w-64"
          />
          <Input
            type="number"
            min={1}
            max={200}
            value={limit}
            onChange={(e) => setLimit(Number(e.target.value) || 40)}
            className="w-24"
          />
          <Button onClick={() => runSearch(1)} disabled={searchMut.isPending}>
            {searchMut.isPending ? (
              <Loader2 className="w-4 h-4 mr-1 animate-spin" />
            ) : (
              <Search className="w-4 h-4 mr-1" />
            )}
            Search
          </Button>
          <div className="flex items-center gap-1">
            <Button
              size="icon"
              variant="outline"
              disabled={page <= 1 || searchMut.isPending}
              onClick={() => runSearch(page - 1)}
            >
              <ChevronLeft className="w-4 h-4" />
            </Button>
            <span className="text-sm px-2 tabular-nums">page {page}</span>
            <Button
              size="icon"
              variant="outline"
              disabled={searchMut.isPending || posts.length === 0}
              onClick={() => runSearch(page + 1)}
            >
              <ChevronRight className="w-4 h-4" />
            </Button>
          </div>
          <Button variant="outline" size="sm" onClick={selectAll} disabled={!posts.length}>
            {selected.size === posts.length && posts.length > 0 ? "Deselect all" : "Select all"}
          </Button>
          <span className="text-sm text-muted-foreground">selected: {selected.size}</span>
          <Input
            placeholder="Exclude tags from ZIP (e.g. fluffy, safety)"
            value={excludeTags}
            onChange={(e) => setExcludeTags(e.target.value)}
            className="w-64 max-w-xs"
          />
          <Button onClick={downloadZip} disabled={!selected.size || downloading}>
            {downloading ? (
              <Loader2 className="w-4 h-4 mr-1 animate-spin" />
            ) : (
              <Download className="w-4 h-4 mr-1" />
            )}
            Download ZIP
          </Button>
        </div>

        {/* Hugging Face Dataset Integration Panel */}
        <div className="max-w-7xl mx-auto px-4 pb-3 flex gap-2 flex-wrap items-center border-t pt-3 mt-1 bg-muted/20">
          <span className="font-semibold text-xs text-muted-foreground mr-1 uppercase tracking-wider">
            Hugging Face Upload:
          </span>
          <Input
            type="password"
            placeholder="HF Write Token"
            value={hfToken}
            onChange={(e) => setHfToken(e.target.value)}
            className="w-48 max-w-xs"
          />
          <Input
            placeholder="dataset-username/dataset-name"
            value={datasetName}
            onChange={(e) => setDatasetName(e.target.value)}
            className="w-64 max-w-xs"
          />
          <Button
            onClick={uploadToDataset}
            disabled={!selected.size || uploading || !hfToken.trim() || !datasetName.trim()}
            variant="secondary"
          >
            {uploading ? (
              <Loader2 className="w-4 h-4 mr-1 animate-spin" />
            ) : (
              <Plus className="w-4 h-4 mr-1" />
            )}
            Upload to HF Dataset
          </Button>
          {uploadStatus && (
            <span
              className={`text-xs font-semibold ${
                uploadStatus.includes("failed") ? "text-destructive" : "text-emerald-600"
              }`}
            >
              {uploadStatus}
            </span>
          )}
        </div>
      </header>

      <main className="max-w-7xl mx-auto p-4">
        {err && (
          <div className="mb-4 p-3 rounded bg-destructive/10 text-destructive text-sm">{err}</div>
        )}
        {searchMut.isPending && (
          <div className="text-center text-muted-foreground py-12">Loading…</div>
        )}
        {!searchMut.isPending && posts.length === 0 && searchMut.isSuccess && !err && (
          <div className="text-center text-muted-foreground py-12">No results.</div>
        )}
        {!searchMut.isPending && !searchMut.isSuccess && (
          <div className="text-center text-muted-foreground py-12">
            Enter tags and press Search.
          </div>
        )}
        <div className="grid grid-cols-2 sm:grid-cols-3 md:grid-cols-4 lg:grid-cols-5 xl:grid-cols-6 gap-3">
          {posts.map((p) => {
            const isSel = selected.has(p.id);
            return (
              <div
                key={p.id}
                className={`relative group rounded overflow-hidden border cursor-pointer transition ${isSel ? "ring-2 ring-primary" : ""}`}
                onClick={() => setLightbox(p)}
              >
                <div className="aspect-square bg-muted">
                  <img
                    src={p.previewUrl}
                    alt={`post ${p.id}`}
                    loading="lazy"
                    className="w-full h-full object-cover"
                    onError={(e) => (e.currentTarget.style.opacity = "0.2")}
                  />
                </div>
                <div
                  className="absolute top-1 left-1"
                  onClick={(e) => {
                    e.stopPropagation();
                    toggle(p.id);
                  }}
                >
                  <Checkbox checked={isSel} className="bg-background/90 border-2" />
                </div>
                <div className="absolute bottom-1 right-1 flex gap-1">
                  {p.rating && (
                    <Badge variant="secondary" className="text-[10px] px-1 py-0">
                      {p.rating}
                    </Badge>
                  )}
                  {p.score !== undefined && (
                    <Badge variant="secondary" className="text-[10px] px-1 py-0">
                      ★{p.score}
                    </Badge>
                  )}
                </div>
              </div>
            );
          })}
        </div>
      </main>

      <AddSiteDialog
        open={addOpen}
        onOpenChange={setAddOpen}
        onAdd={(s) => {
          addSite(s);
          setSiteId(s.id);
        }}
      />
      <Lightbox post={lightbox} onClose={() => setLightbox(null)} />
    </div>
  );
}