src/routes/index.tsx

import { createFileRoute } from "@tanstack/react-router";
import { useState, useMemo, useEffect } 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 [customFilename, setCustomFilename] = useState("");
  const [customSubfolder, setCustomSubfolder] = useState("");

  const [e621Username, setE621Username] = useState("");
  const [e621ApiKey, setE621ApiKey] = useState("");

  const [autocompleteEnabled, setAutocompleteEnabled] = useState(true);
  const [activeTagFile, setActiveTagFile] = useState("tags-selected.csv");
  const [availableTagFiles, setAvailableTagFiles] = useState<string[]>([]);
  const [newTagFileUrl, setNewTagFileUrl] = useState("");
  const [downloadingTagFile, setDownloadingTagFile] = useState(false);
  const [tagSuggestions, setTagSuggestions] = useState<{ name: string; count: number }[]>([]);
  const [activeWordInfo, setActiveWordInfo] = useState<{ word: string; start: number; end: number } | null>(null);
  const [focusedSuggestionIdx, setFocusedSuggestionIdx] = useState(-1);

  const fetchTagFiles = async () => {
    try {
      const res = await fetch("/api/tags?action=list");
      const data = await res.json();
      if (data.success && Array.isArray(data.files)) {
        setAvailableTagFiles(data.files);
        if (data.files.length > 0 && !data.files.includes(activeTagFile)) {
          // Default to tags-selected.csv if available
          if (data.files.includes("tags-selected.csv")) {
            setActiveTagFile("tags-selected.csv");
          } else {
            setActiveTagFile(data.files[0]);
          }
        }
      }
    } catch (err) {
      console.error("Failed to list tag files:", err);
    }
  };

  const handleDownloadTagFile = async () => {
    if (!newTagFileUrl.trim()) return;
    setDownloadingTagFile(true);
    try {
      const res = await fetch("/api/tags", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ downloadUrl: newTagFileUrl.trim() }),
      });
      const data = await res.json();
      if (!res.ok) {
        alert(`Error: ${data.error}`);
      } else {
        alert(`Successfully downloaded and decompressed ${data.filename}!`);
        setNewTagFileUrl("");
        await fetchTagFiles();
        setActiveTagFile(data.filename);
      }
    } catch (err) {
      alert(`Download failed: ${err instanceof Error ? err.message : "unknown error"}`);
    } finally {
      setDownloadingTagFile(false);
    }
  };

  const updateActiveWord = (el: HTMLInputElement) => {
    const val = el.value;
    const pos = el.selectionStart;
    if (pos === null || !autocompleteEnabled) {
      setTagSuggestions([]);
      setActiveWordInfo(null);
      return;
    }

    // Find start and end of word under cursor
    let start = pos;
    while (start > 0 && !/\s/.test(val[start - 1])) {
      start--;
    }
    let end = pos;
    while (end < val.length && !/\s/.test(val[end])) {
      end++;
    }

    const word = val.slice(start, end).trim();
    // Exclude colon commands like rating:safe
    if (word && !word.includes(":")) {
      setActiveWordInfo({ word, start, end });
    } else {
      setTagSuggestions([]);
      setActiveWordInfo(null);
    }
  };

  const selectSuggestion = (tagName: string) => {
    if (!activeWordInfo) return;
    const before = tags.slice(0, activeWordInfo.start);
    const after = tags.slice(activeWordInfo.end);
    
    // Replace word and add a trailing space
    const newTags = before + tagName + " " + (after.trim() ? after : "");
    setTags(newTags);
    setTagSuggestions([]);
    setActiveWordInfo(null);

    // Refocus input and place cursor after inserted word + space
    const inputEl = document.getElementById("search-tags-input") as HTMLInputElement;
    if (inputEl) {
      inputEl.focus();
      const cursorTarget = before.length + tagName.length + 1;
      setTimeout(() => {
        inputEl.setSelectionRange(cursorTarget, cursorTarget);
      }, 0);
    }
  };

  const handleAutocompleteKeyDown = (e: React.KeyboardEvent<HTMLInputElement>) => {
    if (tagSuggestions.length > 0) {
      if (e.key === "ArrowDown") {
        e.preventDefault();
        setFocusedSuggestionIdx((prev) => (prev + 1) % tagSuggestions.length);
      } else if (e.key === "ArrowUp") {
        e.preventDefault();
        setFocusedSuggestionIdx((prev) => (prev - 1 + tagSuggestions.length) % tagSuggestions.length);
      } else if (e.key === "Enter") {
        if (focusedSuggestionIdx >= 0 && focusedSuggestionIdx < tagSuggestions.length) {
          e.preventDefault();
          selectSuggestion(tagSuggestions[focusedSuggestionIdx].name);
        } else {
          runSearch(1);
        }
      } else if (e.key === "Escape") {
        e.preventDefault();
        setTagSuggestions([]);
        setActiveWordInfo(null);
      }
    } else {
      if (e.key === "Enter") {
        runSearch(1);
      }
    }
  };

  useEffect(() => {
    if (typeof window !== "undefined") {
      setE621Username(localStorage.getItem("e621_username") || "");
      setE621ApiKey(localStorage.getItem("e621_api_key") || "");
    }
    fetchTagFiles();
  }, []);

  useEffect(() => {
    if (!autocompleteEnabled || !activeWordInfo || !activeWordInfo.word) {
      setTagSuggestions([]);
      return;
    }

    const delay = setTimeout(async () => {
      try {
        const res = await fetch(
          `/api/tags?action=autocomplete&query=${encodeURIComponent(
            activeWordInfo.word
          )}&file=${encodeURIComponent(activeTagFile)}&limit=15`
        );
        const data = await res.json();
        if (data.success && Array.isArray(data.suggestions)) {
          setTagSuggestions(data.suggestions);
          setFocusedSuggestionIdx(-1);
        }
      } catch (err) {
        console.error("Autocomplete fetch failed:", err);
      }
    }, 150);

    return () => clearTimeout(delay);
  }, [activeWordInfo, activeTagFile, autocompleteEnabled]);

  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,
          login: site.id === "e621" ? e621Username : undefined,
          apiKey: site.id === "e621" ? e621ApiKey : undefined,
        },
      });
    },
  });

  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),
          zipName: customFilename.trim() || undefined,
        }),
      });
      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;
      
      let finalName = "";
      if (customFilename.trim()) {
        finalName = customFilename.trim();
        if (!finalName.endsWith(".zip")) {
          finalName += ".zip";
        }
      } else {
        finalName = `grabber-${site.id}-${Date.now()}.zip`;
      }
      
      a.download = finalName;
      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(),
          zipName: customFilename.trim() || undefined,
          subfolder: customSubfolder.trim() || undefined,
        }),
      });
      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">
          <div className="relative flex-1 min-w-64 z-20">
            <Input
              id="search-tags-input"
              placeholder="tags (e.g. rating:safe fluffy)"
              value={tags}
              onChange={(e) => {
                setTags(e.target.value);
                updateActiveWord(e.target);
              }}
              onKeyUp={(e) => updateActiveWord(e.currentTarget)}
              onSelect={(e) => updateActiveWord(e.currentTarget)}
              onFocus={(e) => updateActiveWord(e.currentTarget)}
              onBlur={() => {
                // Short timeout to allow clicking suggestion buttons
                setTimeout(() => {
                  setTagSuggestions([]);
                  setActiveWordInfo(null);
                }, 200);
              }}
              onKeyDown={handleAutocompleteKeyDown}
              className="w-full font-sans text-sm"
              autoComplete="off"
            />
            {tagSuggestions.length > 0 && (
              <div className="absolute left-0 right-0 top-full mt-1 max-h-60 overflow-y-auto bg-popover text-popover-foreground border rounded-md shadow-lg z-50">
                {tagSuggestions.map((suggestion, idx) => (
                  <button
                    key={suggestion.name}
                    type="button"
                    className={`w-full text-left px-3 py-1.5 text-xs flex justify-between items-center transition hover:bg-accent hover:text-accent-foreground ${
                      idx === focusedSuggestionIdx ? "bg-accent text-accent-foreground" : ""
                    }`}
                    onClick={() => selectSuggestion(suggestion.name)}
                  >
                    <span className="font-medium">{suggestion.name}</span>
                    <span className="text-[10px] text-muted-foreground font-mono">
                      {suggestion.count.toLocaleString()}
                    </span>
                  </button>
                ))}
              </div>
            )}
          </div>
          <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="Custom ZIP / Filename"
            value={customFilename}
            onChange={(e) => setCustomFilename(e.target.value)}
            className="w-64 max-w-xs"
          />
          <Input
            placeholder="Exclude tags from ZIP"
            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>

        {/* e621 Authentication Panel */}
        {siteId === "e621" && (
          <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">
              e621 Authentication:
            </span>
            <Input
              placeholder="Username"
              value={e621Username}
              onChange={(e) => {
                setE621Username(e.target.value);
                localStorage.setItem("e621_username", e.target.value);
              }}
              className="w-48 max-w-xs"
            />
            <Input
              type="password"
              placeholder="API Key"
              value={e621ApiKey}
              onChange={(e) => {
                setE621ApiKey(e.target.value);
                localStorage.setItem("e621_api_key", e.target.value);
              }}
              className="w-48 max-w-xs"
            />
            <span className="text-xs text-muted-foreground italic">
              (Credentials stored only in your local browser storage)
            </span>
          </div>
        )}

        {/* Autocomplete Settings 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">
            Autocomplete Settings:
          </span>
          <div className="flex items-center gap-2 mr-4">
            <label htmlFor="autocomplete-toggle" className="text-xs font-medium cursor-pointer">
              ON
            </label>
            <Checkbox
              id="autocomplete-toggle"
              checked={autocompleteEnabled}
              onCheckedChange={(checked) => setAutocompleteEnabled(!!checked)}
            />
          </div>
          {autocompleteEnabled && availableTagFiles.length > 0 && (
            <>
              <span className="text-xs text-muted-foreground mr-1">Tags File:</span>
              <Select value={activeTagFile} onValueChange={setActiveTagFile}>
                <SelectTrigger className="w-48 h-9 text-xs">
                  <SelectValue />
                </SelectTrigger>
                <SelectContent>
                  {availableTagFiles.map((file) => (
                    <SelectItem key={file} value={file} className="text-xs">
                      {file}
                    </SelectItem>
                  ))}
                </SelectContent>
              </Select>
            </>
          )}
          <span className="text-xs text-muted-foreground ml-2 mr-1">Download Tags Link:</span>
          <Input
            placeholder="https://.../tags.csv.gz"
            value={newTagFileUrl}
            onChange={(e) => setNewTagFileUrl(e.target.value)}
            className="w-64 max-w-xs h-9 text-xs"
          />
          <Button
            onClick={handleDownloadTagFile}
            disabled={downloadingTagFile || !newTagFileUrl.trim()}
            variant="outline"
            size="sm"
            className="h-9 text-xs"
          >
            {downloadingTagFile ? (
              <Loader2 className="w-3 h-3 mr-1 animate-spin" />
            ) : (
              <Download className="w-3 h-3 mr-1" />
            )}
            Download
          </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"
          />
          <Input
            placeholder="Custom Subfolder Path (optional)"
            value={customSubfolder}
            onChange={(e) => setCustomSubfolder(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>
  );
}