src/WatermarkRemover.jsx

import React, { useState, useRef, useCallback } from 'react';
import { Upload, Download, Settings, Play, Eye, EyeOff, RotateCcw, ChevronLeft, ChevronRight } from 'lucide-react';
import * as pdfjsLib from 'pdfjs-dist';
import pdfjsWorker from 'pdfjs-dist/build/pdf.worker.min.mjs?url';
import { jsPDF } from 'jspdf';

pdfjsLib.GlobalWorkerOptions.workerSrc = pdfjsWorker;

export default function WatermarkRemover() {
  const [images, setImages] = useState([]);
  const [originalFileName, setOriginalFileName] = useState('');
  const [currentPage, setCurrentPage] = useState(0);
  const [processing, setProcessing] = useState(false);
  const [showOriginal, setShowOriginal] = useState(false);
  const [settings, setSettings] = useState({
    width: 220,
    height: 80,
    offsetX: 0,
    offsetY: 0,
    blendRadius: 10,
    method: 'telea' // 'telea' or 'ns' (Navier-Stokes)
  });
  const [showSettings, setShowSettings] = useState(false);
  const canvasRef = useRef(null);
  const fileInputRef = useRef(null);

  // Convert PDF pages to images using PDF.js approach (simplified - uses canvas rendering)
  const handleFileUpload = async (e) => {
    const file = e.target.files[0];
    if (!file) return;

    setProcessing(true);
    const loadedImages = [];

    if (file.type === 'application/pdf') {
      try {
        setOriginalFileName(file.name.replace(/\.pdf$/i, ''));
        const arrayBuffer = await file.arrayBuffer();
        const pdf = await pdfjsLib.getDocument({ data: arrayBuffer }).promise;
        const numPages = pdf.numPages;
        
        for (let pageNum = 1; pageNum <= numPages; pageNum++) {
          const page = await pdf.getPage(pageNum);
          const scale = 1.5;
          const viewport = page.getViewport({ scale });
          
          const canvas = document.createElement('canvas');
          const context = canvas.getContext('2d');
          canvas.width = viewport.width;
          canvas.height = viewport.height;
          
          await page.render({
            canvasContext: context,
            viewport: viewport
          }).promise;
          
          const img = new Image();
          await new Promise((resolve) => {
            img.onload = resolve;
            img.src = canvas.toDataURL('image/jpeg', 0.85);
          });
          
          loadedImages.push({
            original: img,
            processed: null,
            name: `${file.name}_page_${pageNum}.png`
          });
        }
        
        setImages(loadedImages);
        setCurrentPage(0);
        setProcessing(false);
        return;
      } catch (err) {
        console.error('PDF loading error:', err);
        alert('Error loading PDF: ' + err.message);
        setProcessing(false);
        return;
      }
    }

    // Handle image files
    if (file.type.startsWith('image/')) {
      const img = new Image();
      img.onload = () => {
        loadedImages.push({
          original: img,
          processed: null,
          name: file.name
        });
        setImages(loadedImages);
        setCurrentPage(0);
        setProcessing(false);
      };
      img.src = URL.createObjectURL(file);
    }
  };

  const handleMultipleImages = async (e) => {
    const files = Array.from(e.target.files);
    if (files.length === 0) return;

    setProcessing(true);
    const loadedImages = [];

    for (const file of files) {
      if (file.type.startsWith('image/')) {
        await new Promise((resolve) => {
          const img = new Image();
          img.onload = () => {
            loadedImages.push({
              original: img,
              processed: null,
              name: file.name
            });
            resolve();
          };
          img.src = URL.createObjectURL(file);
        });
      }
    }

    // Sort by filename
    loadedImages.sort((a, b) => a.name.localeCompare(b.name));
    setImages(loadedImages);
    setCurrentPage(0);
    setProcessing(false);
  };

  // Telea inpainting algorithm (Fast Marching Method)
  const inpaintTelea = (imageData, mask, radius) => {
    const { width, height, data } = imageData;
    const result = new Uint8ClampedArray(data);
    
    // Distance transform for prioritizing pixels
    const distance = new Float32Array(width * height).fill(Infinity);
    const known = new Uint8Array(width * height);
    
    // Initialize known pixels and boundary
    const boundary = [];
    for (let y = 0; y < height; y++) {
      for (let x = 0; x < width; x++) {
        const idx = y * width + x;
        if (!mask[idx]) {
          known[idx] = 1;
          distance[idx] = 0;
        }
      }
    }
    
    // Find initial boundary (unknown pixels adjacent to known)
    for (let y = 0; y < height; y++) {
      for (let x = 0; x < width; x++) {
        const idx = y * width + x;
        if (mask[idx]) {
          // Check 4-neighbors
          for (const [dx, dy] of [[0, 1], [0, -1], [1, 0], [-1, 0]]) {
            const nx = x + dx, ny = y + dy;
            if (nx >= 0 && nx < width && ny >= 0 && ny < height) {
              const nidx = ny * width + nx;
              if (!mask[nidx]) {
                distance[idx] = 1;
                boundary.push({ x, y, dist: 1 });
                break;
              }
            }
          }
        }
      }
    }
    
    // Sort boundary by distance (priority queue simulation)
    boundary.sort((a, b) => a.dist - b.dist);
    
    // Process pixels in order of distance
    while (boundary.length > 0) {
      const { x, y } = boundary.shift();
      const idx = y * width + x;
      
      if (known[idx]) continue;
      
      // Compute color using weighted average of known neighbors
      let totalWeight = 0;
      let r = 0, g = 0, b = 0;
      
      for (let dy = -radius; dy <= radius; dy++) {
        for (let dx = -radius; dx <= radius; dx++) {
          const nx = x + dx, ny = y + dy;
          if (nx >= 0 && nx < width && ny >= 0 && ny < height) {
            const nidx = ny * width + nx;
            if (known[nidx]) {
              const dist = Math.sqrt(dx * dx + dy * dy);
              if (dist <= radius && dist > 0) {
                // Weight based on distance and direction
                const weight = 1 / (dist * dist);
                const pixelIdx = nidx * 4;
                r += result[pixelIdx] * weight;
                g += result[pixelIdx + 1] * weight;
                b += result[pixelIdx + 2] * weight;
                totalWeight += weight;
              }
            }
          }
        }
      }
      
      if (totalWeight > 0) {
        const pixelIdx = idx * 4;
        result[pixelIdx] = r / totalWeight;
        result[pixelIdx + 1] = g / totalWeight;
        result[pixelIdx + 2] = b / totalWeight;
        result[pixelIdx + 3] = 255;
      }
      
      known[idx] = 1;
      
      // Add new boundary pixels
      for (const [dx, dy] of [[0, 1], [0, -1], [1, 0], [-1, 0]]) {
        const nx = x + dx, ny = y + dy;
        if (nx >= 0 && nx < width && ny >= 0 && ny < height) {
          const nidx = ny * width + nx;
          if (mask[nidx] && !known[nidx] && distance[nidx] === Infinity) {
            distance[nidx] = distance[idx] + 1;
            boundary.push({ x: nx, y: ny, dist: distance[nidx] });
            // Keep sorted (simple insertion)
            boundary.sort((a, b) => a.dist - b.dist);
          }
        }
      }
    }
    
    return new ImageData(result, width, height);
  };

  // Navier-Stokes based inpainting (diffusion-based)
  const inpaintNavierStokes = (imageData, mask, radius, iterations = 100) => {
    const { width, height, data } = imageData;
    const result = new Float32Array(data.length);
    for (let i = 0; i < data.length; i++) result[i] = data[i];
    
    // Iterative diffusion
    for (let iter = 0; iter < iterations; iter++) {
      const temp = new Float32Array(result);
      
      for (let y = 1; y < height - 1; y++) {
        for (let x = 1; x < width - 1; x++) {
          const idx = y * width + x;
          if (mask[idx]) {
            const pixelIdx = idx * 4;
            
            // Laplacian diffusion
            for (let c = 0; c < 3; c++) {
              const north = ((y - 1) * width + x) * 4 + c;
              const south = ((y + 1) * width + x) * 4 + c;
              const east = (y * width + (x + 1)) * 4 + c;
              const west = (y * width + (x - 1)) * 4 + c;
              
              result[pixelIdx + c] = (temp[north] + temp[south] + temp[east] + temp[west]) / 4;
            }
            result[pixelIdx + 3] = 255;
          }
        }
      }
    }
    
    const output = new Uint8ClampedArray(result.length);
    for (let i = 0; i < result.length; i++) {
      output[i] = Math.max(0, Math.min(255, Math.round(result[i])));
    }
    
    return new ImageData(output, width, height);
  };

  // Create feathered mask for smooth blending
  const createFeatheredMask = (width, height, region, featherRadius) => {
    const mask = new Uint8Array(width * height);
    const featherWeight = new Float32Array(width * height);
    
    const { x: rx, y: ry, w: rw, h: rh } = region;
    
    for (let y = 0; y < height; y++) {
      for (let x = 0; x < width; x++) {
        const idx = y * width + x;
        
        // Check if inside the watermark region
        if (x >= rx && x < rx + rw && y >= ry && y < ry + rh) {
          mask[idx] = 1;
          
          // Calculate feather weight (distance from edge)
          const distFromLeft = x - rx;
          const distFromRight = (rx + rw) - x;
          const distFromTop = y - ry;
          const distFromBottom = (ry + rh) - y;
          const minDist = Math.min(distFromLeft, distFromRight, distFromTop, distFromBottom);
          
          if (minDist < featherRadius) {
            featherWeight[idx] = minDist / featherRadius;
          } else {
            featherWeight[idx] = 1;
          }
        }
      }
    }
    
    return { mask, featherWeight };
  };

  // Process single image
  const processImage = useCallback((imgData, ctx, imgWidth, imgHeight) => {
    const { width: wmWidth, height: wmHeight, offsetX, offsetY, blendRadius, method } = settings;
    
    // Watermark region (bottom-right corner)
    const region = {
      x: imgWidth - wmWidth - offsetX,
      y: imgHeight - wmHeight - offsetY,
      w: wmWidth,
      h: wmHeight
    };
    
    // Get image data
    const imageData = ctx.getImageData(0, 0, imgWidth, imgHeight);
    
    // Create feathered mask
    const { mask, featherWeight } = createFeatheredMask(imgWidth, imgHeight, region, blendRadius);
    
    // Apply inpainting
    let result;
    if (method === 'telea') {
      result = inpaintTelea(imageData, mask, blendRadius);
    } else {
      result = inpaintNavierStokes(imageData, mask, blendRadius, 150);
    }
    
    // Blend with feathering for smooth edges
    const blended = new Uint8ClampedArray(imageData.data);
    for (let y = 0; y < imgHeight; y++) {
      for (let x = 0; x < imgWidth; x++) {
        const idx = y * imgWidth + x;
        if (mask[idx]) {
          const pixelIdx = idx * 4;
          const weight = featherWeight[idx];
          for (let c = 0; c < 3; c++) {
            blended[pixelIdx + c] = Math.round(
              result.data[pixelIdx + c] * weight + 
              imageData.data[pixelIdx + c] * (1 - weight)
            );
          }
        }
      }
    }
    
    return new ImageData(blended, imgWidth, imgHeight);
  }, [settings]);

  // Remove watermark from current page
  const removeWatermark = useCallback(() => {
    if (images.length === 0) return;
    
    setProcessing(true);
    
    setTimeout(() => {
      const canvas = canvasRef.current;
      const ctx = canvas.getContext('2d');
      const img = images[currentPage].original;
      
      canvas.width = img.width;
      canvas.height = img.height;
      ctx.drawImage(img, 0, 0);
      
      const result = processImage(null, ctx, img.width, img.height);
      ctx.putImageData(result, 0, 0);
      
      // Store processed result
      const processedImg = new Image();
      processedImg.src = canvas.toDataURL('image/png');
      
      setImages(prev => {
        const updated = [...prev];
        updated[currentPage] = {
          ...updated[currentPage],
          processed: processedImg,
          processedDataUrl: canvas.toDataURL('image/jpeg', 0.85)
        };
        return updated;
      });
      
      setProcessing(false);
    }, 50);
  }, [images, currentPage, processImage]);

  // Process all pages
  const processAllPages = useCallback(async () => {
    if (images.length === 0) return;
    
    setProcessing(true);
    
    const canvas = canvasRef.current;
    const ctx = canvas.getContext('2d');
    
    for (let i = 0; i < images.length; i++) {
      await new Promise(resolve => {
        setTimeout(() => {
          const img = images[i].original;
          canvas.width = img.width;
          canvas.height = img.height;
          ctx.drawImage(img, 0, 0);
          
          const result = processImage(null, ctx, img.width, img.height);
          ctx.putImageData(result, 0, 0);
          
          const processedImg = new Image();
          processedImg.src = canvas.toDataURL('image/png');
          
          setImages(prev => {
            const updated = [...prev];
            updated[i] = {
              ...updated[i],
              processed: processedImg,
              processedDataUrl: canvas.toDataURL('image/jpeg', 0.85)
            };
            return updated;
          });
          
          resolve();
        }, 50);
      });
    }
    
    setProcessing(false);
  }, [images, processImage]);

  // Download current processed image
  const downloadCurrent = () => {
    if (!images[currentPage]?.processedDataUrl) return;
    
    const link = document.createElement('a');
    link.download = `processed_${images[currentPage].name}`;
    link.href = images[currentPage].processedDataUrl;
    link.click();
  };

  // Download all as PDF
  const downloadAll = () => {
    const processedImages = images.filter(img => img.processedDataUrl);
    if (processedImages.length === 0) return;
    
    const firstImg = processedImages[0].processed;
    const pdf = new jsPDF({
      orientation: firstImg.width > firstImg.height ? 'landscape' : 'portrait',
      unit: 'px',
      format: [firstImg.width, firstImg.height]
    });
    
    processedImages.forEach((img, idx) => {
      if (idx > 0) {
        pdf.addPage([img.processed.width, img.processed.height], 
          img.processed.width > img.processed.height ? 'landscape' : 'portrait');
      }
      pdf.addImage(img.processedDataUrl, 'JPEG', 0, 0, img.processed.width, img.processed.height, undefined, 'FAST');
    });
    
    pdf.save(`${originalFileName || 'processed'}_new.pdf`);
  };

  // Draw preview with watermark region indicator
  const drawPreview = useCallback(() => {
    if (images.length === 0 || !canvasRef.current) return;
    
    const canvas = canvasRef.current;
    const ctx = canvas.getContext('2d');
    const currentImg = images[currentPage];
    const img = showOriginal || !currentImg.processed ? currentImg.original : currentImg.processed;
    
    canvas.width = img.width;
    canvas.height = img.height;
    ctx.drawImage(img, 0, 0);
    
    // Draw watermark region indicator
    if (showOriginal || !currentImg.processed) {
      const { width: wmWidth, height: wmHeight, offsetX, offsetY } = settings;
      ctx.strokeStyle = '#ff4444';
      ctx.lineWidth = 2;
      ctx.setLineDash([5, 5]);
      ctx.strokeRect(
        img.width - wmWidth - offsetX,
        img.height - wmHeight - offsetY,
        wmWidth,
        wmHeight
      );
      ctx.setLineDash([]);
    }
  }, [images, currentPage, showOriginal, settings]);

  React.useEffect(() => {
    drawPreview();
  }, [drawPreview]);

  return (
    <div className="min-h-screen bg-gray-900 text-white p-6">
      <div className="max-w-6xl mx-auto">
        <h1 className="text-3xl font-bold mb-2 text-center">Watermark Remover</h1>
        <p className="text-gray-400 text-center mb-6">
          Classic CV inpainting for bottom-right watermarks
        </p>

        {/* Upload Section */}
        <div className="bg-gray-800 rounded-lg p-6 mb-6">
          <div className="flex flex-wrap gap-4 justify-center">
            <label className="flex items-center gap-2 px-4 py-2 bg-blue-600 hover:bg-blue-700 rounded-lg cursor-pointer transition">
              <Upload size={20} />
              Upload PDF
              <input
                type="file"
                accept="application/pdf"
                onChange={handleFileUpload}
                className="hidden"
              />
            </label>
            
            <label className="flex items-center gap-2 px-4 py-2 bg-blue-600 hover:bg-blue-700 rounded-lg cursor-pointer transition">
              <Upload size={20} />
              Upload Images
              <input
                type="file"
                accept="image/*"
                multiple
                onChange={handleMultipleImages}
                className="hidden"
                ref={fileInputRef}
              />
            </label>
            
            <button
              onClick={() => setShowSettings(!showSettings)}
              className={`flex items-center gap-2 px-4 py-2 rounded-lg transition ${
                showSettings ? 'bg-purple-600' : 'bg-gray-700 hover:bg-gray-600'
              }`}
            >
              <Settings size={20} />
              Settings
            </button>
          </div>

          {/* Settings Panel */}
          {showSettings && (
            <div className="mt-4 p-4 bg-gray-700 rounded-lg">
              <h3 className="font-semibold mb-3">Watermark Region (Bottom-Right Corner)</h3>
              <div className="grid grid-cols-2 md:grid-cols-3 gap-4">
                <div>
                  <label className="block text-sm text-gray-400 mb-1">Width (px)</label>
                  <input
                    type="number"
                    value={settings.width}
                    onChange={(e) => setSettings(s => ({ ...s, width: parseInt(e.target.value) || 0 }))}
                    className="w-full px-3 py-2 bg-gray-600 rounded"
                  />
                </div>
                <div>
                  <label className="block text-sm text-gray-400 mb-1">Height (px)</label>
                  <input
                    type="number"
                    value={settings.height}
                    onChange={(e) => setSettings(s => ({ ...s, height: parseInt(e.target.value) || 0 }))}
                    className="w-full px-3 py-2 bg-gray-600 rounded"
                  />
                </div>
                <div>
                  <label className="block text-sm text-gray-400 mb-1">Offset X (px)</label>
                  <input
                    type="number"
                    value={settings.offsetX}
                    onChange={(e) => setSettings(s => ({ ...s, offsetX: parseInt(e.target.value) || 0 }))}
                    className="w-full px-3 py-2 bg-gray-600 rounded"
                  />
                </div>
                <div>
                  <label className="block text-sm text-gray-400 mb-1">Offset Y (px)</label>
                  <input
                    type="number"
                    value={settings.offsetY}
                    onChange={(e) => setSettings(s => ({ ...s, offsetY: parseInt(e.target.value) || 0 }))}
                    className="w-full px-3 py-2 bg-gray-600 rounded"
                  />
                </div>
                <div>
                  <label className="block text-sm text-gray-400 mb-1">Blend Radius</label>
                  <input
                    type="number"
                    value={settings.blendRadius}
                    onChange={(e) => setSettings(s => ({ ...s, blendRadius: parseInt(e.target.value) || 1 }))}
                    className="w-full px-3 py-2 bg-gray-600 rounded"
                  />
                </div>
                <div>
                  <label className="block text-sm text-gray-400 mb-1">Method</label>
                  <select
                    value={settings.method}
                    onChange={(e) => setSettings(s => ({ ...s, method: e.target.value }))}
                    className="w-full px-3 py-2 bg-gray-600 rounded"
                  >
                    <option value="telea">Telea (Fast Marching)</option>
                    <option value="ns">Navier-Stokes (Diffusion)</option>
                  </select>
                </div>
              </div>
            </div>
          )}
        </div>

        {/* Main Content */}
        {images.length > 0 && (
          <div className="bg-gray-800 rounded-lg p-6">
            {/* Controls */}
            <div className="flex flex-wrap gap-3 mb-4 justify-center">
              <button
                onClick={removeWatermark}
                disabled={processing}
                className="flex items-center gap-2 px-4 py-2 bg-green-600 hover:bg-green-700 disabled:bg-gray-600 rounded-lg transition"
              >
                <Play size={20} />
                {processing ? 'Processing...' : 'Remove Watermark'}
              </button>
              
              <button
                onClick={processAllPages}
                disabled={processing}
                className="flex items-center gap-2 px-4 py-2 bg-green-600 hover:bg-green-700 disabled:bg-gray-600 rounded-lg transition"
              >
                <Play size={20} />
                Process All ({images.length})
              </button>
              
              <button
                onClick={() => setShowOriginal(!showOriginal)}
                className="flex items-center gap-2 px-4 py-2 bg-gray-700 hover:bg-gray-600 rounded-lg transition"
              >
                {showOriginal ? <Eye size={20} /> : <EyeOff size={20} />}
                {showOriginal ? 'Show Processed' : 'Show Original'}
              </button>
              
              {images[currentPage]?.processed && (
                <>
                  <button
                    onClick={downloadCurrent}
                    className="flex items-center gap-2 px-4 py-2 bg-blue-600 hover:bg-blue-700 rounded-lg transition"
                  >
                    <Download size={20} />
                    Download
                  </button>
                  
                  <button
                    onClick={downloadAll}
                    className="flex items-center gap-2 px-4 py-2 bg-blue-600 hover:bg-blue-700 rounded-lg transition"
                  >
                    <Download size={20} />
                    Download full PDF
                  </button>
                </>
              )}
              
              <button
                onClick={() => {
                  setImages([]);
                  setCurrentPage(0);
                }}
                className="flex items-center gap-2 px-4 py-2 bg-red-600 hover:bg-red-700 rounded-lg transition"
              >
                <RotateCcw size={20} />
                Reset
              </button>
            </div>

            {/* Navigation */}
            {images.length > 1 && (
              <div className="flex items-center justify-center gap-4 mb-4">
                <button
                  onClick={() => setCurrentPage(p => Math.max(0, p - 1))}
                  disabled={currentPage === 0}
                  className="p-2 bg-gray-700 hover:bg-gray-600 disabled:bg-gray-800 disabled:text-gray-600 rounded-lg transition"
                >
                  <ChevronLeft size={24} />
                </button>
                <span className="text-lg">
                  Page {currentPage + 1} of {images.length}
                </span>
                <button
                  onClick={() => setCurrentPage(p => Math.min(images.length - 1, p + 1))}
                  disabled={currentPage === images.length - 1}
                  className="p-2 bg-gray-700 hover:bg-gray-600 disabled:bg-gray-800 disabled:text-gray-600 rounded-lg transition"
                >
                  <ChevronRight size={24} />
                </button>
              </div>
            )}

            {/* Canvas Preview */}
            <div className="flex justify-center overflow-auto bg-gray-900 rounded-lg p-4">
              <canvas
                ref={canvasRef}
                className="max-w-full h-auto border border-gray-700"
                style={{ maxHeight: '600px' }}
              />
            </div>

            {/* Status */}
            <div className="mt-4 text-center text-sm text-gray-400">
              {processing && <span className="text-yellow-400">Processing... Please wait.</span>}
              {!processing && images[currentPage]?.processed && (
                <span className="text-green-400">✓ Watermark removed</span>
              )}
              {!processing && !images[currentPage]?.processed && (
                <span>Red dashed box shows watermark region to remove</span>
              )}
            </div>
          </div>
        )}

        {/* Instructions */}
        {images.length === 0 && (
          <div className="bg-gray-800 rounded-lg p-6 text-center">
            <h3 className="text-xl font-semibold mb-4">How to Use</h3>
            <ol className="text-left max-w-md mx-auto space-y-2 text-gray-300">
              <li>1. Convert your PDF pages to PNG/JPG images first</li>
              <li>2. Upload all page images using the button above</li>
              <li>3. Adjust the watermark region settings if needed</li>
              <li>4. Click "Process All" to remove watermarks</li>
              <li>5. Download the processed images</li>
            </ol>
            <p className="mt-4 text-sm text-gray-500">
              Tip: Use tools like <code className="bg-gray-700 px-1 rounded">pdftoppm</code> or online converters to convert PDF to images
            </p>
          </div>
        )}

        {/* Algorithm Info */}
        <div className="mt-6 text-sm text-gray-500 text-center">
          <p>
            Uses Telea (Fast Marching) or Navier-Stokes inpainting algorithms.
            <br />
            Feathered blending ensures smooth edges.
          </p>
        </div>
      </div>
    </div>
  );
}