src/components/grabber/ImageDescriberModal.tsx

import { useState, useEffect } from "react";
import { Dialog, DialogContent, DialogHeader, DialogTitle, DialogDescription } from "@/components/ui/dialog";
import { Button } from "@/components/ui/button";
import { ScrollArea } from "@/components/ui/scroll-area";
import { Loader2, Download, Check, AlertCircle, Brain, RefreshCw } from "lucide-react";
import { zipSync, strToU8 } from "fflate";
import type { NormalizedPost } from "@/lib/grabber/sites";

interface ImageDescriberModalProps {
  open: boolean;
  onOpenChange: (open: boolean) => void;
  posts: NormalizedPost[];
}

export function ImageDescriberModal({ open, onOpenChange, posts }: ImageDescriberModalProps) {
  const [descriptions, setDescriptionState] = useState<
    Record<string, { text: string; loading: boolean; error?: string }>
  >({});
  const [downloadingZip, setDownloadingZip] = useState<Record<string, boolean>>({});

  useEffect(() => {
    if (!open) return;

    let active = true;
    const runDescriptions = async () => {
      for (const post of posts) {
        if (!active) break;
        
        // Skip if already generating or has description
        const currentState = descriptions[post.id];
        if (currentState && (currentState.text || currentState.loading)) {
          continue;
        }

        setDescriptionState((prev) => ({
          ...prev,
          [post.id]: { text: "", loading: true },
        }));

        try {
          const res = await fetch("/api/describe-image", {
            method: "POST",
            headers: { "Content-Type": "application/json" },
            body: JSON.stringify({ imageUrl: post.fileUrl }),
          });
          const data = await res.json();
          if (!active) break;

          if (data.success && data.description) {
            setDescriptionState((prev) => ({
              ...prev,
              [post.id]: { text: data.description, loading: false },
            }));
          } else {
            setDescriptionState((prev) => ({
              ...prev,
              [post.id]: {
                text: "",
                loading: false,
                error: data.error || "Failed to generate description",
              },
            }));
          }
        } catch (err) {
          if (!active) break;
          setDescriptionState((prev) => ({
            ...prev,
            [post.id]: {
              text: "",
              loading: false,
              error: err instanceof Error ? err.message : "Inference error",
            },
          }));
        }
      }
    };

    runDescriptions();

    return () => {
      active = false;
    };
  }, [open, posts]);

  const handleDescribeSingle = async (post: NormalizedPost) => {
    setDescriptionState((prev) => ({
      ...prev,
      [post.id]: { text: "", loading: true, error: undefined },
    }));

    try {
      const res = await fetch("/api/describe-image", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ imageUrl: post.fileUrl }),
      });
      const data = await res.json();

      if (data.success && data.description) {
        setDescriptionState((prev) => ({
          ...prev,
          [post.id]: { text: data.description, loading: false },
        }));
      } else {
        setDescriptionState((prev) => ({
          ...prev,
          [post.id]: {
            text: "",
            loading: false,
            error: data.error || "Failed to generate description",
          },
        }));
      }
    } catch (err) {
      setDescriptionState((prev) => ({
        ...prev,
        [post.id]: {
          text: "",
          loading: false,
          error: err instanceof Error ? err.message : "Inference error",
        },
      }));
    }
  };

  const downloadSingleDescZip = async (post: NormalizedPost, description: string) => {
    setDownloadingZip((prev) => ({ ...prev, [post.id]: true }));
    try {
      const response = await fetch(post.fileUrl);
      if (!response.ok) throw new Error("Failed to fetch image file");
      const buffer = await response.arrayBuffer();
      const imgBytes = new Uint8Array(buffer);

      const zipFiles: Record<string, Uint8Array> = {};
      const ext = post.ext || "jpg";
      zipFiles[`image.${ext}`] = imgBytes;
      zipFiles[`desc.txt`] = strToU8(description);

      const zipped = zipSync(zipFiles, { level: 0 });
      const blob = new Blob([zipped], { type: "application/zip" });
      const url = URL.createObjectURL(blob);

      const a = document.createElement("a");
      a.href = url;
      a.download = `image_${post.id}_with_description.zip`;
      document.body.appendChild(a);
      a.click();
      a.remove();
      URL.revokeObjectURL(url);
    } catch (err) {
      alert(`Download failed: ${err instanceof Error ? err.message : "unknown error"}`);
    } finally {
      setDownloadingZip((prev) => ({ ...prev, [post.id]: false }));
    }
  };

  // Progress stats
  const describedCount = posts.filter((p) => descriptions[p.id]?.text).length;
  const loadingCount = posts.filter((p) => descriptions[p.id]?.loading).length;

  return (
    <Dialog open={open} onOpenChange={onOpenChange}>
      <DialogContent className="max-w-4xl h-[90vh] flex flex-col p-6">
        <DialogHeader className="flex flex-row justify-between items-center border-b pb-4">
          <div>
            <DialogTitle className="text-xl font-bold flex items-center gap-2">
              <Brain className="w-5 h-5 text-purple-600 animate-pulse" />
              AI Image Describer (GGUF LLaVA)
            </DialogTitle>
            <DialogDescription className="text-xs text-muted-foreground mt-1">
              Described {describedCount} of {posts.length} selected images{" "}
              {loadingCount > 0 && "(Processing sequentially...)"}
            </DialogDescription>
          </div>
        </DialogHeader>

        <ScrollArea className="flex-1 pr-4 py-4">
          <div className="space-y-6">
            {posts.map((post) => {
              const state = descriptions[post.id] || { text: "", loading: false };
              const isZipping = downloadingZip[post.id];

              return (
                <div
                  key={post.id}
                  className="flex flex-col md:flex-row gap-4 p-4 rounded-lg border bg-muted/10 hover:bg-muted/20 transition"
                >
                  <div className="w-full md:w-48 h-48 rounded-md bg-muted overflow-hidden relative flex-shrink-0 flex items-center justify-center">
                    <img
                      src={post.previewUrl}
                      alt={`Post ${post.id}`}
                      className="w-full h-full object-cover"
                    />
                    <div className="absolute top-2 left-2 bg-black/60 text-white text-[10px] px-2 py-0.5 rounded font-mono">
                      ID: {post.id}
                    </div>
                  </div>

                  <div className="flex-1 flex flex-col justify-between min-w-0">
                    <div className="space-y-2">
                      <div className="flex items-center gap-2">
                        <span className="font-semibold text-sm uppercase text-muted-foreground">
                          Description
                        </span>
                        {state.loading && (
                          <span className="flex items-center gap-1 text-xs text-purple-600 font-medium">
                            <Loader2 className="w-3.5 h-3.5 animate-spin" />
                            Analyzing image...
                          </span>
                        )}
                        {state.text && (
                          <span className="flex items-center gap-1 text-xs text-emerald-600 font-medium">
                            <Check className="w-3.5 h-3.5" />
                            Completed
                          </span>
                        )}
                        {state.error && (
                          <span className="flex items-center gap-1 text-xs text-destructive font-medium">
                            <AlertCircle className="w-3.5 h-3.5" />
                            Error: {state.error}
                          </span>
                        )}
                      </div>

                      <div className="text-sm text-foreground bg-muted/40 p-3 rounded-md min-h-[100px] whitespace-pre-wrap">
                        {state.loading && (
                          <div className="flex flex-col items-center justify-center h-16 text-muted-foreground gap-2">
                            <Loader2 className="w-6 h-6 animate-spin text-purple-600" />
                            <p className="text-xs">Running GGUF multimodal inference...</p>
                          </div>
                        )}
                        {!state.loading && !state.text && !state.error && (
                          <span className="text-muted-foreground italic">
                            Waiting in sequence queue...
                          </span>
                        )}
                        {state.error && (
                          <div className="flex flex-col gap-2">
                            <span className="text-destructive text-xs">{state.error}</span>
                            <Button
                              variant="outline"
                              size="sm"
                              className="w-24 text-xs h-8"
                              onClick={() => handleDescribeSingle(post)}
                            >
                              <RefreshCw className="w-3 h-3 mr-1" /> Retry
                            </Button>
                          </div>
                        )}
                        {state.text && <p className="leading-relaxed font-sans">{state.text}</p>}
                      </div>
                    </div>

                    {state.text && (
                      <div className="flex justify-end mt-4">
                        <Button
                          onClick={() => downloadSingleDescZip(post, state.text)}
                          disabled={isZipping}
                          size="sm"
                          className="bg-purple-600 hover:bg-purple-700 text-white font-medium text-xs h-9"
                        >
                          {isZipping ? (
                            <Loader2 className="w-3.5 h-3.5 mr-1 animate-spin" />
                          ) : (
                            <Download className="w-3.5 h-3.5 mr-1" />
                          )}
                          Download image.png + desc.txt
                        </Button>
                      </div>
                    )}
                  </div>
                </div>
              );
            })}
          </div>
        </ScrollArea>
      </DialogContent>
    </Dialog>
  );
}