stable-diffusion.cpp-rest/webui/components/inpainting-canvas.tsx

'use client';

import { useRef, useEffect, useState, useCallback } from 'react';
import { Button } from '@/components/ui/button';
import { Card, CardContent } from '@/components/ui/card';
import { Label } from '@/components/ui/label';
import { Input } from '@/components/ui/input';
import { Tabs, TabsContent, TabsList, TabsTrigger } from '@/components/ui/tabs';
import { Upload, Download, Eraser, Brush, RotateCcw, Link as LinkIcon, Loader2 } from 'lucide-react';
import { fileToBase64 } from '@/lib/utils';
import { validateImageUrlWithBase64 } from '@/lib/image-validation';
import { apiClient } from '@/lib/api';

interface InpaintingCanvasProps {
  onSourceImageChange: (image: string) => void;
  onMaskImageChange: (image: string) => void;
  className?: string;
  targetWidth?: number;
  targetHeight?: number;
}

export function InpaintingCanvas({
  onSourceImageChange,
  onMaskImageChange,
  className,
  targetWidth,
  targetHeight
}: InpaintingCanvasProps) {
  const canvasRef = useRef<HTMLCanvasElement>(null);
  const maskCanvasRef = useRef<HTMLCanvasElement>(null); // Keep for mask generation
  const fileInputRef = useRef<HTMLInputElement>(null);

  const [sourceImage, setSourceImage] = useState<string | null>(null);
  const [originalSourceImage, setOriginalSourceImage] = useState<string | null>(null);
  const [isDrawing, setIsDrawing] = useState(false);
  const [brushSize, setBrushSize] = useState(20);
  const [isEraser, setIsEraser] = useState(false);
  const [canvasSize, setCanvasSize] = useState({ width: 512, height: 512 });
  const [inputMode, setInputMode] = useState<'file' | 'url'>('file');
  const [urlInput, setUrlInput] = useState('');
  const [isLoadingUrl, setIsLoadingUrl] = useState(false);
  const [urlError, setUrlError] = useState<string | null>(null);
  const [isResizing, setIsResizing] = useState(false);

  // Initialize canvases
  useEffect(() => {
    const canvas = canvasRef.current;
    const maskCanvas = maskCanvasRef.current;

    if (!canvas || !maskCanvas) return;

    const ctx = canvas.getContext('2d');
    const maskCtx = maskCanvas.getContext('2d');

    if (!ctx || !maskCtx) return;

    // Set canvas size
    canvas.width = canvasSize.width;
    canvas.height = canvasSize.height;
    maskCanvas.width = canvasSize.width;
    maskCanvas.height = canvasSize.height;

    // Initialize mask canvas with black (no inpainting)
    maskCtx.fillStyle = 'black';
    maskCtx.fillRect(0, 0, canvasSize.width, canvasSize.height);

    // Update mask image
    updateMaskImage();
  }, [canvasSize]);

  const updateMaskImage = useCallback(() => {
    const maskCanvas = maskCanvasRef.current;
    if (!maskCanvas) return;

    const maskDataUrl = maskCanvas.toDataURL();
    onMaskImageChange(maskDataUrl);
  }, [onMaskImageChange]);

  const loadImageToCanvas = useCallback((base64Image: string) => {
    // Store original image for resizing
    setOriginalSourceImage(base64Image);
    setSourceImage(base64Image);
    onSourceImageChange(base64Image);

    // Load image to get dimensions and update canvas size
    const img = new Image();
    img.onload = () => {
      // Use target dimensions if provided, otherwise fit within 512x512
      let width: number;
      let height: number;

      if (targetWidth && targetHeight) {
        width = targetWidth;
        height = targetHeight;
      } else {
        // Calculate scaled dimensions to fit within 512x512 while maintaining aspect ratio
        const maxSize = 512;
        width = img.width;
        height = img.height;

        if (width > maxSize || height > maxSize) {
          const aspectRatio = width / height;
          if (width > height) {
            width = maxSize;
            height = maxSize / aspectRatio;
          } else {
            height = maxSize;
            width = maxSize * aspectRatio;
          }
        }
      }

      const newCanvasSize = { width: Math.round(width), height: Math.round(height) };
      setCanvasSize(newCanvasSize);

      // Draw image on main canvas
      const canvas = canvasRef.current;
      if (!canvas) return;

      const ctx = canvas.getContext('2d');
      if (!ctx) return;

      canvas.width = width;
      canvas.height = height;
      ctx.drawImage(img, 0, 0, width, height);

      // Update mask canvas size
      const maskCanvas = maskCanvasRef.current;
      if (!maskCanvas) return;

      const maskCtx = maskCanvas.getContext('2d');
      if (!maskCtx) return;

      maskCanvas.width = width;
      maskCanvas.height = height;
      maskCtx.fillStyle = 'black';
      maskCtx.fillRect(0, 0, width, height);

      updateMaskImage();
    };
    img.src = base64Image;
  }, [onSourceImageChange, updateMaskImage, targetWidth, targetHeight]);

  // Auto-resize image when target dimensions change
  useEffect(() => {
    const resizeImage = async () => {
      if (!originalSourceImage || !targetWidth || !targetHeight) {
        return;
      }

      // Don't resize if we're already resizing
      if (isResizing) {
        return;
      }

      try {
        setIsResizing(true);
        const result = await apiClient.resizeImage(originalSourceImage, targetWidth, targetHeight);
        loadImageToCanvas(result.image);
      } catch (err) {
        console.error('Failed to resize image:', err);
      } finally {
        setIsResizing(false);
      }
    };

    resizeImage();
  }, [targetWidth, targetHeight, originalSourceImage]);

  const handleImageUpload = async (e: React.ChangeEvent<HTMLInputElement>) => {
    const file = e.target.files?.[0];
    if (!file) return;

    try {
      const base64 = await fileToBase64(file);
      loadImageToCanvas(base64);
    } catch (err) {
      console.error('Failed to load image:', err);
    }
  };

  const handleUrlLoad = async () => {
    if (!urlInput.trim()) {
      setUrlError('Please enter a URL');
      return;
    }

    setIsLoadingUrl(true);
    setUrlError(null);

    try {
      const result = await validateImageUrlWithBase64(urlInput);

      if (!result.isValid) {
        setUrlError(result.error || 'Failed to load image from URL');
        setIsLoadingUrl(false);
        return;
      }

      // Use base64 data if available, otherwise use the URL directly
      const imageData = result.base64Data || urlInput;
      loadImageToCanvas(imageData);
      setIsLoadingUrl(false);
    } catch (err) {
      setUrlError(err instanceof Error ? err.message : 'Failed to load image from URL');
      setIsLoadingUrl(false);
    }
  };

  const startDrawing = (e: React.MouseEvent<HTMLCanvasElement>) => {
    if (!sourceImage) return;

    setIsDrawing(true);
    draw(e);
  };

  const stopDrawing = () => {
    setIsDrawing(false);
  };

  const draw = (e: React.MouseEvent<HTMLCanvasElement>) => {
    if (!isDrawing || !sourceImage) return;

    const canvas = canvasRef.current;
    const maskCanvas = maskCanvasRef.current;
    if (!canvas || !maskCanvas) return;

    const ctx = canvas.getContext('2d');
    const maskCtx = maskCanvas.getContext('2d');
    if (!ctx || !maskCtx) return;

    const rect = canvas.getBoundingClientRect();
    const scaleX = canvas.width / rect.width;
    const scaleY = canvas.height / rect.height;

    const x = (e.clientX - rect.left) * scaleX;
    const y = (e.clientY - rect.top) * scaleY;

    // Draw on mask canvas (for API)
    maskCtx.globalCompositeOperation = 'source-over';
    maskCtx.fillStyle = isEraser ? 'black' : 'white';
    maskCtx.beginPath();
    maskCtx.arc(x, y, brushSize, 0, Math.PI * 2);
    maskCtx.fill();

    // Draw visual overlay directly on main canvas
    ctx.save();
    ctx.globalCompositeOperation = 'source-over';

    if (isEraser) {
      // For eraser, just redraw the image at that position
      const img = new Image();
      img.onload = () => {
        // Clear the area and redraw
        ctx.save();
        ctx.globalCompositeOperation = 'destination-out';
        ctx.beginPath();
        ctx.arc(x, y, brushSize, 0, Math.PI * 2);
        ctx.fill();
        ctx.restore();

        // Redraw the image in the cleared area
        ctx.save();
        ctx.globalCompositeOperation = 'destination-over';
        ctx.drawImage(img, 0, 0, canvas.width, canvas.height);
        ctx.restore();

        updateMaskImage();
      };
      img.src = sourceImage;
    } else {
      // For brush, draw a visible overlay
      ctx.globalAlpha = 0.6;
      ctx.fillStyle = 'rgba(255, 105, 180, 0.8)'; // Bright pink for visibility
      ctx.beginPath();
      ctx.arc(x, y, brushSize, 0, Math.PI * 2);
      ctx.fill();

      // Also draw a border for better visibility
      ctx.globalAlpha = 1.0;
      ctx.strokeStyle = 'rgba(255, 0, 0, 0.9)'; // Red border
      ctx.lineWidth = 2;
      ctx.beginPath();
      ctx.arc(x, y, brushSize, 0, Math.PI * 2);
      ctx.stroke();

      ctx.restore();
      updateMaskImage();
    }
  };

  const clearMask = () => {
    const canvas = canvasRef.current;
    const maskCanvas = maskCanvasRef.current;
    if (!canvas || !maskCanvas) return;

    const ctx = canvas.getContext('2d');
    const maskCtx = maskCanvas.getContext('2d');
    if (!ctx || !maskCtx) return;

    // Clear mask canvas
    maskCtx.fillStyle = 'black';
    maskCtx.fillRect(0, 0, maskCanvas.width, maskCanvas.height);

    // Redraw source image on main canvas
    if (sourceImage) {
      const img = new Image();
      img.onload = () => {
        ctx.drawImage(img, 0, 0, canvas.width, canvas.height);
      };
      img.src = sourceImage;
    }

    updateMaskImage();
  };

  const downloadMask = () => {
    const canvas = maskCanvasRef.current;
    if (!canvas) return;

    const link = document.createElement('a');
    link.download = 'inpainting-mask.png';
    link.href = canvas.toDataURL();
    link.click();
  };

  return (
    <div className={`space-y-4 ${className}`}>
      <Card>
        <CardContent className="pt-6">
          <div className="space-y-4">
            <div className="space-y-2">
              <Label>Source Image</Label>
              <Tabs value={inputMode} onValueChange={(value) => setInputMode(value as 'file' | 'url')}>
                <TabsList className="grid w-full grid-cols-2">
                  <TabsTrigger value="file">
                    <Upload className="w-4 h-4 mr-2" />
                    Upload File
                  </TabsTrigger>
                  <TabsTrigger value="url">
                    <LinkIcon className="w-4 h-4 mr-2" />
                    From URL
                  </TabsTrigger>
                </TabsList>

                <TabsContent value="file" className="space-y-4 mt-4">
                  <Button
                    type="button"
                    variant="outline"
                    onClick={() => fileInputRef.current?.click()}
                    className="w-full"
                  >
                    <Upload className="h-4 w-4 mr-2" />
                    {sourceImage ? 'Change Image' : 'Upload Image'}
                  </Button>
                  <input
                    ref={fileInputRef}
                    type="file"
                    accept="image/*"
                    onChange={handleImageUpload}
                    className="hidden"
                  />
                </TabsContent>

                <TabsContent value="url" className="space-y-4 mt-4">
                  <div className="space-y-2">
                    <Input
                      type="url"
                      value={urlInput}
                      onChange={(e) => {
                        setUrlInput(e.target.value);
                        setUrlError(null);
                      }}
                      placeholder="https://example.com/image.png"
                      disabled={isLoadingUrl}
                    />
                    <p className="text-xs text-muted-foreground">
                      Enter a URL that ends with an image extension (.jpg, .png, .gif, etc.)
                    </p>
                  </div>
                  <Button
                    type="button"
                    variant="outline"
                    onClick={handleUrlLoad}
                    disabled={isLoadingUrl || !urlInput.trim()}
                    className="w-full"
                  >
                    {isLoadingUrl ? (
                      <>
                        <Download className="h-4 w-4 mr-2 animate-spin" />
                        Loading...
                      </>
                    ) : (
                      <>
                        <Download className="h-4 w-4 mr-2" />
                        Load from URL
                      </>
                    )}
                  </Button>
                  {urlError && (
                    <p className="text-sm text-destructive">{urlError}</p>
                  )}
                </TabsContent>
              </Tabs>
            </div>

            {isResizing && (
              <div className="text-sm text-muted-foreground flex items-center gap-2">
                <Loader2 className="h-4 w-4 animate-spin" />
                Resizing image...
              </div>
            )}

            {sourceImage && (
              <>
                <div className="space-y-2">
                  <Label>Mask Editor</Label>
                  <div className="relative">
                    <canvas
                      ref={canvasRef}
                      className="rounded-lg border border-border cursor-crosshair"
                      style={{ maxWidth: '512px', height: 'auto' }}
                      onMouseDown={startDrawing}
                      onMouseUp={stopDrawing}
                      onMouseMove={draw}
                      onMouseLeave={stopDrawing}
                    />
                  </div>
                  <p className="text-xs text-muted-foreground">
                    Draw on the image to mark areas for inpainting. White areas will be inpainted, black areas will be preserved.
                  </p>
                </div>

                <div className="space-y-2">
                  <Label htmlFor="brush_size">
                    Brush Size: {brushSize}px
                  </Label>
                  <input
                    id="brush_size"
                    type="range"
                    value={brushSize}
                    onChange={(e) => setBrushSize(Number(e.target.value))}
                    min={1}
                    max={100}
                    className="w-full"
                  />
                </div>

                <div className="flex gap-2">
                  <Button
                    type="button"
                    variant={isEraser ? "default" : "outline"}
                    onClick={() => setIsEraser(true)}
                    className="flex-1"
                  >
                    <Eraser className="h-4 w-4 mr-2" />
                    Eraser
                  </Button>
                  <Button
                    type="button"
                    variant={!isEraser ? "default" : "outline"}
                    onClick={() => setIsEraser(false)}
                    className="flex-1"
                  >
                    <Brush className="h-4 w-4 mr-2" />
                    Brush
                  </Button>
                </div>

                <div className="flex gap-2">
                  <Button
                    type="button"
                    variant="outline"
                    onClick={clearMask}
                    className="flex-1"
                  >
                    <RotateCcw className="h-4 w-4 mr-2" />
                    Clear Mask
                  </Button>
                  <Button
                    type="button"
                    variant="outline"
                    onClick={downloadMask}
                    className="flex-1"
                  >
                    <Download className="h-4 w-4 mr-2" />
                    Download Mask
                  </Button>
                </div>
              </>
            )}
          </div>
        </CardContent>
      </Card>
    </div>
  );
}