import { useEffect, useState } from "react";
import { invoke, Channel } from "@tauri-apps/api/core";
import { ImageOff } from "lucide-react";
import { useToast } from "./Toast";

interface Props {
  /** Text-to-image prompt. Changing it triggers a new generation. */
  prompt: string | null;
  /** Bump to force a regenerate with the same prompt. */
  nonce?: number;
  className?: string;
  /** Aspect for the placeholder box while loading/empty. */
  aspect?: string;
}

// ponytail: matches the Rust ImageEvent enum (tagged union).
type ImageEvent =
  | { type: "progress"; data: { step: number | null; total: number | null } }
  | { type: "done"; data: string }
  | { type: "error"; data: string };

/**
 * Generates and caches an image for `prompt` via the Ollama image model.
 * Results are cached on disk by the backend, so re-renders are instant.
 * macOS-only: on other OSes the backend returns an error and we show a
 * placeholder instead of a confusing timeout.
 *
 * Uses the streaming command so the DM sees a real progress bar (driven from
 * Ollama's NDJSON step/total) during the multi-second wait. The image still
 * loads even if no progress events arrive — Done carries the data URL.
 */
export function GeneratedImage({ prompt, nonce = 0, className = "", aspect = "aspect-square" }: Props) {
  const [dataUrl, setDataUrl] = useState<string | null>(null);
  const [loading, setLoading] = useState(false);
  const [unsupported, setUnsupported] = useState(false);
  const [expanded, setExpanded] = useState(false);
  // ponytail: progress for the in-flight generation. step/total may be null
  // if Ollama omits them; we show an indeterminate shimmer then.
  const [progress, setProgress] = useState<{ step: number | null; total: number | null } | null>(null);
  const { addToast } = useToast();

  useEffect(() => {
    if (!prompt) return;
    let cancelled = false;
    setLoading(true);
    setUnsupported(false);
    setProgress(null);

    const channel = new Channel<ImageEvent>();
    channel.onmessage = (event) => {
      if (cancelled) return;
      if (event.type === "progress") {
        setProgress(event.data);
      } else if (event.type === "done") {
        setDataUrl(event.data);
        setLoading(false);
        setProgress(null);
      } else if (event.type === "error") {
        const msg = String(event.data);
        if (msg.includes("macOS-only")) setUnsupported(true);
        else addToast(`Image failed: ${msg}`, "error");
        setLoading(false);
        setProgress(null);
      }
    };

    invoke("generate_image_stream", { req: { prompt }, channel }).catch((e) => {
      if (!cancelled) {
        const msg = String(e);
        if (msg.includes("macOS-only")) setUnsupported(true);
        else addToast(`Image failed: ${msg}`, "error");
        setLoading(false);
        setProgress(null);
      }
    });

    return () => {
      cancelled = true;
    };
  }, [prompt, nonce, addToast]);

  if (!prompt) return null;

  const box = `relative ${aspect} w-full rounded-lg overflow-hidden border border-[var(--color-border-glass)] bg-[var(--color-bg-deep)] flex items-center justify-center ${className}`;

  if (unsupported) {
    return (
      <div className={box}>
        <div className="flex flex-col items-center gap-1 text-[var(--color-text-dim)]">
          <ImageOff size={20} />
          <span className="text-[10px] text-center px-2">Image gen is macOS-only via Ollama</span>
        </div>
      </div>
    );
  }

  if (loading) {
    // ponytail: determinate bar when step/total are known, indeterminate shimmer otherwise.
    const pct = progress && progress.step != null && progress.total ? Math.round((progress.step / progress.total) * 100) : null;
    return (
      <div className={box}>
        <div className="absolute inset-0 skeleton" />
        <div className="relative flex flex-col items-center gap-1 w-full px-3">
          {pct != null ? (
            <div className="w-full h-1 rounded-full bg-[var(--color-bg-surface)] overflow-hidden">
              <div className="h-full bg-[var(--color-gold-bright)] transition-all duration-200" style={{ width: `${pct}%` }} />
            </div>
          ) : (
            <span className="text-[10px] text-[var(--color-text-dim)]">Painting…</span>
          )}
        </div>
      </div>
    );
  }

  return dataUrl ? (
    <div className={box}>
      <img src={dataUrl} alt={prompt} className="w-full h-full object-cover cursor-zoom-in" onClick={() => setExpanded(true)} />
      {expanded && (
        <div
          className="fixed inset-0 z-[60] bg-black/80 flex items-center justify-center p-8"
          onClick={() => setExpanded(false)}
          aria-label="Close image preview"
        >
          <img src={dataUrl} alt={prompt} className="max-w-full max-h-full rounded-lg object-contain" />
        </div>
      )}
    </div>
  ) : null;
}