frontend/components/generation/CorrectionModal.tsx

"use client";

import React, { useState, useEffect } from "react";
import { X, Wand2, Image as ImageIcon, CheckCircle2, AlertCircle, Loader2, Sparkles, Download } from "lucide-react";
import { correctImage } from "@/lib/api/endpoints";
import type { ImageCorrectResponse, AdCreativeDB } from "@/types/api";
import { getImageUrlFallback } from "@/lib/utils/formatters";
import { ProgressBar } from "@/components/ui/ProgressBar";
import { Button } from "@/components/ui/Button";
import { Input } from "@/components/ui/Input";
import { Card, CardContent } from "@/components/ui/Card";

interface CorrectionModalProps {
  isOpen: boolean;
  onClose: () => void;
  adId: string;
  imageUrl?: string | null;
  initialInstructions?: string;
  ad?: AdCreativeDB | null;
  onSuccess?: (result: ImageCorrectResponse, keepCorrected: boolean) => void;
}

type CorrectionStep = "idle" | "input" | "analyzing" | "correcting" | "regenerating" | "complete" | "error";

export const CorrectionModal: React.FC<CorrectionModalProps> = ({
  isOpen,
  onClose,
  adId,
  imageUrl,
  initialInstructions = "",
  ad,
  onSuccess,
}) => {
  const [step, setStep] = useState<CorrectionStep>("idle");
  const [progress, setProgress] = useState(0);
  const [result, setResult] = useState<ImageCorrectResponse | null>(null);
  const [error, setError] = useState<string | null>(null);
  const [userInstructions, setUserInstructions] = useState("");
  const [useAutoAnalyze, setUseAutoAnalyze] = useState(false);
  const [keepCorrected, setKeepCorrected] = useState(true);
  const [correctedImgUseFallback, setCorrectedImgUseFallback] = useState(false);

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

  useEffect(() => {
    if (isOpen) {
      setStep("input");
      setProgress(0);
      setResult(null);
      setError(null);
      setUserInstructions(initialInstructions);
      setUseAutoAnalyze(false);
    } else {
      // Reset state when modal closes
      setStep("idle");
      setProgress(0);
      setResult(null);
      setError(null);
      setUserInstructions("");
      setUseAutoAnalyze(false);
      setKeepCorrected(true);
    }
  }, [isOpen]);

  const handleCorrection = async () => {
    if (!userInstructions && !useAutoAnalyze) {
      setError("Please specify what you want to correct or enable auto-analysis");
      return;
    }

    setStep("analyzing");
    setProgress(0);
    setError(null);
    setResult(null);

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

      setStep("analyzing");
      await new Promise((resolve) => setTimeout(resolve, 1000));

      setStep("correcting");
      await new Promise((resolve) => setTimeout(resolve, 1000));

      setStep("regenerating");
      await new Promise((resolve) => setTimeout(resolve, 2000));

      // Actually perform the correction
      const response = await correctImage({
        image_id: adId,
        image_url: imageUrl || undefined,
        user_instructions: userInstructions || undefined,
        auto_analyze: useAutoAnalyze,
      });

      clearInterval(progressInterval);
      setProgress(100);

      if (response.status === "success") {
        setStep("complete");
        setResult(response);
        // Don't call onSuccess immediately - let user see the corrections first
        // onSuccess will be called when user clicks "Done"
      } else {
        setStep("error");
        setError(response.error || "Correction failed");
      }
    } catch (err: any) {
      setStep("error");
      setError(err.response?.data?.detail || err.message || "Failed to correct image");
      setProgress(0);
    }
  };

  const getStepLabel = () => {
    switch (step) {
      case "input":
        return "Specify Corrections";
      case "analyzing":
        return "Analyzing image...";
      case "correcting":
        return "Generating corrections...";
      case "regenerating":
        return "Regenerating with image...";
      case "complete":
        return "Correction complete!";
      case "error":
        return "Error occurred";
      default:
        return "Starting correction...";
    }
  };

  const getStepIcon = () => {
    switch (step) {
      case "complete":
        return <CheckCircle2 className="h-6 w-6 text-green-500" />;
      case "error":
        return <AlertCircle className="h-6 w-6 text-red-500" />;
      default:
        return <Loader2 className="h-6 w-6 text-blue-500 animate-spin" />;
    }
  };

  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-2xl 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-blue-500 to-cyan-500 rounded-lg">
                <Wand2 className="h-5 w-5 text-white" />
              </div>
              <div>
                <h2 className="text-xl font-bold text-gray-900">Correct Image</h2>
                <p className="text-sm text-gray-500">
                  {step === "input"
                    ? "Specify what you want to correct"
                    : "Analyzing and correcting your ad creative"}
                </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">
                    <div>
                      <label className="block text-sm font-semibold text-gray-700 mb-2">
                        What would you like to correct?
                      </label>
                      <Input
                        type="text"
                        placeholder="e.g., 'Fix spelling: change Save 50% to Save 60%' or 'Adjust colors to be brighter' or 'Change headline text to X'"
                        value={userInstructions}
                        onChange={(e) => setUserInstructions(e.target.value)}
                        className="w-full"
                      />
                      <p className="text-xs text-gray-500 mt-2">
                        Be specific about what you want to change. Only the specified changes will be made.
                      </p>
                    </div>

                    <div className="flex items-center gap-3 pt-2 border-t border-gray-200">
                      <input
                        type="checkbox"
                        id="auto-analyze"
                        checked={useAutoAnalyze}
                        onChange={(e) => setUseAutoAnalyze(e.target.checked)}
                        className="w-4 h-4 text-blue-600 border-gray-300 rounded focus:ring-blue-500"
                      />
                      <label htmlFor="auto-analyze" className="text-sm text-gray-700 cursor-pointer">
                        Or let AI automatically analyze and suggest corrections
                      </label>
                    </div>

                    {useAutoAnalyze && (
                      <div className="bg-blue-50 border border-blue-200 rounded-lg p-3 text-sm text-blue-800">
                        <p className="font-semibold mb-1">Auto-Analysis Mode</p>
                        <p>AI will analyze the image for spelling mistakes and visual issues, then suggest corrections.</p>
                      </div>
                    )}
                  </div>
                </CardContent>
              </Card>
            </div>
          )}

          {/* Progress Section */}
          {step !== "input" && step !== "complete" && step !== "error" && (
            <div className="space-y-4">
              <div className="flex items-center gap-4">
                {getStepIcon()}
                <div className="flex-1">
                  <p className="font-semibold text-gray-900">{getStepLabel()}</p>
                  <ProgressBar progress={progress} showPercentage={true} className="mt-2" />
                </div>
              </div>
            </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">Correction Failed</h3>
                  <p className="text-sm text-red-700">{error}</p>
                </div>
              </div>
            </div>
          )}

          {/* Success State */}
          {step === "complete" && result && (
            <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-5 w-5 text-green-500" />
                  <div>
                    <h3 className="font-semibold text-green-900">Correction Complete!</h3>
                    <p className="text-sm text-green-700">Your image has been corrected successfully</p>
                  </div>
                </div>
              </div>

              <div className="space-y-6">
                <div className="grid grid-cols-1 md:grid-cols-2 gap-6">
                  {/* Original Image */}
                  <div className={`space-y-3 p-4 rounded-2xl border-2 transition-all ${!keepCorrected ? 'border-blue-500 bg-blue-50/50 shadow-lg' : 'border-gray-100 hover:border-gray-200'}`}>
                    <div className="flex items-center justify-between">
                      <div className="flex items-center gap-2">
                        <div className={`w-4 h-4 rounded-full border-2 flex items-center justify-center ${!keepCorrected ? 'border-blue-500 bg-blue-500' : 'border-gray-300'}`}>
                          {!keepCorrected && <div className="w-1.5 h-1.5 rounded-full bg-white" />}
                        </div>
                        <h3 className="font-bold text-gray-900">Original</h3>
                      </div>
                      <Button
                        variant={!keepCorrected ? "primary" : "secondary"}
                        size="sm"
                        onClick={() => setKeepCorrected(false)}
                      >
                        {!keepCorrected ? "Selected" : "Keep Original"}
                      </Button>
                    </div>
                    <div className="aspect-square relative rounded-xl overflow-hidden border border-gray-200 cursor-pointer" onClick={() => setKeepCorrected(false)}>
                      {(imageUrl || ad?.r2_url || ad?.image_url) ? (
                        <img
                          src={(imageUrl || ad?.r2_url || ad?.image_url)!}
                          alt="Original"
                          className="w-full h-full object-cover"
                        />
                      ) : (
                        <div className="w-full h-full bg-gray-50 flex flex-col items-center justify-center text-gray-400">
                          <ImageIcon className="h-8 w-8 mb-2" />
                          <span className="text-xs">No image</span>
                        </div>
                      )}
                    </div>
                  </div>

                  {/* Corrected Image */}
                  <div className={`space-y-3 p-4 rounded-2xl border-2 transition-all ${keepCorrected ? 'border-blue-500 bg-blue-50/50 shadow-lg' : 'border-gray-100 hover:border-gray-200'}`}>
                    <div className="flex items-center justify-between">
                      <div className="flex items-center gap-2">
                        <div className={`w-4 h-4 rounded-full border-2 flex items-center justify-center ${keepCorrected ? 'border-blue-500 bg-blue-500' : 'border-gray-300'}`}>
                          {keepCorrected && <div className="w-1.5 h-1.5 rounded-full bg-white" />}
                        </div>
                        <h3 className="font-bold text-gray-900">Corrected</h3>
                      </div>
                      <div className="flex gap-2">
                        <Button
                          variant="secondary"
                          size="sm"
                          onClick={async (e) => {
                            e.stopPropagation();
                            try {
                              const response = await fetch(result.corrected_image!.image_url!);
                              const blob = await response.blob();
                              const url = window.URL.createObjectURL(blob);
                              const link = document.createElement("a");
                              link.href = url;
                              link.download = result.corrected_image!.filename || "corrected-image.png";
                              document.body.appendChild(link);
                              link.click();
                              document.body.removeChild(link);
                              window.URL.revokeObjectURL(url);
                            } catch (err) {
                              window.open(result.corrected_image!.image_url!, "_blank");
                            }
                          }}
                        >
                          <Download className="h-4 w-4" />
                        </Button>
                        <Button
                          variant={keepCorrected ? "primary" : "secondary"}
                          size="sm"
                          onClick={() => setKeepCorrected(true)}
                        >
                          {keepCorrected ? "Selected" : "Use Corrected"}
                        </Button>
                      </div>
                    </div>
                    <div className="aspect-square relative rounded-xl overflow-hidden border border-gray-200 cursor-pointer" onClick={() => setKeepCorrected(true)}>
                      {result.corrected_image && (result.corrected_image.image_url || result.corrected_image.filename) && (() => {
                        const { primary, fallback } = getImageUrlFallback(result.corrected_image?.image_url, result.corrected_image?.filename, result.corrected_image?.r2_url);
                        const src = (correctedImgUseFallback && fallback ? fallback : primary ?? fallback) ?? "";
                        return src ? (
                          <img
                            src={src}
                            alt="Corrected"
                            className="w-full h-full object-cover"
                            onError={() => { if (fallback) setCorrectedImgUseFallback(true); }}
                          />
                        ) : null;
                      })()}
                      <div className="absolute top-2 right-2 px-2 py-1 bg-green-500 text-white text-xs font-bold rounded-md shadow-sm">
                        NEW
                      </div>
                    </div>
                  </div>
                </div>
              </div>

              {/* Show corrections details */}
              {result.corrections && (
                <div className="space-y-4">
                  <h3 className="font-semibold text-gray-900 text-lg">Correction Details</h3>

                  {result.corrections.spelling_corrections && result.corrections.spelling_corrections.length > 0 && (
                    <div>
                      <h4 className="font-semibold text-gray-900 mb-2 flex items-center gap-2">
                        <span className="text-yellow-600">✏️</span>
                        Spelling Corrections
                      </h4>
                      <div className="space-y-2">
                        {result.corrections.spelling_corrections.map((correction, idx) => (
                          <div
                            key={idx}
                            className="bg-yellow-50 border border-yellow-200 rounded-lg p-3 text-sm"
                          >
                            <div className="flex items-center gap-2">
                              <span className="font-medium text-red-600 line-through">
                                {correction.detected}
                              </span>
                              <span className="text-gray-400">→</span>
                              <span className="font-medium text-green-600">
                                {correction.corrected}
                              </span>
                            </div>
                            {correction.context && (
                              <p className="text-xs text-gray-500 mt-1 italic">{correction.context}</p>
                            )}
                          </div>
                        ))}
                      </div>
                    </div>
                  )}

                  {result.corrections.visual_corrections && result.corrections.visual_corrections.length > 0 && (
                    <div>
                      <h4 className="font-semibold text-gray-900 mb-2 flex items-center gap-2">
                        <span className="text-blue-600">🎨</span>
                        Visual Improvements
                      </h4>
                      <div className="space-y-2">
                        {result.corrections.visual_corrections.map((correction, idx) => (
                          <div
                            key={idx}
                            className="bg-blue-50 border border-blue-200 rounded-lg p-3 text-sm"
                          >
                            <div className="flex items-start justify-between gap-2">
                              <div className="flex-1">
                                <p className="font-medium text-gray-900">{correction.issue}</p>
                                <p className="text-gray-600 mt-1">{correction.suggestion}</p>
                              </div>
                              {correction.priority && (
                                <span className={`text-xs px-2 py-1 rounded ${correction.priority === "high" ? "bg-red-100 text-red-700" :
                                  correction.priority === "medium" ? "bg-yellow-100 text-yellow-700" :
                                    "bg-gray-100 text-gray-700"
                                  }`}>
                                  {correction.priority}
                                </span>
                              )}
                            </div>
                          </div>
                        ))}
                      </div>
                    </div>
                  )}

                  {result.corrections.corrected_prompt && (
                    <div>
                      <h4 className="font-semibold text-gray-900 mb-2">Corrected Prompt</h4>
                      <div className="bg-gray-50 border border-gray-200 rounded-lg p-3 text-sm">
                        <p className="text-gray-700">{result.corrections.corrected_prompt}</p>
                      </div>
                    </div>
                  )}

                  {result.analysis && (
                    <div className="pt-4 border-t border-gray-100">
                      <h4 className="font-bold text-gray-900 mb-3 flex items-center gap-2">
                        <ImageIcon className="h-5 w-5 text-blue-500" />
                        Detailed AI Analysis
                      </h4>
                      <div className="bg-blue-50/50 border border-blue-100 rounded-xl p-4 text-sm whitespace-pre-wrap leading-relaxed text-gray-700 max-h-60 overflow-y-auto custom-scrollbar">
                        {result.analysis}
                      </div>
                    </div>
                  )}

                  {/* Show message if no corrections were made */}
                  {(!result.corrections.spelling_corrections || result.corrections.spelling_corrections.length === 0) &&
                    (!result.corrections.visual_corrections || result.corrections.visual_corrections.length === 0) && (
                      <div className="bg-gray-50 border border-gray-200 rounded-lg p-3 text-sm text-gray-600">
                        <p>No specific corrections were identified. The image was regenerated based on your instructions.</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-end gap-3">
            {step === "input" && (
              <Button
                onClick={handleCorrection}
                variant="primary"
                disabled={!userInstructions && !useAutoAnalyze}
              >
                <Sparkles className="h-4 w-4 mr-2" />
                Start Correction
              </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, keepCorrected);
                }
                onClose();
              }}
              variant={step === "complete" ? "primary" : "secondary"}
            >
              {step === "complete" ? `Use ${keepCorrected ? 'Corrected' : 'Original'}` : "Close"}
            </Button>
          </div>
        </div>
      </div>
    </div>
  );
};