update

Home.tsx

@@ -1,133 +1,79 @@
 /**
  * @license
  * SPDX-License-Identifier: Apache-2.0
- */
+*/
 /* tslint:disable */
+import {Content, GoogleGenAI, Modality} from '@google/genai';
 import {
-  ChevronDown,
   Library,
   LoaderCircle,
   Paintbrush,
   PictureInPicture,
-  Redo2,
   SendHorizontal,
   Sparkles,
   Trash2,
-  Undo2,
   X,
 } from 'lucide-react';
 import {useEffect, useRef, useState} from 'react';
 
-// This function remains useful for parsing potential error messages
 function parseError(error: string) {
+  const regex = /{"error":(.*)}/gm;
+  const m = regex.exec(error);
   try {
-    // Attempt to parse the error as a JSON object which the proxy might send
-    const errObj = JSON.parse(error);
-    return errObj.message || error;
+    const e = m[1];
+    const err = JSON.parse(e);
+    return err.message || error;
   } catch (e) {
-    // If it's not JSON, return the original error string
-    const regex = /{"error":(.*)}/gm;
-    const m = regex.exec(error);
-    try {
-      const e = m[1];
-      const err = JSON.parse(e);
-      return err.message || error;
-    } catch (e) {
-      return error;
-    }
+    return error;
   }
 }
 
 export default function Home() {
-  const canvasRef = useRef<HTMLCanvasElement>(null);
-  const fileInputRef = useRef<HTMLInputElement>(null);
-  const backgroundImageRef = useRef<HTMLImageElement | null>(null);
-  const dropdownRef = useRef<HTMLDivElement>(null);
+  const canvasRef = useRef(null);
+  const fileInputRef = useRef(null);
+  const backgroundImageRef = useRef(null);
   const [isDrawing, setIsDrawing] = useState(false);
   const [prompt, setPrompt] = useState('');
   const [generatedImage, setGeneratedImage] = useState<string | null>(null);
-  const [multiImages, setMultiImages] = useState<
-    {url: string; type: string}[]
-  >([]);
+  const [multiImages, setMultiImages] = useState<string[]>([]);
   const [isLoading, setIsLoading] = useState(false);
   const [showErrorModal, setShowErrorModal] = useState(false);
   const [errorMessage, setErrorMessage] = useState('');
   const [mode, setMode] = useState<
     'canvas' | 'editor' | 'imageGen' | 'multi-img-edit'
   >('editor');
-  const [isDropdownOpen, setIsDropdownOpen] = useState(false);
+
   const [apiKey, setApiKey] = useState('');
   const [showApiKeyModal, setShowApiKeyModal] = useState(false);
-
-  // State for canvas history
-  const [history, setHistory] = useState<string[]>([]);
-  const [historyIndex, setHistoryIndex] = useState(-1);
-
-  // When switching to canvas mode, initialize it and its history
-  useEffect(() => {
-    if (mode === 'canvas' && canvasRef.current) {
-      const canvas = canvasRef.current;
-      const ctx = canvas.getContext('2d');
-      ctx.fillStyle = '#FFFFFF';
-      ctx.fillRect(0, 0, canvas.width, canvas.height);
-
-      // If an image already exists from another mode, draw it.
-      if (generatedImage) {
-        const img = new window.Image();
-        img.onload = () => {
-          ctx.drawImage(img, 0, 0, canvas.width, canvas.height);
-          // Save this as the initial state for this session
-          const dataUrl = canvas.toDataURL();
-          setHistory([dataUrl]);
-          setHistoryIndex(0);
-        };
-        img.src = generatedImage;
-      } else {
-        // Otherwise, save the blank state as initial
-        const dataUrl = canvas.toDataURL();
-        setHistory([dataUrl]);
-        setHistoryIndex(0);
-      }
-    }
-  }, [mode, generatedImage]);
+  const [tempApiKey, setTempApiKey] = useState('');
+  const submissionRef = useRef<(() => Promise<void>) | null>(null);
 
   // Load background image when generatedImage changes
   useEffect(() => {
     if (generatedImage && canvasRef.current) {
+      // Use the window.Image constructor to avoid conflict with Next.js Image component
       const img = new window.Image();
       img.onload = () => {
         backgroundImageRef.current = img;
         drawImageToCanvas();
-        if (mode === 'canvas') {
-          // A small timeout to let the draw happen before saving
-          setTimeout(saveCanvasState, 50);
-        }
       };
       img.src = generatedImage;
     }
-  }, [generatedImage, mode]);
+  }, [generatedImage]);
 
-  // Handle clicks outside the dropdown to close it
+  // Initialize canvas with white background when component mounts
   useEffect(() => {
-    function handleClickOutside(event: MouseEvent) {
-      if (
-        dropdownRef.current &&
-        !dropdownRef.current.contains(event.target as Node)
-      ) {
-        setIsDropdownOpen(false);
-      }
+    if (canvasRef.current) {
+      initializeCanvas();
     }
-    document.addEventListener('mousedown', handleClickOutside);
-    return () => {
-      document.removeEventListener('mousedown', handleClickOutside);
-    };
-  }, [dropdownRef]);
+  }, []);
 
   // Initialize canvas with white background
   const initializeCanvas = () => {
     const canvas = canvasRef.current;
-    if (!canvas) return;
     const ctx = canvas.getContext('2d');
+
+    // Fill canvas with white background
     ctx.fillStyle = '#FFFFFF';
     ctx.fillRect(0, 0, canvas.width, canvas.height);
   };
@@ -138,8 +84,12 @@ export default function Home() {
 
     const canvas = canvasRef.current;
     const ctx = canvas.getContext('2d');
+
+    // Fill with white background first
     ctx.fillStyle = '#FFFFFF';
     ctx.fillRect(0, 0, canvas.width, canvas.height);
+
+    // Draw the background image
     ctx.drawImage(
       backgroundImageRef.current,
       0,
@@ -149,52 +99,16 @@ export default function Home() {
     );
   };
 
-  // Canvas history functions
-  const saveCanvasState = () => {
-    if (!canvasRef.current) return;
-    const canvas = canvasRef.current;
-    const dataUrl = canvas.toDataURL();
-    const newHistory = history.slice(0, historyIndex + 1);
-    newHistory.push(dataUrl);
-    setHistory(newHistory);
-    setHistoryIndex(newHistory.length - 1);
-  };
-
-  const restoreCanvasState = (index: number) => {
-    if (!canvasRef.current || !history[index]) return;
-    const canvas = canvasRef.current;
-    const ctx = canvas.getContext('2d');
-    const dataUrl = history[index];
-    const img = new window.Image();
-    img.onload = () => {
-      ctx.clearRect(0, 0, canvas.width, canvas.height);
-      ctx.drawImage(img, 0, 0, canvas.width, canvas.height);
-    };
-    img.src = dataUrl;
-  };
-
-  const handleUndo = () => {
-    if (historyIndex > 0) {
-      const newIndex = historyIndex - 1;
-      setHistoryIndex(newIndex);
-      restoreCanvasState(newIndex);
-    }
-  };
-
-  const handleRedo = () => {
-    if (historyIndex < history.length - 1) {
-      const newIndex = historyIndex + 1;
-      setHistoryIndex(newIndex);
-      restoreCanvasState(newIndex);
-    }
-  };
-
   // Get the correct coordinates based on canvas scaling
-  const getCoordinates = (e: any) => {
-    const canvas = canvasRef.current!;
+  const getCoordinates = (e) => {
+    const canvas = canvasRef.current;
     const rect = canvas.getBoundingClientRect();
+
+    // Calculate the scaling factor between the internal canvas size and displayed size
     const scaleX = canvas.width / rect.width;
     const scaleY = canvas.height / rect.height;
+
+    // Apply the scaling to get accurate coordinates
     return {
       x:
         (e.nativeEvent.offsetX ||
@@ -205,26 +119,34 @@ export default function Home() {
     };
   };
 
-  const startDrawing = (e: any) => {
-    const canvas = canvasRef.current!;
-    const ctx = canvas.getContext('2d')!;
+  const startDrawing = (e) => {
+    const canvas = canvasRef.current;
+    const ctx = canvas.getContext('2d');
     const {x, y} = getCoordinates(e);
+
+    // Prevent default behavior to avoid scrolling on touch devices
     if (e.type === 'touchstart') {
       e.preventDefault();
     }
+
+    // Start a new path without clearing the canvas
     ctx.beginPath();
     ctx.moveTo(x, y);
     setIsDrawing(true);
   };
 
-  const draw = (e: any) => {
+  const draw = (e) => {
     if (!isDrawing) return;
+
+    // Prevent default behavior to avoid scrolling on touch devices
     if (e.type === 'touchmove') {
       e.preventDefault();
     }
-    const canvas = canvasRef.current!;
-    const ctx = canvas.getContext('2d')!;
+
+    const canvas = canvasRef.current;
+    const ctx = canvas.getContext('2d');
     const {x, y} = getCoordinates(e);
+
     ctx.lineWidth = 5;
     ctx.lineCap = 'round';
     ctx.strokeStyle = '#000000';
@@ -233,22 +155,21 @@ export default function Home() {
   };
 
   const stopDrawing = () => {
-    if (!isDrawing) return;
     setIsDrawing(false);
-    saveCanvasState();
   };
 
   const handleClear = () => {
+    // For canvas, we must clear the actual drawing strokes
     if (mode === 'canvas' && canvasRef.current) {
-      initializeCanvas();
-      const dataUrl = canvasRef.current.toDataURL();
-      setHistory([dataUrl]);
-      setHistoryIndex(0);
+      const canvas = canvasRef.current;
+      const ctx = canvas.getContext('2d');
+      ctx.fillStyle = '#FFFFFF';
+      ctx.fillRect(0, 0, canvas.width, canvas.height);
     }
+    // For all modes, we clear the background/uploaded/generated image state
     setGeneratedImage(null);
     setMultiImages([]);
     backgroundImageRef.current = null;
-    setPrompt('');
   };
 
   const processFiles = (files: FileList | null) => {
@@ -258,16 +179,11 @@ export default function Home() {
     );
     if (fileArray.length === 0) return;
 
-    if (!apiKey) {
-      setShowApiKeyModal(true);
-    }
-
     if (mode === 'multi-img-edit') {
       const readers = fileArray.map((file) => {
-        return new Promise<{url: string; type: string}>((resolve, reject) => {
+        return new Promise<string>((resolve, reject) => {
           const reader = new FileReader();
-          reader.onload = () =>
-            resolve({url: reader.result as string, type: file.type});
+          reader.onload = () => resolve(reader.result as string);
           reader.onerror = reject;
           reader.readAsDataURL(file);
         });
@@ -276,18 +192,17 @@ export default function Home() {
         setMultiImages((prev) => [...prev, ...newImages]);
       });
     } else {
-      const file = fileArray[0];
+      // single file modes
       const reader = new FileReader();
       reader.onload = () => {
         setGeneratedImage(reader.result as string);
       };
-      reader.readAsDataURL(file);
+      reader.readAsDataURL(fileArray[0]);
     }
   };
 
   const handleFileChange = (e: React.ChangeEvent<HTMLInputElement>) => {
     processFiles(e.target.files);
-    e.target.value = '';
   };
 
   const handleDrop = (e: React.DragEvent<HTMLDivElement>) => {
@@ -314,149 +229,165 @@ export default function Home() {
     );
   };
 
-  // *** MODIFIED FUNCTION ***
-  const handleSubmit = async (e: React.FormEvent) => {
+  const handleApiKeySubmit = async (e: React.FormEvent) => {
     e.preventDefault();
-
-    if (!apiKey) {
-      setShowApiKeyModal(true);
-      return;
+    if (tempApiKey) {
+      setErrorMessage(''); // Clear previous error message
+      setApiKey(tempApiKey);
+      setShowApiKeyModal(false);
+      setTempApiKey('');
+      if (submissionRef.current) {
+        await submissionRef.current();
+        submissionRef.current = null;
+      }
     }
+  };
 
-    setIsLoading(true);
+  const handleSubmit = async (e) => {
+    e.preventDefault();
 
-    try {
-      if (mode === 'editor' && !generatedImage) {
-        setErrorMessage('Please upload an image to edit.');
-        setShowErrorModal(true);
-        return;
-      }
+    // Pre-submission checks that don't require an API key
+    if (mode === 'editor' && !generatedImage) {
+      setErrorMessage('Please upload an image to edit.');
+      setShowErrorModal(true);
+      return;
+    }
 
-      if (mode === 'multi-img-edit' && multiImages.length === 0) {
-        setErrorMessage('Please upload at least one image to edit.');
-        setShowErrorModal(true);
-        return;
-      }
+    if (mode === 'multi-img-edit' && multiImages.length === 0) {
+      setErrorMessage('Please upload at least one image to edit.');
+      setShowErrorModal(true);
+      return;
+    }
 
-      const parts: any[] = [];
-
-      // This logic for building the 'parts' array is correct.
-      if (mode === 'imageGen') {
-        const tempCanvas = document.createElement('canvas');
-        tempCanvas.width = 960;
-        tempCanvas.height = 540;
-        const tempCtx = tempCanvas.getContext('2d')!;
-        tempCtx.fillStyle = '#FFFFFF';
-        tempCtx.fillRect(0, 0, tempCanvas.width, tempCanvas.height);
-        tempCtx.fillStyle = '#FEFEFE';
-        tempCtx.fillRect(0, 0, 1, 1);
-        const imageB64 = tempCanvas.toDataURL('image/png').split(',')[1];
-        parts.push({inlineData: {data: imageB64, mimeType: 'image/png'}});
-      } else if (mode === 'canvas') {
-        if (!canvasRef.current) return;
-        const canvas = canvasRef.current;
-        const imageB64 = canvas.toDataURL('image/png').split(',')[1];
-        parts.push({inlineData: {data: imageB64, mimeType: 'image/png'}});
-      } else if (mode === 'editor' && generatedImage) {
-        const mimeType = generatedImage.substring(
-          generatedImage.indexOf(':') + 1,
-          generatedImage.indexOf(';'),
-        );
-        const imageB64 = generatedImage.split(',')[1];
-        parts.push({inlineData: {data: imageB64, mimeType}});
-      } else if (mode === 'multi-img-edit') {
-        multiImages.forEach((img) => {
-          parts.push({
-            inlineData: {data: img.url.split(',')[1], mimeType: img.type},
+    const submitAction = async () => {
+      setIsLoading(true);
+      try {
+        const ai = new GoogleGenAI({apiKey});
+
+        if (mode === 'imageGen') {
+          const response = await ai.models.generateImages({
+            model: 'imagen-4.0-generate-001',
+            prompt: prompt,
+            config: {
+              numberOfImages: 1,
+            },
           });
-        });
-      }
-
-      parts.push({text: prompt});
 
-      // Construct the request body for the Gemini REST API
-      const requestBody = {
-        contents: [{role: 'USER', parts}],
-      };
-
-      // Define the proxy endpoint
-      const proxyUrl = `/api-proxy/v1beta/models/gemini-2.5-flash-image-preview:generateContent?key=${apiKey}`;
+          const base64ImageBytes: string =
+            response.generatedImages[0].image.imageBytes;
+          const imageUrl = `data:image/png;base64,${base64ImageBytes}`;
+          setGeneratedImage(imageUrl);
+        } else {
+          // Editor, Canvas, and Multi-Img modes
+          const parts: any[] = [];
+          if (mode === 'canvas') {
+            if (!canvasRef.current) return;
+            const canvas = canvasRef.current;
+            const tempCanvas = document.createElement('canvas');
+            tempCanvas.width = canvas.width;
+            tempCanvas.height = canvas.height;
+            const tempCtx = tempCanvas.getContext('2d');
+            tempCtx.fillStyle = '#FFFFFF';
+            tempCtx.fillRect(0, 0, tempCanvas.width, tempCanvas.height);
+            tempCtx.drawImage(canvas, 0, 0);
+            const imageB64 = tempCanvas.toDataURL('image/png').split(',')[1];
+            parts.push({inlineData: {data: imageB64, mimeType: 'image/png'}});
+          } else if (mode === 'editor') {
+            const imageB64 = generatedImage.split(',')[1];
+            parts.push({inlineData: {data: imageB64, mimeType: 'image/png'}});
+          } else if (mode === 'multi-img-edit') {
+            multiImages.forEach((img) => {
+              parts.push({
+                inlineData: {data: img.split(',')[1], mimeType: 'image/png'},
+              });
+            });
+          }
 
-      // Use fetch to send the request to your proxy server
-      const response = await fetch(proxyUrl, {
-        method: 'POST',
-        headers: {
-          'Content-Type': 'application/json',
-        },
-        body: JSON.stringify(requestBody),
-      });
+          parts.push({text: prompt});
 
-      if (!response.ok) {
-        const errorData = await response.json();
-        throw new Error(
-          errorData.error?.message || `HTTP error! status: ${response.status}`,
-        );
-      }
+          const contents: Content[] = [{role: 'USER', parts}];
 
-      const responseData = await response.json();
+          const response = await ai.models.generateContent({
+            model: 'gemini-2.5-flash-image-preview',
+            contents,
+            config: {
+              responseModalities: [Modality.TEXT, Modality.IMAGE],
+            },
+          });
 
-      // Process the response
-      const result = {message: '', imageData: null};
+          const data = {
+            success: true,
+            message: '',
+            imageData: null,
+            error: undefined,
+          };
+          for (const part of response.candidates[0].content.parts) {
+            if (part.text) {
+              data.message = part.text;
+            } else if (part.inlineData) {
+              data.imageData = part.inlineData.data;
+            }
+          }
 
-      if (responseData.candidates && responseData.candidates.length > 0) {
-        for (const part of responseData.candidates[0].content.parts) {
-          if (part.text) {
-            result.message = part.text;
-          } else if (part.inlineData) {
-            result.imageData = part.inlineData.data;
+          if (data.imageData) {
+            const imageUrl = `data:image/png;base64,${data.imageData}`;
+            if (mode === 'multi-img-edit') {
+              setGeneratedImage(imageUrl);
+              setMultiImages([]);
+              setMode('editor');
+            } else {
+              setGeneratedImage(imageUrl);
+            }
+          } else {
+            setErrorMessage(
+              data.message || 'Failed to generate image. Please try again.',
+            );
+            setShowErrorModal(true);
           }
         }
-      } else {
-        throw new Error('Invalid response structure from API.');
-      }
-
-      if (result.imageData) {
-        const imageUrl = `data:image/png;base64,${result.imageData}`;
-        if (mode === 'multi-img-edit') {
-          setGeneratedImage(imageUrl);
-          setMultiImages([]);
-          setMode('editor');
+      } catch (error) {
+        console.error('Error submitting:', error);
+        const parsedError = parseError(error.message);
+        if (
+          parsedError &&
+          (parsedError.includes('API_KEY_INVALID') ||
+            parsedError.includes('API key not valid'))
+        ) {
+          setErrorMessage(
+            'Your API key is not valid. Please enter a valid key to continue.',
+          );
+          setApiKey(''); // Clear the invalid key
+          setShowApiKeyModal(true); // Re-show the modal
+          // submissionRef is not cleared, so user can retry.
         } else {
-          setGeneratedImage(imageUrl);
+          setErrorMessage(parsedError || 'An unexpected error occurred.');
+          setShowErrorModal(true);
         }
-      } else {
-        setErrorMessage(
-          result.message || 'Failed to generate image. Please try again.',
-        );
-        setShowErrorModal(true);
+      } finally {
+        setIsLoading(false);
       }
-    } catch (error: any) {
-      console.error('Error submitting:', error);
-      setErrorMessage(error.message || 'An unexpected error occurred.');
-      setShowErrorModal(true);
-    } finally {
-      setIsLoading(false);
+    };
+
+    if (!apiKey) {
+      submissionRef.current = submitAction;
+      setShowApiKeyModal(true);
+    } else {
+      await submitAction();
     }
   };
 
+  // Close the error modal
   const closeErrorModal = () => {
     setShowErrorModal(false);
   };
 
-  const handleApiKeySubmit = (e: React.FormEvent) => {
-    e.preventDefault();
-    const newApiKey = (e.target as any).apiKey.value;
-    if (newApiKey) {
-      setApiKey(newApiKey);
-      setShowApiKeyModal(false);
-    }
-  };
-
+  // Add touch event prevention function
   useEffect(() => {
     const canvas = canvasRef.current;
     if (!canvas) return;
 
-    const preventTouchDefault = (e: TouchEvent) => {
+    // Function to prevent default touch behavior on canvas
+    const preventTouchDefault = (e) => {
       if (isDrawing) {
         e.preventDefault();
       }
@@ -469,6 +400,7 @@ export default function Home() {
       passive: false,
     });
 
+    // Remove event listener when component unmounts
     return () => {
       canvas.removeEventListener('touchstart', preventTouchDefault);
       canvas.removeEventListener('touchmove', preventTouchDefault);
@@ -476,15 +408,16 @@ export default function Home() {
   }, [isDrawing]);
 
   const baseDisplayClass =
-    'w-full sm:h-[60vh] h-[40vh] min-h-[320px] bg-white/90 touch-none flex items-center justify-center p-4 transition-colors';
+    'w-full sm:h-[60vh] h-[30vh] min-h-[320px] bg-white/90 touch-none flex items-center justify-center p-4 transition-colors';
 
   return (
     <>
-      <div className="min-h-screen text-gray-900 flex flex-col justify-start items-center">
+      <div className="min-h-screen notebook-paper-bg text-gray-900 flex flex-col justify-start items-center">
         <main className="container mx-auto px-3 sm:px-6 py-5 sm:py-10 pb-32 max-w-5xl w-full">
+          {/* Header section with title and tools */}
           <div className="flex flex-col sm:flex-row sm:justify-between sm:items-end mb-2 sm:mb-6 gap-2">
             <div>
-              <h1 className="text-2xl sm:text-3xl font-bold mb-0 leading-tight">
+              <h1 className="text-2xl sm:text-3xl font-bold mb-0 leading-tight font-mega">
                 Nano Banana AIO
               </h1>
               <p className="text-sm sm:text-base text-gray-500 mt-1">
@@ -499,7 +432,7 @@ export default function Home() {
                 by{' '}
                 <a
                   className="underline"
-                  href="https://huggingface.co/prithivMLmods"
+                  href="https://www.linkedin.com/in/prithiv-sakthi/"
                   target="_blank"
                   rel="noopener noreferrer">
                   prithivsakthi-ur
@@ -508,85 +441,46 @@ export default function Home() {
             </div>
 
             <menu className="flex items-center bg-gray-300 rounded-full p-2 shadow-sm self-start sm:self-auto">
-              <div className="flex flex-wrap justify-center items-center bg-gray-200/80 rounded-full p-1 mr-2">
-                <div className="relative" ref={dropdownRef}>
-                  <button
-                    onClick={() => setIsDropdownOpen(!isDropdownOpen)}
-                    className={`p-2 sm:px-3 sm:py-1.5 rounded-full text-sm font-medium flex items-center gap-2 transition-colors ${
-                      mode === 'editor' || mode === 'multi-img-edit'
-                        ? 'bg-white shadow'
-                        : 'text-gray-600 hover:bg-gray-300/50'
-                    }`}
-                    aria-haspopup="true"
-                    aria-expanded={isDropdownOpen}>
-                    {mode === 'multi-img-edit' ? (
-                      <>
-                        <Library className="w-4 h-4" />
-                        <span className="hidden sm:inline">Multi-Image</span>
-                      </>
-                    ) : (
-                      <>
-                        <PictureInPicture className="w-4 h-4" />
-                        <span className="hidden sm:inline">Editor</span>
-                      </>
-                    )}
-                    <ChevronDown className="w-4 h-4 opacity-70" />
-                  </button>
-                  {isDropdownOpen && (
-                    <div className="absolute top-full mt-2 w-48 bg-white rounded-lg shadow-xl z-10 border border-gray-200 py-1">
-                      <button
-                        onClick={() => {
-                          setMode('editor');
-                          setIsDropdownOpen(false);
-                        }}
-                        className={`w-full text-left px-3 py-2 text-sm flex items-center gap-3 transition-colors ${
-                          mode === 'editor'
-                            ? 'bg-gray-100 text-gray-900'
-                            : 'text-gray-700 hover:bg-gray-50'
-                        }`}
-                        aria-pressed={mode === 'editor'}>
-                        <PictureInPicture className="w-4 h-4" />
-                        <span>Single Image Edit</span>
-                      </button>
-                      <button
-                        onClick={() => {
-                          setMode('multi-img-edit');
-                          setIsDropdownOpen(false);
-                        }}
-                        className={`w-full text-left px-3 py-2 text-sm flex items-center gap-3 transition-colors ${
-                          mode === 'multi-img-edit'
-                            ? 'bg-gray-100 text-gray-900'
-                            : 'text-gray-700 hover:bg-gray-50'
-                        }`}
-                        aria-pressed={mode === 'multi-img-edit'}>
-                        <Library className="w-4 h-4" />
-                        <span>Multi-Image Edit</span>
-                      </button>
-                    </div>
-                  )}
-                </div>
-
+              <div className="flex items-center bg-gray-200/80 rounded-full p-1 mr-2">
+                <button
+                  onClick={() => setMode('editor')}
+                  className={`px-3 py-1.5 rounded-full text-sm font-medium flex items-center gap-2 transition-colors ${
+                    mode === 'editor'
+                      ? 'bg-white shadow'
+                      : 'text-gray-600 hover:bg-gray-300/50'
+                  }`}
+                  aria-pressed={mode === 'editor'}>
+                  <PictureInPicture className="w-4 h-4" /> Editor
+                </button>
+                <button
+                  onClick={() => setMode('multi-img-edit')}
+                  className={`px-3 py-1.5 rounded-full text-sm font-medium flex items-center gap-2 transition-colors ${
+                    mode === 'multi-img-edit'
+                      ? 'bg-white shadow'
+                      : 'text-gray-600 hover:bg-gray-300/50'
+                  }`}
+                  aria-pressed={mode === 'multi-img-edit'}>
+                  <Library className="w-4 h-4" /> Multi-Image
+                </button>
                 <button
                   onClick={() => setMode('canvas')}
-                  className={`p-2 sm:px-3 sm:py-1.5 rounded-full text-sm font-medium flex items-center gap-2 transition-colors ${
+                  className={`px-3 py-1.5 rounded-full text-sm font-medium flex items-center gap-2 transition-colors ${
                     mode === 'canvas'
                       ? 'bg-white shadow'
                       : 'text-gray-600 hover:bg-gray-300/50'
                   }`}
                   aria-pressed={mode === 'canvas'}>
-                  <Paintbrush className="w-4 h-4" />
-                  <span className="hidden sm:inline">Canvas</span>
+                  <Paintbrush className="w-4 h-4" /> Canvas
                 </button>
                 <button
                   onClick={() => setMode('imageGen')}
-                  className={`p-2 sm:px-3 sm:py-1.5 rounded-full text-sm font-medium flex items-center gap-2 transition-colors ${
+                  className={`px-3 py-1.5 rounded-full text-sm font-medium flex items-center gap-2 transition-colors ${
                     mode === 'imageGen'
                       ? 'bg-white shadow'
                       : 'text-gray-600 hover:bg-gray-300/50'
                   }`}
                   aria-pressed={mode === 'imageGen'}>
-                  <Sparkles className="w-4 h-4" />
-                  <span className="hidden sm:inline">Image Gen</span>
+                  <Sparkles className="w-4 h-4" /> Image Gen
                 </button>
               </div>
               <button
@@ -601,6 +495,7 @@ export default function Home() {
             </menu>
           </div>
 
+          {/* Main display section */}
           <div className="w-full mb-6">
             <input
               ref={fileInputRef}
@@ -612,41 +507,23 @@ export default function Home() {
               multiple={mode === 'multi-img-edit'}
             />
             {mode === 'canvas' ? (
-              <div className="relative w-full">
-                <canvas
-                  ref={canvasRef}
-                  width={960}
-                  height={540}
-                  onMouseDown={startDrawing}
-                  onMouseMove={draw}
-                  onMouseUp={stopDrawing}
-                  onMouseLeave={stopDrawing}
-                  onTouchStart={startDrawing}
-                  onTouchMove={draw}
-                  onTouchEnd={stopDrawing}
-                  className="border-2 border-black w-full sm:h-[60vh] h-[40vh] min-h-[320px] bg-white/90 touch-none"
-                  style={{
-                    cursor:
-                      "url('data:image/svg+xml;utf8,<svg xmlns=\"http://www.w3.org/2000/svg\" width=\"24\" height=\"24\" viewBox=\"0 0 24 24\" fill=\"none\" stroke=\"%23FF0000\" stroke-width=\"2\" stroke-linecap=\"round\" stroke-linejoin=\"round\"><path d=\"M12 5v14M5 12h14\"/></svg>') 12 12, crosshair",
-                  }}
-                />
-                <div className="absolute top-2 right-2 flex gap-2">
-                  <button
-                    onClick={handleUndo}
-                    disabled={historyIndex <= 0}
-                    className="p-2 bg-white rounded-md shadow disabled:opacity-50 disabled:cursor-not-allowed hover:bg-gray-100 transition-colors"
-                    aria-label="Undo">
-                    <Undo2 className="w-5 h-5" />
-                  </button>
-                  <button
-                    onClick={handleRedo}
-                    disabled={historyIndex >= history.length - 1}
-                    className="p-2 bg-white rounded-md shadow disabled:opacity-50 disabled:cursor-not-allowed hover:bg-gray-100 transition-colors"
-                    aria-label="Redo">
-                    <Redo2 className="w-5 h-5" />
-                  </button>
-                </div>
-              </div>
+              <canvas
+                ref={canvasRef}
+                width={960}
+                height={540}
+                onMouseDown={startDrawing}
+                onMouseMove={draw}
+                onMouseUp={stopDrawing}
+                onMouseLeave={stopDrawing}
+                onTouchStart={startDrawing}
+                onTouchMove={draw}
+                onTouchEnd={stopDrawing}
+                className="border-2 border-black w-full sm:h-[60vh] h-[30vh] min-h-[320px] bg-white/90 touch-none"
+                style={{
+                  cursor:
+                    "url('data:image/svg+xml;utf8,<svg xmlns=\"http://www.w3.org/2000/svg\" width=\"24\" height=\"24\" viewBox=\"0 0 24 24\" fill=\"none\" stroke=\"%23FF0000\" stroke-width=\"2\" stroke-linecap=\"round\" stroke-linejoin=\"round\"><path d=\"M12 5v14M5 12h14\"/></svg>') 12 12, crosshair",
+                }}
+              />
             ) : mode === 'editor' ? (
               <div
                 className={`${baseDisplayClass} ${
@@ -688,7 +565,7 @@ export default function Home() {
                     {multiImages.map((image, index) => (
                       <div key={index} className="relative group aspect-square">
                         <img
-                          src={image.url}
+                          src={image}
                           alt={`upload preview ${index + 1}`}
                           className="w-full h-full object-cover rounded-md"
                         />
@@ -722,7 +599,7 @@ export default function Home() {
             ) : (
               // Image Gen mode display
               <div
-                className={`relative ${baseDisplayClass} border-2 ${
+                className={`${baseDisplayClass} border-2 ${
                   generatedImage ? 'border-black' : 'border-gray-400'
                 }`}>
                 {generatedImage ? (
@@ -751,11 +628,9 @@ export default function Home() {
                 placeholder={
                   mode === 'imageGen'
                     ? 'Describe the image you want to create...'
-                    : mode === 'multi-img-edit'
-                    ? 'Describe how to edit or combine the images...'
                     : 'Add your change...'
                 }
-                className="w-full p-3 sm:p-4 pr-12 sm:pr-14 text-sm sm:text-base border-2 border-black bg-white text-gray-800 shadow-sm focus:ring-2 focus:ring-gray-200 focus:outline-none transition-all"
+                className="w-full p-3 sm:p-4 pr-12 sm:pr-14 text-sm sm:text-base border-2 border-black bg-white text-gray-800 shadow-sm focus:ring-2 focus:ring-gray-200 focus:outline-none transition-all font-mono"
                 required
               />
               <button
@@ -777,58 +652,71 @@ export default function Home() {
             </div>
           </form>
         </main>
-        {/* Error Modal */}
-        {showErrorModal && (
+        {/* API Key Modal */}
+        {showApiKeyModal && (
           <div className="fixed inset-0 bg-black/50 flex items-center justify-center z-50 p-4">
             <div className="bg-white rounded-lg shadow-xl max-w-md w-full p-6">
-              <div className="flex justify-between items-start mb-4">
-                <h3 className="text-xl font-bold text-gray-700">
-                  Failed to generate
-                </h3>
+              <form onSubmit={handleApiKeySubmit}>
+                <div className="flex justify-between items-start mb-4">
+                  <h3 className="text-xl font-bold text-gray-700">
+                    Add Gemini API Key
+                  </h3>
+                  <button
+                    type="button"
+                    onClick={() => {
+                      setShowApiKeyModal(false);
+                      setErrorMessage('');
+                      submissionRef.current = null; // Cancel submission
+                    }}
+                    className="text-gray-400 hover:text-gray-500"
+                    aria-label="Close">
+                    <X className="w-5 h-5" />
+                  </button>
+                </div>
+                <p className="text-gray-600 mb-4 text-sm">
+                  Add the API key to process the request. The API key will be
+                  removed if the app page is refreshed or closed.
+                </p>
+                {errorMessage && (
+                  <p className="text-red-500 text-sm mb-2 font-medium">
+                    {errorMessage}
+                  </p>
+                )}
+                <input
+                  type="password"
+                  value={tempApiKey}
+                  onChange={(e) => setTempApiKey(e.target.value)}
+                  placeholder="Enter your Gemini API Key"
+                  className="w-full p-2 mb-4 border-2 border-gray-300 rounded focus:ring-2 focus:ring-gray-400 focus:outline-none transition-all"
+                  required
+                  aria-label="Gemini API Key"
+                />
                 <button
-                  onClick={closeErrorModal}
-                  className="text-gray-400 hover:text-gray-500">
-                  <X className="w-5 h-5" />
+                  type="submit"
+                  className="w-full p-2 bg-black text-white rounded hover:bg-gray-800 transition-colors">
+                  Save & Continue
                 </button>
-              </div>
-              <p className="font-medium text-gray-600">
-                {parseError(errorMessage)}
-              </p>
+              </form>
             </div>
           </div>
         )}
-        {/* API Key Modal */}
-        {showApiKeyModal && (
+        {/* Error Modal */}
+        {showErrorModal && (
           <div className="fixed inset-0 bg-black/50 flex items-center justify-center z-50 p-4">
             <div className="bg-white rounded-lg shadow-xl max-w-md w-full p-6">
               <div className="flex justify-between items-start mb-4">
                 <h3 className="text-xl font-bold text-gray-700">
-                  Add Gemini API Key
+                  Failed to generate
                 </h3>
                 <button
-                  onClick={() => setShowApiKeyModal(false)}
+                  onClick={closeErrorModal}
                   className="text-gray-400 hover:text-gray-500">
                   <X className="w-5 h-5" />
                 </button>
               </div>
-              <p className="text-gray-600 mb-4">
-                Add the API key to process the request. The API key will be
-                removed if the app page is refreshed or closed.
+              <p className="font-medium text-gray-600">
+                {parseError(errorMessage)}
               </p>
-              <form onSubmit={handleApiKeySubmit}>
-                <input
-                  type="password"
-                  name="apiKey"
-                  className="w-full p-2 border-2 border-gray-300 rounded-md mb-4"
-                  placeholder="Enter your Gemini API Key"
-                  required
-                />
-                <button
-                  type="submit"
-                  className="w-full bg-black text-white p-2 rounded-md hover:bg-gray-800 transition-colors">
-                  Submit
-                </button>
-              </form>
             </div>
           </div>
         )}