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 { Upload, Download, Eraser, Brush, RotateCcw } from 'lucide-react';
import { fileToBase64 } from '@/lib/utils';

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

export function InpaintingCanvas({
  onSourceImageChange,
  onMaskImageChange,
  className
}: 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 [isDrawing, setIsDrawing] = useState(false);
  const [brushSize, setBrushSize] = useState(20);
  const [isEraser, setIsEraser] = useState(false);
  const [canvasSize, setCanvasSize] = useState({ width: 512, height: 512 });

  // 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 handleImageUpload = async (e: React.ChangeEvent<HTMLInputElement>) => {
    const file = e.target.files?.[0];
    if (!file) return;

    try {
      const base64 = await fileToBase64(file);
      setSourceImage(base64);
      onSourceImageChange(base64);

      // Load image to get dimensions and update canvas size
      const img = new Image();
      img.onload = () => {
        // Calculate scaled dimensions to fit within 512x512 while maintaining aspect ratio
        const maxSize = 512;
        let width = img.width;
        let 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 = base64;
    } catch (err) {
      console.error('Failed to load image:', err);
    }
  };

  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>
              <div className="space-y-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"
                />
              </div>
            </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>
  );
}