frontend/components/generation/RegenerationModal.tsx

"use client";

import React, { useState, useEffect } from "react";
import { X, RefreshCw, CheckCircle2, AlertCircle, Loader2, Sparkles, ChevronDown, Check, ArrowLeft, Image as ImageIcon } from "lucide-react";
import { regenerateImage, getImageModels, confirmImageSelection } from "@/lib/api/endpoints";
import type { ImageRegenerateResponse, AdCreativeDB, ImageModel } from "@/types/api";
import { getImageUrlFallback } from "@/lib/utils/formatters";
import { ProgressBar } from "@/components/ui/ProgressBar";
import { Button } from "@/components/ui/Button";
import { Card, CardContent } from "@/components/ui/Card";

interface RegenerationModalProps {
  isOpen: boolean;
  onClose: () => void;
  adId: string;
  ad?: AdCreativeDB | null;
  onSuccess?: (result: ImageRegenerateResponse) => void;
}

type RegenerationStep = "idle" | "input" | "regenerating" | "compare" | "saving" | "complete" | "error";

export const RegenerationModal: React.FC<RegenerationModalProps> = ({
  isOpen,
  onClose,
  adId,
  ad,
  onSuccess,
}) => {
  const [step, setStep] = useState<RegenerationStep>("idle");
  const [progress, setProgress] = useState(0);
  const [result, setResult] = useState<ImageRegenerateResponse | null>(null);
  const [error, setError] = useState<string | null>(null);
  const [selectedModel, setSelectedModel] = useState<string | null>(null);
  const [models, setModels] = useState<ImageModel[]>([]);
  const [defaultModel, setDefaultModel] = useState<string>("");
  const [loadingModels, setLoadingModels] = useState(false);
  const [selectedImage, setSelectedImage] = useState<"original" | "new" | null>(null);
  const [savingSelection, setSavingSelection] = useState(false);
  const [regenImgUseFallback, setRegenImgUseFallback] = useState(false);

  useEffect(() => {
    if (result) setRegenImgUseFallback(false);
  }, [result]);

  useEffect(() => {
    if (isOpen) {
      setStep("input");
      setProgress(0);
      setResult(null);
      setError(null);
      setSelectedImage(null);
      // Load available models
      loadModels();
    } else {
      // Reset state when modal closes
      setStep("idle");
      setProgress(0);
      setResult(null);
      setError(null);
      setSelectedModel(null);
      setSelectedImage(null);
    }
  }, [isOpen]);

  const loadModels = async () => {
    setLoadingModels(true);
    try {
      const response = await getImageModels();
      setModels(response.models);
      setDefaultModel(response.default);
      // Set initial selection to current ad's model or default
      setSelectedModel(ad?.image_model || response.default);
    } catch (err: any) {
      console.error("Failed to load models:", err);
      // Use fallback models if API fails
      setModels([
        { key: "nano-banana", id: "google/nano-banana", uses_dimensions: false },
        { key: "nano-banana-pro", id: "google/nano-banana-pro", uses_dimensions: false },
        { key: "z-image-turbo", id: "prunaai/z-image-turbo", uses_dimensions: true },
        { key: "imagen-4-ultra", id: "google/imagen-4-ultra", uses_dimensions: false },
        { key: "recraft-v3", id: "recraft-ai/recraft-v3", uses_dimensions: false },
        { key: "ideogram-v3", id: "ideogram-ai/ideogram-v3-quality", uses_dimensions: false },
        { key: "photon", id: "luma/photon", uses_dimensions: false },
        { key: "seedream-3", id: "bytedance/seedream-3", uses_dimensions: false },
      ]);
      setDefaultModel("nano-banana");
      setSelectedModel(ad?.image_model || "nano-banana");
    } finally {
      setLoadingModels(false);
    }
  };

  const handleRegenerate = async () => {
    setStep("regenerating");
    setProgress(0);
    setError(null);
    setResult(null);
    setSelectedImage(null);

    try {
      // Simulate progress updates
      const progressInterval = setInterval(() => {
        setProgress((prev) => {
          if (prev < 90) {
            return prev + 3;
          }
          return prev;
        });
      }, 400);

      // Actually perform the regeneration (preview mode)
      const response = await regenerateImage({
        image_id: adId,
        image_model: selectedModel,
        preview_only: true,
      });

      clearInterval(progressInterval);
      setProgress(100);

      if (response.status === "success") {
        setResult(response);
        setStep("compare"); // Go to comparison step
      } else {
        setStep("error");
        setError(response.error || "Regeneration failed");
      }
    } catch (err: any) {
      setStep("error");
      setError(err.response?.data?.detail || err.message || "Failed to regenerate image");
      setProgress(0);
    }
  };

  const handleConfirmSelection = async () => {
    if (!selectedImage || !result) return;

    setSavingSelection(true);
    setStep("saving");

    try {
      if (selectedImage === "new") {
        // Save the new image
        await confirmImageSelection({
          image_id: adId,
          selection: "new",
          new_image_url: result.regenerated_image?.image_url,
          new_r2_url: result.regenerated_image?.r2_url,
          new_filename: result.regenerated_image?.filename,
          new_model: result.regenerated_image?.model_used,
          new_seed: result.regenerated_image?.seed_used,
        });
      } else {
        // Keep the original
        await confirmImageSelection({
          image_id: adId,
          selection: "original",
        });
      }

      setStep("complete");
    } catch (err: any) {
      setStep("error");
      setError(err.response?.data?.detail || err.message || "Failed to save selection");
    } finally {
      setSavingSelection(false);
    }
  };

  const getModelDisplayName = (key: string) => {
    const displayNames: Record<string, string> = {
      "z-image-turbo": "Z-Image Turbo",
      "nano-banana": "Nano Banana",
      "nano-banana-pro": "Nano Banana Pro",
      "imagen-4": "Imagen 4",
      "imagen-4-ultra": "Imagen 4 Ultra",
      "recraft-v3": "Recraft V3",
      "ideogram-v3": "Ideogram V3 Quality",
      "photon": "Luma Photon",
      "seedream-3": "SeedReam 3",
      "gpt-image-1.5": "GPT Image 1.5 (OpenAI)",
    };
    return displayNames[key] || key;
  };

  if (!isOpen) return null;

  return (
    <div className="fixed inset-0 z-50 flex items-center justify-center p-4 bg-black/60 backdrop-blur-sm">
      <div className="bg-white rounded-2xl shadow-2xl max-w-4xl w-full max-h-[90vh] overflow-y-auto">
        {/* Header */}
        <div className="sticky top-0 bg-white border-b border-gray-200 px-6 py-4 rounded-t-2xl z-10">
          <div className="flex items-center justify-between">
            <div className="flex items-center gap-3">
              <div className="p-2 bg-gradient-to-r from-purple-500 to-pink-500 rounded-lg">
                <RefreshCw className="h-5 w-5 text-white" />
              </div>
              <div>
                <h2 className="text-xl font-bold text-gray-900">Regenerate Image</h2>
                <p className="text-sm text-gray-500">
                  {step === "input" && "Generate a new version of your image"}
                  {step === "regenerating" && "Creating a new image..."}
                  {step === "compare" && "Compare and choose your preferred image"}
                  {step === "saving" && "Saving your selection..."}
                  {step === "complete" && "Selection saved!"}
                  {step === "error" && "Something went wrong"}
                </p>
              </div>
            </div>
            <button
              onClick={onClose}
              className="p-2 hover:bg-gray-100 rounded-lg transition-colors"
            >
              <X className="h-5 w-5 text-gray-500" />
            </button>
          </div>
        </div>

        {/* Content */}
        <div className="p-6 space-y-6">
          {/* Input Step */}
          {step === "input" && (
            <div className="space-y-4">
              <Card variant="glass">
                <CardContent className="pt-6">
                  <div className="space-y-4">
                    {/* Current Model Info */}
                    {ad?.image_model && (
                      <div className="bg-gray-50 border border-gray-200 rounded-lg p-3">
                        <p className="text-xs text-gray-500 mb-1">Current Model</p>
                        <p className="font-medium text-gray-900">{getModelDisplayName(ad.image_model)}</p>
                      </div>
                    )}

                    {/* Model Selection */}
                    <div>
                      <label className="block text-sm font-semibold text-gray-700 mb-2">
                        Select Image Model
                      </label>
                      {loadingModels ? (
                        <div className="flex items-center justify-center py-4">
                          <Loader2 className="h-5 w-5 text-blue-500 animate-spin" />
                          <span className="ml-2 text-sm text-gray-500">Loading models...</span>
                        </div>
                      ) : (
                        <div className="relative">
                          <select
                            value={selectedModel || ""}
                            onChange={(e) => setSelectedModel(e.target.value)}
                            className="w-full px-4 py-3 bg-white border border-gray-300 rounded-lg appearance-none focus:outline-none focus:ring-2 focus:ring-purple-500 focus:border-transparent text-gray-900 font-medium"
                          >
                            {models.map((model) => (
                              <option key={model.key} value={model.key}>
                                {getModelDisplayName(model.key)}
                                {model.key === defaultModel ? " (Default)" : ""}
                                {model.key === ad?.image_model ? " (Current)" : ""}
                              </option>
                            ))}
                          </select>
                          <ChevronDown className="absolute right-3 top-1/2 -translate-y-1/2 h-5 w-5 text-gray-400 pointer-events-none" />
                        </div>
                      )}
                      <p className="text-xs text-gray-500 mt-2">
                        Different models produce different styles. Try a new model for variety!
                      </p>
                    </div>

                    {/* Info about what will happen */}
                    <div className="bg-purple-50 border border-purple-200 rounded-lg p-3 text-sm text-purple-800">
                      <p className="font-semibold mb-1">How it works</p>
                      <ul className="list-disc list-inside space-y-1 text-purple-700">
                        <li>Uses the same prompt as the original image</li>
                        <li>Generates a completely new image with a fresh seed</li>
                        <li>You can compare both and choose which one to keep</li>
                      </ul>
                    </div>
                  </div>
                </CardContent>
              </Card>
            </div>
          )}

          {/* Progress Section */}
          {step === "regenerating" && (
            <div className="space-y-4">
              <div className="flex items-center gap-4">
                <Loader2 className="h-6 w-6 text-blue-500 animate-spin" />
                <div className="flex-1">
                  <p className="font-semibold text-gray-900">Generating new image...</p>
                  <p className="text-sm text-gray-500 mb-2">Using {getModelDisplayName(selectedModel || defaultModel)}</p>
                  <ProgressBar progress={progress} showPercentage={true} className="mt-2" />
                </div>
              </div>
            </div>
          )}

          {/* Comparison Step */}
          {step === "compare" && result && (
            <div className="space-y-6">
              <div className="bg-blue-50 border border-blue-200 rounded-xl p-4">
                <p className="text-sm text-blue-800 font-medium">
                  Click on the image you want to keep. Your selection will be saved when you confirm.
                </p>
              </div>

              {/* Side by side comparison */}
              <div className="grid grid-cols-2 gap-4">
                {/* Original Image */}
                <div
                  className={`relative cursor-pointer rounded-xl overflow-hidden border-4 transition-all ${selectedImage === "original"
                    ? "border-green-500 ring-4 ring-green-200"
                    : "border-gray-200 hover:border-gray-400"
                    }`}
                  onClick={() => setSelectedImage("original")}
                >
                  <div className="absolute top-3 left-3 z-10">
                    <span className="bg-gray-900/80 text-white text-xs font-semibold px-3 py-1.5 rounded-full">
                      Original
                    </span>
                  </div>
                  {selectedImage === "original" && (
                    <div className="absolute top-3 right-3 z-10">
                      <div className="bg-green-500 text-white p-1.5 rounded-full">
                        <Check className="h-4 w-4" />
                      </div>
                    </div>
                  )}
                  {(result.original_image_url || ad?.r2_url || ad?.image_url) ? (
                    <img
                      src={(result.original_image_url || ad?.r2_url || ad?.image_url)!}
                      alt="Original"
                      className="w-full aspect-square object-cover"
                    />
                  ) : (
                    <div className="w-full aspect-square bg-gray-50 flex items-center justify-center text-gray-400">
                      <ImageIcon className="h-8 w-8" />
                    </div>
                  )}
                  <div className="p-3 bg-gray-50">
                    <p className="text-xs text-gray-500">Model</p>
                    <p className="text-sm font-medium text-gray-900">
                      {getModelDisplayName(ad?.image_model || "Unknown")}
                    </p>
                  </div>
                </div>

                {/* New Image */}
                <div
                  className={`relative cursor-pointer rounded-xl overflow-hidden border-4 transition-all ${selectedImage === "new"
                    ? "border-green-500 ring-4 ring-green-200"
                    : "border-gray-200 hover:border-gray-400"
                    }`}
                  onClick={() => setSelectedImage("new")}
                >
                  <div className="absolute top-3 left-3 z-10">
                    <span className="bg-purple-600 text-white text-xs font-semibold px-3 py-1.5 rounded-full">
                      New
                    </span>
                  </div>
                  {selectedImage === "new" && (
                    <div className="absolute top-3 right-3 z-10">
                      <div className="bg-green-500 text-white p-1.5 rounded-full">
                        <Check className="h-4 w-4" />
                      </div>
                    </div>
                  )}
                  <img
                    src={(() => {
                      const { primary, fallback } = getImageUrlFallback(
                        result.regenerated_image?.image_url,
                        result.regenerated_image?.filename,
                        result.regenerated_image?.r2_url
                      );
                      return (regenImgUseFallback && fallback ? fallback : primary ?? fallback) ?? "";
                    })()}
                    alt="Regenerated"
                    className="w-full aspect-square object-cover"
                    onError={() => {
                      const { fallback } = getImageUrlFallback(
                        result.regenerated_image?.image_url,
                        result.regenerated_image?.filename,
                        result.regenerated_image?.r2_url
                      );
                      if (fallback) setRegenImgUseFallback(true);
                    }}
                  />
                  <div className="p-3 bg-purple-50">
                    <p className="text-xs text-gray-500">Model</p>
                    <p className="text-sm font-medium text-gray-900">
                      {getModelDisplayName(result.regenerated_image?.model_used || "Unknown")}
                    </p>
                  </div>
                </div>
              </div>

              {/* Selection hint */}
              {!selectedImage && (
                <p className="text-center text-sm text-gray-500">
                  Click on an image to select it
                </p>
              )}
            </div>
          )}

          {/* Saving Step */}
          {step === "saving" && (
            <div className="flex flex-col items-center justify-center py-8">
              <Loader2 className="h-10 w-10 text-purple-500 animate-spin mb-4" />
              <p className="text-lg font-semibold text-gray-900">Saving your selection...</p>
              <p className="text-sm text-gray-500">
                {selectedImage === "new" ? "Updating to new image" : "Keeping original image"}
              </p>
            </div>
          )}

          {/* Error State */}
          {step === "error" && (
            <div className="bg-red-50 border border-red-200 rounded-xl p-4">
              <div className="flex items-start gap-3">
                <AlertCircle className="h-5 w-5 text-red-500 mt-0.5" />
                <div className="flex-1">
                  <h3 className="font-semibold text-red-900 mb-1">Something went wrong</h3>
                  <p className="text-sm text-red-700">{error}</p>
                </div>
              </div>
            </div>
          )}

          {/* Success State */}
          {step === "complete" && (
            <div className="space-y-4">
              <div className="bg-green-50 border border-green-200 rounded-xl p-4">
                <div className="flex items-center gap-3">
                  <CheckCircle2 className="h-6 w-6 text-green-500" />
                  <div>
                    <h3 className="font-semibold text-green-900">Selection Saved!</h3>
                    <p className="text-sm text-green-700">
                      {selectedImage === "new"
                        ? "Your ad has been updated with the new image."
                        : "The original image has been kept."}
                    </p>
                  </div>
                </div>
              </div>

              {/* Show the selected image */}
              <div className="flex justify-center">
                <div className="max-w-md">
                  {((selectedImage === "new"
                    ? (() => {
                        const { primary, fallback } = getImageUrlFallback(
                          result?.regenerated_image?.image_url,
                          result?.regenerated_image?.filename,
                          result?.regenerated_image?.r2_url
                        );
                        return (regenImgUseFallback && fallback ? fallback : primary ?? fallback) ?? null;
                      })()
                    : (result?.original_image_url || ad?.r2_url || ad?.image_url))) ? (
                    <img
                      src={(selectedImage === "new"
                        ? (() => {
                            const { primary, fallback } = getImageUrlFallback(
                              result?.regenerated_image?.image_url,
                              result?.regenerated_image?.filename,
                              result?.regenerated_image?.r2_url
                            );
                            return (regenImgUseFallback && fallback ? fallback : primary ?? fallback) ?? "";
                          })()
                        : (result?.original_image_url || ad?.r2_url || ad?.image_url))!}
                      alt="Selected"
                      className="w-full rounded-xl border border-gray-200"
                      onError={() => {
                        if (selectedImage === "new") {
                          const { fallback } = getImageUrlFallback(
                            result?.regenerated_image?.image_url,
                            result?.regenerated_image?.filename,
                            result?.regenerated_image?.r2_url
                          );
                          if (fallback) setRegenImgUseFallback(true);
                        }
                      }}
                    />
                  ) : (
                    <div className="w-full aspect-square bg-gray-50 flex items-center justify-center text-gray-400 rounded-xl border border-gray-200">
                      <ImageIcon className="h-8 w-8" />
                    </div>
                  )}
                  <p className="text-center text-sm text-gray-500 mt-2">
                    {selectedImage === "new" ? "New image saved" : "Original image kept"}
                  </p>
                </div>
              </div>
            </div>
          )}
        </div>

        {/* Footer */}
        <div className="sticky bottom-0 bg-gray-50 border-t border-gray-200 px-6 py-4 rounded-b-2xl">
          <div className="flex justify-between gap-3">
            {/* Left side buttons */}
            <div>
              {step === "compare" && (
                <Button
                  onClick={() => {
                    setStep("input");
                    setResult(null);
                    setSelectedImage(null);
                  }}
                  variant="secondary"
                >
                  <ArrowLeft className="h-4 w-4 mr-2" />
                  Generate Another
                </Button>
              )}
            </div>

            {/* Right side buttons */}
            <div className="flex gap-3">
              {step === "input" && (
                <Button
                  onClick={handleRegenerate}
                  variant="primary"
                  disabled={!selectedModel || loadingModels}
                >
                  <Sparkles className="h-4 w-4 mr-2" />
                  Regenerate Image
                </Button>
              )}

              {step === "compare" && (
                <Button
                  onClick={handleConfirmSelection}
                  variant="primary"
                  disabled={!selectedImage || savingSelection}
                >
                  {savingSelection ? (
                    <>
                      <Loader2 className="h-4 w-4 mr-2 animate-spin" />
                      Saving...
                    </>
                  ) : (
                    <>
                      <Check className="h-4 w-4 mr-2" />
                      Confirm Selection
                    </>
                  )}
                </Button>
              )}

              {step === "error" && (
                <Button onClick={() => setStep("input")} variant="primary">
                  Try Again
                </Button>
              )}

              <Button
                onClick={() => {
                  if (step === "complete" && result) {
                    // Call onSuccess when user clicks "Done" to reload the ad
                    onSuccess?.(result);
                  }
                  onClose();
                }}
                variant={step === "complete" ? "primary" : "secondary"}
              >
                {step === "complete" ? "Done" : "Close"}
              </Button>
            </div>
          </div>
        </div>
      </div>
    </div>
  );
};