Blame for src/main.js

/**
 * Pulpie WebGPU - Main Application Controller
 *
 * Orchestrates WebGPU hardware detection, model caching, EuroBERT tokenization,
 * block segmentation, and rendering for sizzlebop/pulpie-orange-small-onnx.
 */

import { checkWebGpuSupport, createSession, runWebGpuInference } from './lib/webgpu-engine.js';
import { loadModelBuffer } from './lib/model-cache.js';
import { getTokenizer, tokenizeBlocks } from './lib/tokenizer.js';
import { simplifyAndExtractBlocks } from './lib/html-simplifier.js';
import { packChunks } from './lib/chunker.js';
import { reconstructMainContent } from './lib/reconstructor.js';
import { fetchWebpage } from './lib/url-fetcher.js';
import { SAMPLE_HTMLS } from './lib/samples.js';
import { displayResults } from './lib/ui-renderer.js';

// Application State
const state = {
  gpuSupported: false,
  gpuInfo: null,
  isModelReady: false,
  isLoadingModel: false,
  session: null,
  provider: 'webgpu', // 'webgpu' | 'wasm'
  tokenizer: null,
  currentMode: 'url', // 'url' | 'html'
  activeOutputTab: 'preview', // 'preview' | 'markdown' | 'html' | 'blocks'
  lastResult: null,
};

// DOM Elements
const elements = {
  webgpuBadge: document.getElementById('webgpuBadge'),

  modelBanner: document.getElementById('modelBanner'),
  progressBarContainer: document.getElementById('progressBarContainer'),
  progressBarFill: document.getElementById('progressBarFill'),
  progressStatusText: document.getElementById('progressStatusText'),
  progressSpeedText: document.getElementById('progressSpeedText'),

  modeUrlTab: document.getElementById('modeUrlTab'),
  modeHtmlTab: document.getElementById('modeHtmlTab'),
  urlView: document.getElementById('urlView'),
  htmlView: document.getElementById('htmlView'),
  urlInput: document.getElementById('urlInput'),
  clearUrlBtn: document.getElementById('clearUrlBtn'),
  pasteUrlBtn: document.getElementById('pasteUrlBtn'),
  extractBtn: document.getElementById('extractBtn'),

  htmlInput: document.getElementById('htmlInput'),
  htmlDropZone: document.getElementById('htmlDropZone'),
  dropOverlay: document.getElementById('dropOverlay'),
  baseUrlInput: document.getElementById('baseUrlInput'),
  clearHtmlBtn: document.getElementById('clearHtmlBtn'),
  extractHtmlBtn: document.getElementById('extractHtmlBtn'),

  statusAlert: document.getElementById('statusAlert'),
  statusAlertIcon: document.getElementById('statusAlertIcon'),
  statusAlertMsg: document.getElementById('statusAlertMsg'),

  resultsCard: document.getElementById('resultsCard'),
  pageTitleDisplay: document.getElementById('pageTitleDisplay'),
  sourceUrlDisplay: document.getElementById('sourceUrlDisplay'),
  statKeptBlocks: document.getElementById('statKeptBlocks'),
  statRetention: document.getElementById('statRetention'),
  statTokens: document.getElementById('statTokens'),
  statTimeLabel: document.getElementById('statTimeLabel'),
  statGpuTime: document.getElementById('statGpuTime'),

  copyBtn: document.getElementById('copyBtn'),
  copyBtnText: document.getElementById('copyBtnText'),
  downloadMdBtn: document.getElementById('downloadMdBtn'),

  tabViewPreview: document.getElementById('tabViewPreview'),
  tabViewMarkdown: document.getElementById('tabViewMarkdown'),
  tabViewHtml: document.getElementById('tabViewHtml'),
  tabViewBlocks: document.getElementById('tabViewBlocks'),
  readerContent: document.getElementById('readerContent'),
  markdownContent: document.getElementById('markdownContent'),
  htmlContent: document.getElementById('htmlContent'),
  blockInspectorList: document.getElementById('blockInspectorList'),
};

function showStatus(msg, type = 'info') {
  if (!elements.statusAlert) return;
  elements.statusAlert.className = `status-alert ${type}`;
  elements.statusAlertMsg.textContent = msg;
  elements.statusAlert.style.display = 'flex';
}

function hideStatus() {
  if (elements.statusAlert) elements.statusAlert.style.display = 'none';
}

async function initializeSystem() {
  const gpuCheck = await checkWebGpuSupport();
  state.gpuSupported = gpuCheck.supported;
  state.gpuInfo = gpuCheck;
}

async function loadAndPrepareModel(providedBuffer = null) {
  if (state.isLoadingModel) return;
  state.isLoadingModel = true;
  hideStatus();

  if (elements.progressBarContainer) {
    elements.progressBarContainer.style.display = 'flex';
  }
  if (elements.progressStatusText) {
    elements.progressStatusText.textContent = 'Preparing model...';
  }

  try {
    let modelBuffer = providedBuffer;
    if (!modelBuffer) {
      modelBuffer = await loadModelBuffer(null, ({ loaded, total, percent, speedMbps, fromCache }) => {
        if (elements.progressBarFill) {
          elements.progressBarFill.style.width = `${percent}%`;
        }
        if (fromCache) {
          if (elements.progressStatusText) {
            elements.progressStatusText.textContent = 'Loading model from browser cache...';
          }
          if (elements.progressSpeedText) {
            elements.progressSpeedText.textContent = 'Cached';
          }
        } else {
          const loadedMb = (loaded / (1024 * 1024)).toFixed(1);
          const totalMb = (total / (1024 * 1024)).toFixed(1);
          if (elements.progressStatusText) {
            elements.progressStatusText.textContent = `Downloading model: ${loadedMb} / ${totalMb} MB (${speedMbps} Mbps)`;
          }
          if (elements.progressSpeedText) {
            elements.progressSpeedText.textContent = `${percent}%`;
          }
        }
      });
    }

    if (elements.progressStatusText) {
      elements.progressStatusText.textContent = 'Compiling model session...';
    }
    if (elements.progressBarFill) {
      elements.progressBarFill.style.width = '100%';
    }

    const [{ session, provider }, tokenizer] = await Promise.all([
      createSession(modelBuffer, (msg) => {
        if (elements.progressStatusText) {
          elements.progressStatusText.textContent = msg;
        }
      }),
      getTokenizer((status) => {
        if (elements.progressStatusText) {
          elements.progressStatusText.textContent = status.status;
        }
      }),
    ]);

    state.session = session;
    state.provider = provider;
    state.tokenizer = tokenizer;
    state.isModelReady = true;

    if (elements.webgpuBadge) {
      if (provider === 'webgpu') {
        elements.webgpuBadge.textContent = 'WebGPU';
        elements.webgpuBadge.className = 'badge badge-gpu';
      } else {
        elements.webgpuBadge.textContent = 'WASM';
        elements.webgpuBadge.className = 'badge badge-wasm';
      }
    }

    if (elements.modelBanner) {
      elements.modelBanner.style.display = 'none';
    }
  } catch (err) {
    console.error('[App] Model setup failed:', err);
    showStatus(`Failed to initialize model: ${err.message}`, 'error');
    if (elements.progressStatusText) {
      elements.progressStatusText.textContent = 'Model loading failed.';
    }
  } finally {
    state.isLoadingModel = false;
  }
}

async function handleExtract() {
  hideStatus();

  let rawHtml = '';
  let sourceUrl = '';
  let baseUrl = '';

  if (state.currentMode === 'url') {
    sourceUrl = elements.urlInput.value.trim();
    if (!sourceUrl) {
      showStatus('Please enter a valid webpage URL to extract.', 'error');
      elements.urlInput.focus();
      return;
    }
    baseUrl = sourceUrl;
  } else {
    rawHtml = elements.htmlInput.value.trim();
    baseUrl = elements.baseUrlInput.value.trim();
    if (!rawHtml) {
      showStatus('Please paste HTML content or load a sample preset.', 'error');
      elements.htmlInput.focus();
      return;
    }
  }

  if (!state.isModelReady) {
    showStatus('Loading model into WebGPU before running extraction...', 'info');
    await loadAndPrepareModel();
    if (!state.isModelReady) return;
  }

  const submitBtns = [elements.extractBtn, elements.extractHtmlBtn];
  submitBtns.forEach((b) => {
    b.disabled = true;
    b.querySelector('.btn-text').textContent = 'Extracting...';
  });

  const overallStartTime = performance.now();

  try {
    if (state.currentMode === 'url') {
      showStatus('Fetching webpage content...', 'info');
      const fetchResult = await fetchWebpage(sourceUrl, (msg) => showStatus(msg, 'info'));
      rawHtml = fetchResult.html;
      sourceUrl = fetchResult.finalUrl;
    }

    showStatus('Segmenting HTML into semantic content blocks...', 'info');
    const { doc, blocks, itemIds, title } = simplifyAndExtractBlocks(rawHtml, baseUrl);

    if (blocks.length === 0) {
      throw new Error('No content blocks identified in HTML.');
    }

    showStatus(`Tokenizing ${blocks.length} blocks with EuroBERT tokenizer...`, 'info');
    const blockTokens = tokenizeBlocks(blocks, state.tokenizer);
    const maxTokens = state.provider === 'webgpu' ? 2048 : 1024;
    const chunks = packChunks(blockTokens, { maxTokens });

    const engineName = state.provider === 'webgpu' ? 'WebGPU' : 'WASM';
    showStatus(`Running ${engineName} inference across ${chunks.length} sequence chunk(s)...`, 'info');
    const inferenceResult = await runWebGpuInference(state.session, chunks, itemIds, {
      onProgress: ({ currentChunk, totalChunks }) => {
        showStatus(`${engineName} forward pass: chunk ${currentChunk}/${totalChunks}...`, 'info');
      },
    });

    showStatus('Reconstructing main content and generating Markdown...', 'info');
    const { html: cleanHtml, markdown } = reconstructMainContent(doc, inferenceResult.labels);

    const totalTimeMs = Number((performance.now() - overallStartTime).toFixed(1));

    state.lastResult = {
      title: title || 'Extracted Document',
      sourceUrl,
      cleanHtml,
      markdown,
      blocks,
      itemIds,
      labels: inferenceResult.labels,
      rawPredictions: inferenceResult.rawPredictions,
      inferenceTimeMs: inferenceResult.inferenceTimeMs,
      totalTokens: inferenceResult.totalTokens,
      totalTimeMs,
      provider: state.provider,
    };

    displayResults(state.lastResult, elements);
    hideStatus();
  } catch (err) {
    console.error('[App] Extraction failed:', err);
    showStatus(`Extraction failed: ${err.message}`, 'error');
  } finally {
    submitBtns.forEach((b) => {
      b.disabled = false;
      b.querySelector('.btn-text').textContent = 'Extract';
    });
  }
}

function setupEventListeners() {
  elements.modeUrlTab.addEventListener('click', () => {
    state.currentMode = 'url';
    elements.modeUrlTab.classList.add('active');
    elements.modeHtmlTab.classList.remove('active');
    elements.urlView.style.display = 'block';
    elements.htmlView.style.display = 'none';
  });

  elements.modeHtmlTab.addEventListener('click', () => {
    state.currentMode = 'html';
    elements.modeHtmlTab.classList.add('active');
    elements.modeUrlTab.classList.remove('active');
    elements.htmlView.style.display = 'block';
    elements.urlView.style.display = 'none';
  });

  elements.urlInput.addEventListener('input', () => {
    elements.clearUrlBtn.style.display = elements.urlInput.value ? 'flex' : 'none';
  });

  elements.clearUrlBtn.addEventListener('click', () => {
    elements.urlInput.value = '';
    elements.clearUrlBtn.style.display = 'none';
    elements.urlInput.focus();
  });

  elements.pasteUrlBtn.addEventListener('click', async () => {
    try {
      const text = await navigator.clipboard.readText();
      if (text) {
        elements.urlInput.value = text;
        elements.clearUrlBtn.style.display = 'flex';
      }
    } catch {
      showStatus('Unable to access clipboard. Please paste manually.', 'info');
    }
  });

  elements.clearHtmlBtn.addEventListener('click', () => {
    elements.htmlInput.value = '';
    elements.htmlInput.focus();
  });

  document.querySelectorAll('.preset-btn').forEach((btn) => {
    btn.addEventListener('click', () => {
      const presetKey = btn.dataset.preset;
      const sample = SAMPLE_HTMLS[presetKey];
      if (sample) {
        elements.modeHtmlTab.click();
        elements.htmlInput.value = sample;
        elements.baseUrlInput.value = 'https://usefeyn.com';
        showStatus(`Loaded '${presetKey}' sample HTML document.`, 'info');
      }
    });
  });

  elements.htmlDropZone.addEventListener('dragover', (e) => {
    e.preventDefault();
    elements.dropOverlay.style.display = 'flex';
  });

  elements.htmlDropZone.addEventListener('dragleave', (e) => {
    if (!elements.htmlDropZone.contains(e.relatedTarget)) {
      elements.dropOverlay.style.display = 'none';
    }
  });

  elements.htmlDropZone.addEventListener('drop', (e) => {
    e.preventDefault();
    elements.dropOverlay.style.display = 'none';
    const files = e.dataTransfer?.files;
    if (files && files[0]) {
      const reader = new FileReader();
      reader.onload = () => {
        elements.htmlInput.value = reader.result;
        elements.modeHtmlTab.click();
        showStatus(`Loaded file: ${files[0].name}`, 'info');
      };
      reader.readAsText(files[0]);
    }
  });

  document.querySelectorAll('.output-tab').forEach((tab) => {
    tab.addEventListener('click', () => {
      document.querySelectorAll('.output-tab').forEach((t) => t.classList.remove('active'));
      tab.classList.add('active');
      const tabName = tab.dataset.tab;
      state.activeOutputTab = tabName;

      elements.tabViewPreview.style.display = tabName === 'preview' ? 'block' : 'none';
      elements.tabViewMarkdown.style.display = tabName === 'markdown' ? 'block' : 'none';
      elements.tabViewHtml.style.display = tabName === 'html' ? 'block' : 'none';
      elements.tabViewBlocks.style.display = tabName === 'blocks' ? 'block' : 'none';
    });
  });

  elements.copyBtn.addEventListener('click', async () => {
    if (!state.lastResult) return;
    const content =
      state.activeOutputTab === 'html' ? state.lastResult.cleanHtml : state.lastResult.markdown;
    try {
      await navigator.clipboard.writeText(content);
      elements.copyBtnText.textContent = 'Copied!';
      setTimeout(() => (elements.copyBtnText.textContent = 'Copy'), 2000);
    } catch {
      showStatus('Clipboard copy failed. Please copy manually.', 'info');
    }
  });

  elements.downloadMdBtn.addEventListener('click', () => {
    if (!state.lastResult?.markdown) return;
    const blob = new Blob([state.lastResult.markdown], { type: 'text/markdown;charset=utf-8' });
    const url = URL.createObjectURL(blob);
    const a = document.createElement('a');
    const safeTitle = (state.lastResult.title || 'extracted')
      .toLowerCase()
      .replace(/[^a-z0-9]+/g, '-')
      .replace(/^-|-$/g, '');
    a.href = url;
    a.download = `${safeTitle || 'pulpie-content'}.md`;
    a.click();
    URL.revokeObjectURL(url);
  });

  elements.extractBtn.addEventListener('click', handleExtract);
  elements.extractHtmlBtn.addEventListener('click', handleExtract);

  window.addEventListener('keydown', (e) => {
    if ((e.ctrlKey || e.metaKey) && e.key === 'Enter') {
      e.preventDefault();
      handleExtract();
    }
  });
}

document.addEventListener('DOMContentLoaded', async () => {
  setupEventListeners();
  await initializeSystem();
  loadAndPrepareModel();
});