app.js

import {
  AutoModel, AutoProcessor, env, RawImage
} from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.5.1/dist/transformers.min.js';

env.allowLocalModels = false;
env.backends.onnx.wasm.proxy = false;

document.addEventListener('DOMContentLoaded', () => {
  const $ = id => document.getElementById(id);

  const html          = document.documentElement;
  const themeToggle   = $('themeToggle');
  const iconSun       = $('iconSun');
  const iconMoon      = $('iconMoon');
  const uploadZone    = $('uploadZone');
  const fileInput     = $('fileInput');
  const panel         = $('panel');
  const originalImg   = $('originalImg');
  const resultBody    = $('resultBody');
  const resultStatus  = $('resultStatus');
  const processBtn    = $('processBtn');
  const resetBtn      = $('resetBtn');
  const progressWrap  = $('progressWrap');
  const progressFill  = $('progressFill');
  const progressLabel = $('progressLabel');
  const uploadHint    = $('uploadHint');
  const messageBanner = $('messageBanner');
  const queueList = $('queueList');
  const historyList = $('historyList');
  const downloadZipBtn = $('downloadZipBtn');
  const downloadHistoryZipBtn = $('downloadHistoryZipBtn');
  const clearHistoryBtn = $('clearHistoryBtn');

  const MAX_FILE_SIZE = 20 * 1024 * 1024;
  const ACCEPTED_TYPES = ['image/png', 'image/jpeg', 'image/webp'];

  let model = null;
  let processor = null;
  let currentFile = null;
  let currentDataURL = null;
  let filesQueue = [];
  let history = [];
  let currentBatchResults = [];
  let isDark = false;
  let isBusy = false;

  function applyTheme(dark) {
    isDark = dark;
    html.setAttribute('data-theme', isDark ? 'dark' : 'light');
    iconSun.style.display  = isDark ? 'none' : '';
    iconMoon.style.display = isDark ? '' : 'none';
  }

  function initTheme() {
    const savedTheme = localStorage.getItem('theme');

    if (savedTheme !== null) {
      applyTheme(savedTheme === 'dark');
    } else {
      const prefersDark = window.matchMedia &&
                         window.matchMedia('(prefers-color-scheme: dark)').matches;
      applyTheme(prefersDark);
    }
  }

  function setMessage(text, type = 'info') {
    if (!text) {
      messageBanner.className = 'message-banner';
      messageBanner.textContent = '';
      return;
    }

    messageBanner.className = `message-banner visible ${type}`;
    messageBanner.textContent = text;
  }

  function setProgress(pct, label) {
    progressFill.style.width = pct + '%';
    progressLabel.textContent = label;
  }

  function resetProgress() {
    setProgress(0, 'Waiting for image...');
    progressWrap.classList.remove('visible');
  }

  function fileToDataURL(file) {
    return new Promise((resolve, reject) => {
      try {
        const reader = new FileReader();
        reader.onload = (e) => resolve(e.target.result);
        reader.onerror = (err) => reject(err);
        reader.readAsDataURL(file);
      } catch (err) {
        reject(err);
      }
    });
  }

  function dataURLtoBlob(dataURL) {
    const parts = dataURL.split(',');
    const meta = parts[0] || '';
    const bstr = atob(parts[1] || '');
    const mimeMatch = meta.match(/:(.*?);/);
    const mime = mimeMatch ? mimeMatch[1] : 'application/octet-stream';
    const n = bstr.length;
    const u8 = new Uint8Array(n);
    for (let i = 0; i < n; i++) u8[i] = bstr.charCodeAt(i);
    return new Blob([u8], { type: mime });
  }

  function imageDataFromDataURL(dataURL) {
    return new Promise((resolve, reject) => {
      const img = new Image();
      img.onload = () => {
        const c = document.createElement('canvas');
        c.width = img.naturalWidth;
        c.height = img.naturalHeight;
        const ctx = c.getContext('2d');
        ctx.drawImage(img, 0, 0);
        resolve({ imageData: ctx.getImageData(0, 0, c.width, c.height), width: c.width, height: c.height });
      };
      img.onerror = reject;
      img.src = dataURL;
    });
  }

  function setUploadState(state) {
    uploadZone.classList.remove('is-error', 'is-ready', 'is-processing');
    if (state) uploadZone.classList.add(`is-${state}`);
  }

  function validateFile(file) {
    if (!file) return { valid: false, message: 'No file was selected.' };

    const extension = (file.name || '').split('.').pop()?.toLowerCase();
    const normalizedType = (file.type || '').toLowerCase();
    const hasSupportedType = ACCEPTED_TYPES.includes(normalizedType) || ['png', 'jpg', 'jpeg', 'webp'].includes(extension);

    if (!hasSupportedType) {
      return { valid: false, message: 'Please choose a PNG, JPEG, or WebP image.' };
    }

    if (file.size > MAX_FILE_SIZE) {
      return { valid: false, message: 'The selected file is larger than 20 MB. Please choose a smaller image.' };
    }

    return { valid: true };
  }

  function resetResultArea() {
    resultBody.innerHTML = `
      <div class="result-empty">
        <svg width="32" height="32" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="1.2">
          <rect x="3" y="3" width="18" height="18" rx="2"/>
          <circle cx="8.5" cy="8.5" r="1.5"/><polyline points="21 15 16 10 5 21"/>
        </svg>
        <span>Your result will appear here</span>
      </div>`;
    resultBody.style.flexDirection = '';
    resultBody.classList.remove('has-result');
    resultStatus.style.display = 'none';
    resultStatus.className = 'status-badge processing';
    resultStatus.style.background = '';
    resultStatus.style.color = '';
  }

  function setBusy(busy) {
    isBusy = busy;
    processBtn.disabled = busy;
    fileInput.disabled = busy;
    uploadZone.classList.toggle('is-processing', busy);
  }

  themeToggle.addEventListener('click', () => {
    const newIsDark = !isDark;
    applyTheme(newIsDark);
    localStorage.setItem('theme', newIsDark ? 'dark' : 'light');
  });

  if (window.matchMedia) {
    window.matchMedia('(prefers-color-scheme: dark)').addEventListener('change', e => {
      if (!localStorage.getItem('theme')) {
        applyTheme(e.matches);
      }
    });
  }

  uploadZone.addEventListener('dragover', e => {
    e.preventDefault();
    uploadZone.classList.add('drag-over');
  });
  uploadZone.addEventListener('dragleave', () => {
    uploadZone.classList.remove('drag-over');
  });
  uploadZone.addEventListener('drop', e => {
    e.preventDefault();
    uploadZone.classList.remove('drag-over');
    const files = e.dataTransfer.files;
    if (files && files.length) handleFile(files);
  });
  uploadZone.addEventListener('keydown', e => {
    if (e.key === 'Enter' || e.key === ' ') {
      e.preventDefault();
      fileInput.click();
    }
  });
  fileInput.addEventListener('change', () => {
    if (fileInput.files && fileInput.files.length) handleFile(fileInput.files);
  });

  async function handleFile(fileOrList) {
    const files = fileOrList instanceof FileList || Array.isArray(fileOrList) ? Array.from(fileOrList) : [fileOrList];
    let added = 0;
    for (const file of files) {
      const validation = validateFile(file);
      if (!validation.valid) {
        setUploadState('error');
        setMessage(validation.message, 'error');
        continue;
      }
      try {
        const dataURL = await fileToDataURL(file);
        filesQueue.push({ file, name: file.name, dataURL });
        added++;
      } catch (err) {
        console.error('[Upload] error:', err);
      }
    }

    if (filesQueue.length > 0) {
      // show first item as preview
      currentFile = filesQueue[0].file;
      currentDataURL = filesQueue[0].dataURL;
      originalImg.src = currentDataURL;
      uploadZone.style.display = 'none';
      panel.classList.add('visible');
      resetResultArea();
      resultStatus.style.display = 'none';
      resetProgress();
      setMessage(`${filesQueue.length} image(s) added to queue. Click Remove Background to process.`, 'success');
      setUploadState('ready');
      renderQueue();
      updateOriginalLabel();
      setBusy(false);
    } else if (added === 0) {
      setMessage('No valid images were added.', 'warning');
    }
  }

  function renderQueue() {
    if (!queueList) return;
    queueList.innerHTML = '';
    filesQueue.forEach((it, idx) => {
      const el = document.createElement('div');
      el.className = 'item';
      el.innerHTML = `<img src="${it.dataURL}" alt="${it.name}"><div class="meta">${it.name}</div><div class="remove" data-idx="${idx}">✕</div>`;
      queueList.appendChild(el);
    });
    // attach remove handlers
    queueList.querySelectorAll('.remove').forEach(btn => {
      btn.addEventListener('click', (e) => {
        const idx = Number(e.currentTarget.getAttribute('data-idx'));
        filesQueue.splice(idx, 1);
        renderQueue();
        if (filesQueue.length === 0) {
          // restore upload view
          panel.classList.remove('visible');
          uploadZone.style.display = '';
        } else {
          currentFile = filesQueue[0].file;
          currentDataURL = filesQueue[0].dataURL;
          originalImg.src = currentDataURL;
        }
      });
    });
    updateOriginalLabel();
  }

  function loadHistory() {
    try {
      const raw = localStorage.getItem('rb_history');
      history = raw ? JSON.parse(raw) : [];
    } catch (e) {
      history = [];
    }
    renderHistory();
  }

  function updateOriginalLabel() {
    const lbl = document.getElementById('originalLabel');
    if (!lbl) return;
    const count = filesQueue.length || (currentFile ? 1 : 0);
    lbl.textContent = count > 1 ? `Original (${count})` : 'Original';
  }

  function saveHistory() {
    try {
      localStorage.setItem('rb_history', JSON.stringify(history));
    } catch (e) {
      console.warn('Could not save history', e);
    }
  }

  function addToHistory(item) {
    history.unshift(item);
    // limit history
    if (history.length > 40) history.length = 40;
    saveHistory();
    renderHistory();
  }

  function renderHistory() {
    if (!historyList) return;
    historyList.innerHTML = '';
    history.forEach((h, idx) => {
      const el = document.createElement('div');
      el.className = 'item';
      el.innerHTML = `<img src="${h.dataURL}" alt="${h.name}"><div class="meta">${h.name}</div>`;
      el.title = 'Download image';
      el.addEventListener('click', () => {
        // download
        const a = document.createElement('a');
        a.href = h.dataURL;
        a.download = h.name || `removed-${idx}.png`;
        a.click();
      });
      historyList.appendChild(el);
    });

    // show/hide history action buttons
    if (downloadHistoryZipBtn) downloadHistoryZipBtn.style.display = history.length > 0 ? 'inline-block' : 'none';
    if (clearHistoryBtn) clearHistoryBtn.style.display = history.length > 0 ? 'inline-block' : 'none';
  }

  if (downloadHistoryZipBtn) {
    downloadHistoryZipBtn.addEventListener('click', async () => {
      if (!history || history.length === 0) return;
      try {
        setMessage('Preparing ZIP of history...', 'info');
        const zip = new JSZip();
        for (const it of history) {
          zip.file(it.name, dataURLtoBlob(it.dataURL));
        }
        const blob = await zip.generateAsync({ type: 'blob' });
        const a = document.createElement('a');
        a.href = URL.createObjectURL(blob);
        a.download = 'history-removed-images.zip';
        a.click();
        URL.revokeObjectURL(a.href);
        setMessage('ZIP downloaded', 'success');
      } catch (e) {
        console.error('ZIP error', e);
        setMessage('Failed to create ZIP from history.', 'error');
      }
    });
  }

  if (clearHistoryBtn) {
    clearHistoryBtn.addEventListener('click', () => {
      try {
        if (!confirm('Clear history? This cannot be undone.')) return;
        history = [];
        saveHistory();
        renderHistory();
        setMessage('History cleared', 'success');
      } catch (e) {
        console.error('Clear history error', e);
        setMessage('Failed to clear history.', 'error');
      }
    });
  }

  async function ensureModelReady() {
    if (model && processor) return;

    setProgress(8, 'Downloading model files...');
    try {
      model = await AutoModel.from_pretrained('briaai/RMBG-1.4', {
        config: { model_type: 'custom' },
      });
      setProgress(36, 'Preparing image processor...');
      processor = await AutoProcessor.from_pretrained('briaai/RMBG-1.4', {
        config: {
          do_normalize: true,
          do_pad: false,
          do_rescale: true,
          do_resize: true,
          image_mean: [0.5, 0.5, 0.5],
          image_std: [1, 1, 1],
          resample: 2,
          rescale_factor: 1 / 255,
          size: { width: 1024, height: 1024 },
        },
      });
    } catch (error) {
      throw new Error(`The AI model could not be loaded. Please check your connection and try again. ${error?.message || ''}`.trim());
    }
  }

  async function runSegmentation(dataURL) {
    await ensureModelReady();

    setProgress(46, 'Reading image...');
    const image = await RawImage.fromURL(dataURL);
    const origW = image.width;
    const origH = image.height;

    setProgress(58, 'Preparing mask...');
    const { pixel_values } = await processor(image);

    setProgress(74, 'Running inference...');
    const { output } = await model({ input: pixel_values });

    setProgress(86, 'Refining mask...');
    const tensor = output[0];
    const dims = tensor.dims;
    const raw = tensor.data;

    const tH = dims[dims.length - 2];
    const tW = dims[dims.length - 1];
    const nPx = tH * tW;
    const offset = raw.length - nPx;

    let dMin = Infinity;
    let dMax = -Infinity;
    for (let i = 0; i < nPx; i++) {
      const v = raw[offset + i];
      if (v < dMin) dMin = v;
      if (v > dMax) dMax = v;
    }

    const needsSigmoid = dMin < -0.01 || dMax > 1.01;
    const sigmoid = x => 1 / (1 + Math.exp(-x));
    const maskTensor = new Float32Array(nPx);
    for (let i = 0; i < nPx; i++) {
      const v = raw[offset + i];
      maskTensor[i] = needsSigmoid ? sigmoid(v) : v;
    }

    const maskFull = new Float32Array(origW * origH);
    const scaleX = (tW - 1) / Math.max(origW - 1, 1);
    const scaleY = (tH - 1) / Math.max(origH - 1, 1);

    for (let y = 0; y < origH; y++) {
      const ty = y * scaleY;
      const ty0 = Math.floor(ty);
      const ty1 = Math.min(ty0 + 1, tH - 1);
      const dy = ty - ty0;
      for (let x = 0; x < origW; x++) {
        const tx = x * scaleX;
        const tx0 = Math.floor(tx);
        const tx1 = Math.min(tx0 + 1, tW - 1);
        const dx = tx - tx0;
        const v00 = maskTensor[ty0 * tW + tx0];
        const v10 = maskTensor[ty0 * tW + tx1];
        const v01 = maskTensor[ty1 * tW + tx0];
        const v11 = maskTensor[ty1 * tW + tx1];
        maskFull[y * origW + x] = v00 * (1 - dx) * (1 - dy) + v10 * dx * (1 - dy) + v01 * (1 - dx) * dy + v11 * dx * dy;
      }
    }

    let borderSum = 0;
    let borderCount = 0;
    for (let x = 0; x < origW; x++) {
      borderSum += maskFull[x] + maskFull[(origH - 1) * origW + x];
      borderCount += 2;
    }
    for (let y = 1; y < origH - 1; y++) {
      borderSum += maskFull[y * origW] + maskFull[y * origW + origW - 1];
      borderCount += 2;
    }
    const borderMean = borderCount > 0 ? borderSum / borderCount : 0;
    const isInverted = borderMean > 0.5;
    if (isInverted) {
      for (let i = 0; i < maskFull.length; i++) maskFull[i] = 1 - maskFull[i];
    }

    setProgress(94, 'Compositing image...');
    const { imageData: origImageData } = await imageDataFromDataURL(dataURL);
    const origPx = origImageData.data;
    const canvas = document.createElement('canvas');
    canvas.width = origW;
    canvas.height = origH;
    const ctx = canvas.getContext('2d');
    const out = ctx.createImageData(origW, origH);

    for (let i = 0; i < origW * origH; i++) {
      out.data[4 * i] = origPx[4 * i];
      out.data[4 * i + 1] = origPx[4 * i + 1];
      out.data[4 * i + 2] = origPx[4 * i + 2];
      out.data[4 * i + 3] = Math.round(maskFull[i] * 255);
    }
    ctx.putImageData(out, 0, 0);

    return canvas.toDataURL('image/png');
  }

  resetBtn.addEventListener('click', () => {
    currentFile = null;
    currentDataURL = null;
    fileInput.value = '';
    originalImg.src = '';
    panel.classList.remove('visible');
    uploadZone.style.display = '';
    uploadZone.classList.remove('drag-over', 'is-processing', 'is-ready', 'is-error');
    resetResultArea();
    resultStatus.style.display = 'none';
    resetProgress();
    setMessage('Choose a PNG, JPEG, or WebP image to begin.', 'info');
    setBusy(false);
    uploadHint.textContent = 'Supports PNG, JPEG, and WebP — up to 20 MB.';
    filesQueue = [];
    currentBatchResults = [];
    if (downloadZipBtn) downloadZipBtn.style.display = 'none';
    updateOriginalLabel();
  });

  processBtn.addEventListener('click', async () => {
    if (isBusy) return;
    if ((!currentDataURL || !currentDataURL.length) && filesQueue.length === 0) return;

    currentBatchResults = [];
    if (downloadZipBtn) downloadZipBtn.style.display = 'none';
    setBusy(true);
    resultStatus.style.display = 'inline-flex';
    resultStatus.className = 'status-badge processing';
    resultStatus.innerHTML = '<div class="dot pulse"></div>Processing';
    progressWrap.classList.add('visible');
    setMessage('Processing images...', 'info');

    try {
      if (filesQueue.length > 0) {
        const total = filesQueue.length;
        while (filesQueue.length > 0) {
          const idx = total - filesQueue.length;
          const item = filesQueue[0];
          setProgress(Math.round((idx / total) * 90) + 5, `Processing ${idx + 1}/${total}...`);
          const resultDataURL = await runSegmentation(item.dataURL);

          // show result for this item
          const resultImg = new Image();
          resultImg.src = resultDataURL;
          resultImg.style.cssText = 'max-width:100%;max-height:320px;object-fit:contain;border-radius:4px;';
          const dlBtn = document.createElement('button');
          dlBtn.className = 'btn btn-secondary';
          dlBtn.innerHTML = '<svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round"><path d="M21 15v4a2 2 0 0 1-2 2H5a2 2 0 0 1-2-2v-4"/><polyline points="7 10 12 15 17 10"/><line x1="12" y1="15" x2="12" y2="3"/></svg> Download PNG';
          dlBtn.style.cssText = 'margin-top:12px;width:fit-content;';
          dlBtn.addEventListener('click', () => {
            const a = document.createElement('a');
            a.href = resultDataURL;
            a.download = (item.name ? item.name.replace(/\.[^.]+$/, '') : 'removed') + '-bg.png';
            a.click();
          });

          resultBody.innerHTML = '';
          resultBody.style.flexDirection = 'column';
          resultBody.classList.add('has-result');
          resultBody.appendChild(resultImg);
          resultBody.appendChild(dlBtn);

          resultStatus.innerHTML = '<div class="dot"></div>Ready';
          resultStatus.className = 'status-badge ready';

          // save to history
          addToHistory({ name: (item.name ? item.name.replace(/\.[^.]+$/, '') : 'removed') + '.png', dataURL: resultDataURL, ts: Date.now() });
          currentBatchResults.push({ name: (item.name ? item.name.replace(/\.[^.]+$/, '') : 'removed') + '.png', dataURL: resultDataURL });

          // remove processed item
          filesQueue.shift();
          renderQueue();
        }

        if (currentBatchResults.length > 1 && downloadZipBtn) {
          downloadZipBtn.style.display = 'inline-flex';
          downloadZipBtn.onclick = async () => {
            try {
              setMessage('Preparing ZIP...', 'info');
              const zip = new JSZip();
              for (const it of currentBatchResults) {
                zip.file(it.name, dataURLtoBlob(it.dataURL));
              }
              const blob = await zip.generateAsync({ type: 'blob' });
              const a = document.createElement('a');
              a.href = URL.createObjectURL(blob);
              a.download = 'removed-images.zip';
              a.click();
              URL.revokeObjectURL(a.href);
              setMessage('ZIP downloaded', 'success');
            } catch (e) {
              console.error('ZIP error', e);
              setMessage('Failed to create ZIP.', 'error');
            }
          };
        } else if (downloadZipBtn) {
          downloadZipBtn.style.display = 'none';
          downloadZipBtn.onclick = null;
        }
        setProgress(100, 'All done ✓');
        setMessage('All images processed and saved to history.', 'success');
      } else {
        // single image
        setProgress(5, 'Preparing segmentation...');
        const resultDataURL = await runSegmentation(currentDataURL);

        const resultImg = new Image();
        resultImg.src = resultDataURL;
        resultImg.style.cssText = 'max-width:100%;max-height:320px;object-fit:contain;border-radius:4px;';
        const dlBtn = document.createElement('button');
        dlBtn.className = 'btn btn-secondary';
        dlBtn.innerHTML = '<svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round"><path d="M21 15v4a2 2 0 0 1-2 2H5a2 2 0 0 1-2-2v-4"/><polyline points="7 10 12 15 17 10"/><line x1="12" y1="15" x2="12" y2="3"/></svg> Download PNG';
        dlBtn.style.cssText = 'margin-top:12px;width:fit-content;';
        dlBtn.addEventListener('click', () => {
          const a = document.createElement('a');
          a.href = resultDataURL;
          a.download = (currentFile && currentFile.name ? currentFile.name.replace(/\.[^.]+$/, '') : 'removed') + '-bg.png';
          a.click();
        });

        resultBody.innerHTML = '';
        resultBody.style.flexDirection = 'column';
        resultBody.classList.add('has-result');
        resultBody.appendChild(resultImg);
        resultBody.appendChild(dlBtn);

        resultStatus.innerHTML = '<div class="dot"></div>Ready';
        resultStatus.className = 'status-badge ready';
        addToHistory({ name: (currentFile && currentFile.name ? currentFile.name.replace(/\.[^.]+$/, '') : 'removed') + '.png', dataURL: resultDataURL, ts: Date.now() });
        setProgress(100, 'Done ✓');
        setMessage('Background removed successfully. Download your PNG whenever you want.', 'success');
      }
    } catch (error) {
      console.error('[RMBG] error:', error);
      const detail = error?.message || 'Something went wrong while processing the image.';
      setProgress(0, '⚠ ' + detail);
      resultStatus.innerHTML = '<div class="dot"></div>Error';
      resultStatus.className = 'status-badge error';
      setMessage(`Processing failed: ${detail}`, 'error');
    } finally {
      setBusy(false);
      progressWrap.classList.remove('visible');
    }
  });

  initTheme();
  setMessage('Choose a PNG, JPEG, or WebP image to begin.', 'info');
  uploadHint.textContent = 'Supports PNG, JPEG, and WebP — up to 20 MB.';
  loadHistory();
});