LexAI/entrypoints/content.ts

import { defineContentScript } from 'wxt/utils/define-content-script';
import { isExtensionValid, safeSendMessage } from '@lib/messaging';

export default defineContentScript({
  matches: ['<all_urls>'],
  main() {
    // Debug logging is stripped in production builds — selection text must never
    // reach the host page's console outside dev.
    const isDev = import.meta.env.COMMAND === 'serve';
    const debugLog = (...args: unknown[]) => {
      if (isDev) console.log(...args);
    };
    const debugWarn = (...args: unknown[]) => {
      if (isDev) console.warn(...args);
    };
    debugLog('LexAI content script loaded');

    // ─── State ───────────────────────────────────────────────────────────────
    let toolbar: HTMLElement | null = null;
    let modal: HTMLElement | null = null;

    // Stored immediately on mouseup — used during Replace (selection is lost by then)
    let selectedText = '';
    let storedStart = -1;
    let storedEnd = -1;
    let storedElement: HTMLTextAreaElement | HTMLInputElement | null = null;
    let storedRange: Range | null = null;

    // Lock flag — prevents mouseup from resetting stored selection while modal is open
    let selectionLocked = false;

    // Writing style — persisted in chrome.storage.local
    let currentStyle = 'Default';
    chrome.storage.local.get(['writingStyle'], (res) => {
      currentStyle = (res.writingStyle as string) || 'Default';
    });

    // ─── Selection capture (called immediately on mouseup) ────────────────────

    function captureSelectionNow(): boolean {
      if (selectionLocked) return true; // selection is locked — don't reset
      const el = document.activeElement;

      // 1) textarea / input — grab positions while they're still valid
      if (el instanceof HTMLTextAreaElement || el instanceof HTMLInputElement) {
        const start = el.selectionStart ?? -1;
        const end = el.selectionEnd ?? -1;
        if (end - start >= 2) {
          const text = el.value.substring(start, end).trim();
          if (text.length > 10) {
            selectedText   = text;
            storedStart    = start;
            storedEnd      = end;
            storedElement  = el;
            storedRange    = null;
            return true;
          }
        }
        return false;
      }

      // 2) contenteditable / regular DOM
      const sel = window.getSelection();
      if (sel && sel.rangeCount > 0) {
        const text = sel.toString().trim();
        if (text.length > 10) {
          selectedText  = text;
          storedRange   = sel.getRangeAt(0).cloneRange(); // CLONE — selection will be lost later
          storedElement = null;
          storedStart   = -1;
          storedEnd     = -1;
          return true;
        }
      }

      return false;
    }

    // ─── Capture selection on toolbar button mousedown ────────────────────────
    // Called from each toolbar button's mousedown — at this point focus has NOT
    // shifted yet (even without e.preventDefault), so selection is guaranteed valid.

    function captureForButton() {
      if (selectionLocked) return;

      const activeEl = document.activeElement;
      if (activeEl instanceof HTMLTextAreaElement || activeEl instanceof HTMLInputElement) {
        const start = activeEl.selectionStart ?? -1;
        const end   = activeEl.selectionEnd   ?? -1;
        if (start >= 0 && end > start) {
          storedStart   = start;
          storedEnd     = end;
          storedElement = activeEl;
          storedRange   = null;
          selectedText  = activeEl.value.substring(start, end);
          debugLog('[LexAI captureForButton] textarea captured', { start, end, text: selectedText.substring(0, 40) });
        }
      } else {
        const sel = window.getSelection();
        if (sel && sel.rangeCount > 0 && sel.toString().trim().length > 0) {
          storedRange   = sel.getRangeAt(0).cloneRange();
          storedElement = null;
          storedStart   = -1;
          storedEnd     = -1;
          selectedText  = sel.toString().trim();
          debugLog('[LexAI captureForButton] DOM range captured', { text: selectedText.substring(0, 40) });
        }
      }
    }

    // ─── Replace using stored state ───────────────────────────────────────────

    function replaceText(newText: string) {
      // textarea / input path
      if (storedElement) {
        const el = storedElement;
        const before = el.value.substring(0, storedStart);
        const after  = el.value.substring(storedEnd);
        el.value = before + newText + after;
        el.setSelectionRange(storedStart, storedStart + newText.length);
        el.dispatchEvent(new Event('input',  { bubbles: true }));
        el.dispatchEvent(new Event('change', { bubbles: true }));
        el.focus();
        return;
      }

      // DOM range path (contenteditable etc.)
      if (storedRange) {
        try {
          storedRange.deleteContents();
          const textNode = document.createTextNode(newText);
          storedRange.insertNode(textNode);

          // move caret after inserted text
          const sel = window.getSelection();
          if (sel) {
            sel.removeAllRanges();
            const newRange = document.createRange();
            newRange.setStartAfter(textNode);
            newRange.collapse(true);
            sel.addRange(newRange);
          }
        } catch (_) {
          document.execCommand('insertText', false, newText);
        } finally {
          storedRange = null;
        }
      }
    }

    // ─── Toolbar ──────────────────────────────────────────────────────────────

    function getToolbarPosition(rect: DOMRect) {
      const TOOLBAR_W = 280;
      const TOOLBAR_H = 40;
      const MARGIN    = 8;
      const scrollX   = window.scrollX;
      const scrollY   = window.scrollY;

      // Center above the selection rect
      let left = rect.left + scrollX + rect.width / 2 - TOOLBAR_W / 2;
      let top  = rect.top  + scrollY - TOOLBAR_H - MARGIN;

      // Clamp horizontally to viewport
      left = Math.max(scrollX + MARGIN, Math.min(left, scrollX + window.innerWidth - TOOLBAR_W - MARGIN));

      // Flip below if not enough space above
      if (top < scrollY + MARGIN) {
        top = rect.bottom + scrollY + MARGIN;
      }

      return { top, left };
    }

    // Inject shared keyframe CSS once
    function ensureLexAIStyles() {
      if (document.getElementById('lexai-styles')) return;
      const style = document.createElement('style');
      style.id = 'lexai-styles';
      style.textContent = `
        @keyframes lexai-spin { to { transform: rotate(360deg); } }
        @keyframes lexai-fadein { from { opacity: 0; transform: translateY(4px); } to { opacity: 1; transform: translateY(0); } }
      `;
      document.head.appendChild(style);
    }

    function showToolbar(rect: DOMRect) {
      ensureLexAIStyles();
      hideToolbar();

      const { top, left } = getToolbarPosition(rect);

      toolbar = document.createElement('div');
      toolbar.id = 'lexai-toolbar';
      toolbar.setAttribute('data-lexai', 'true');

      Object.assign(toolbar.style, {
        position:    'absolute',
        top:         `${top}px`,
        left:        `${left}px`,
        zIndex:      '2147483647',
        background:  'linear-gradient(135deg, #1e1e2e 0%, #181825 100%)',
        borderRadius:'10px',
        padding:     '6px 8px',
        display:     'flex',
        gap:         '4px',
        boxShadow:   '0 4px 24px rgba(0,0,0,0.45), 0 1px 3px rgba(0,0,0,0.3)',
        border:      '1px solid rgba(205,214,244,0.12)',
        fontFamily:  '-apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif',
        alignItems:  'center',
      });

      const actions: { label: string; action: string; color: string }[] = [
        { label: '✓ Fix',      action: 'grammar',  color: '#a6e3a1' },
        { label: '↺ Rephrase', action: 'rephrase', color: '#89b4fa' },
        { label: '↓ Shorten',  action: 'shorten',  color: '#fab387' },
        { label: '↑ Expand',   action: 'expand',   color: '#cba6f7' },
        { label: '💡 Explain', action: 'explain',  color: '#2dd4bf' },
      ];

      actions.forEach(({ label, action, color }) => {
        const btn = document.createElement('button');
        btn.textContent = label;
        btn.setAttribute('data-lexai', 'true');

        Object.assign(btn.style, {
          background:    'rgba(49,50,68,0.8)',
          color:         color,
          border:        '1px solid rgba(205,214,244,0.1)',
          borderRadius:  '7px',
          padding:       '4px 10px',
          fontSize:      '12px',
          fontWeight:    '600',
          cursor:        'pointer',
          transition:    'all 0.15s ease',
          letterSpacing: '0.01em',
          whiteSpace:    'nowrap',
        });

        btn.addEventListener('mouseenter', () => {
          btn.style.background   = 'rgba(69,71,90,0.9)';
          btn.style.borderColor  = color + '50';
          btn.style.transform    = 'translateY(-1px)';
        });
        btn.addEventListener('mouseleave', () => {
          btn.style.background   = 'rgba(49,50,68,0.8)';
          btn.style.borderColor  = 'rgba(205,214,244,0.1)';
          btn.style.transform    = 'translateY(0)';
        });
        btn.addEventListener('mousedown', (e) => {
          e.preventDefault();       // prevent focus shift away from textarea
          captureForButton();       // re-capture selection while it's guaranteed valid
        });
        btn.addEventListener('click', (e) => {
          e.stopPropagation();
          runAction(action);
        });

        toolbar!.appendChild(btn);
      });

      // ─── More button (Copy As + Download) ────────────────────────────────────
      const moreBtn = document.createElement('button');
      moreBtn.textContent = '⋯ More';
      moreBtn.setAttribute('data-lexai', 'true');
      Object.assign(moreBtn.style, {
        background:   'rgba(49,50,68,0.8)',
        color:        '#a6adc8',
        border:       '1px solid rgba(205,214,244,0.15)',
        borderRadius: '6px',
        padding:      '5px 10px',
        fontSize:     '12px',
        fontWeight:   '600',
        cursor:       'pointer',
        whiteSpace:   'nowrap',
      });

      // Dropdown menu
      let moreMenu: HTMLElement | null = null;

      function closeMoveMenu() {
        if (moreMenu) { moreMenu.remove(); moreMenu = null; }
      }

      moreBtn.addEventListener('mousedown', (e) => e.preventDefault());
      moreBtn.addEventListener('click', (e) => {
        e.stopPropagation();
        if (moreMenu) { closeMoveMenu(); return; }

        moreMenu = document.createElement('div');
        moreMenu.setAttribute('data-lexai', 'true');
        Object.assign(moreMenu.style, {
          position:      'fixed',
          background:    'rgba(30,30,46,0.98)',
          border:        '1px solid rgba(205,214,244,0.15)',
          borderRadius:  '8px',
          padding:       '4px',
          zIndex:        '2147483647',
          minWidth:      '140px',
          boxShadow:     '0 4px 20px rgba(0,0,0,0.4)',
          display:       'flex',
          flexDirection: 'column',
          gap:           '2px',
        });

        // Position above the More button
        const rect = moreBtn.getBoundingClientRect();
        moreMenu.style.left      = rect.left + 'px';
        moreMenu.style.top       = (rect.top - 8) + 'px';
        moreMenu.style.transform = 'translateY(-100%)';

        const menuItems = [
          { label: '⎘ Copy As', action: 'copy-as' },
          { label: '⬇ Download', action: 'download' },
        ];

        menuItems.forEach(({ label, action }) => {
          const item = document.createElement('button');
          item.textContent = label;
          item.setAttribute('data-lexai', 'true');
          Object.assign(item.style, {
            background:   'transparent',
            color:        '#cdd6f4',
            border:       'none',
            borderRadius: '6px',
            padding:      '7px 12px',
            fontSize:     '13px',
            cursor:       'pointer',
            textAlign:    'left',
            width:        '100%',
          });
          item.addEventListener('mouseenter', () => { item.style.background = 'rgba(137,180,250,0.1)'; });
          item.addEventListener('mouseleave', () => { item.style.background = 'transparent'; });
          item.addEventListener('mousedown', (e) => e.preventDefault());
          item.addEventListener('click', () => {
            closeMoveMenu();
            captureForButton(); // ensure selection is captured
            if (action === 'download') {
              downloadAsText(selectedText);
            } else {
              showCopyAsModal(selectedText);
            }
          });
          moreMenu!.appendChild(item);
        });

        document.body.appendChild(moreMenu);
      });

      toolbar!.appendChild(moreBtn);

      // Close dropdown on outside click / scroll
      document.addEventListener('click', closeMoveMenu);
      document.addEventListener('scroll', closeMoveMenu, { passive: true });

      document.body.appendChild(toolbar);
    }

    function hideToolbar() {
      if (toolbar) {
        toolbar.remove();
        toolbar = null;
      }
    }

    // Extension context guard: isExtensionValid is imported from ~/lib/messaging.

    function showErrorToast(msg: string) {
      const toast = document.createElement('div');
      toast.style.cssText = `
        position: fixed; bottom: 20px; right: 20px; z-index: 999999;
        background: #f38ba8; color: #1e1e2e; padding: 10px 16px;
        border-radius: 8px; font-family: system-ui, sans-serif;
        font-size: 13px; box-shadow: 0 4px 12px rgba(0,0,0,0.3);
      `;
      toast.textContent = msg;
      document.body.appendChild(toast);
      setTimeout(() => toast.remove(), 4000);
    }

    function showToast(msg: string, isError = false) {
      const toast = document.createElement('div');
      toast.setAttribute('data-lexai', 'true');
      toast.style.cssText = `
        position: fixed; bottom: 20px; right: 20px; z-index: 999999;
        background: ${isError ? '#f38ba8' : '#a6e3a1'}; color: #1e1e2e;
        padding: 10px 16px; border-radius: 8px;
        font-family: system-ui, sans-serif; font-size: 13px; font-weight: 600;
        box-shadow: 0 4px 12px rgba(0,0,0,0.3);
        animation: lexai-fadein 0.15s ease;
      `;
      toast.textContent = msg;
      document.body.appendChild(toast);
      setTimeout(() => toast.remove(), 2500);
    }

    // ─── LEXAI-20: Download as .txt ───────────────────────────────────────────

    function downloadAsText(text: string) {
      const timestamp = new Date().toISOString().replace(/[:.]/g, '-').slice(0, 19);
      const filename = `lexai-note-${timestamp}.txt`;
      const blob = new Blob([text], { type: 'text/plain' });
      const url = URL.createObjectURL(blob);
      const a = document.createElement('a');
      a.href = url;
      a.download = filename;
      document.body.appendChild(a);
      a.click();
      document.body.removeChild(a);
      URL.revokeObjectURL(url);
      showToast(`Downloaded: ${filename}`);
    }

    // ─── LEXAI-19: Copy As modal ──────────────────────────────────────────────

    function showCopyAsModal(text: string) {
      // Remove any existing copy-as modal
      document.getElementById('lexai-copyAs-overlay')?.remove();

      const formats = ['JSON', 'Markdown', 'Bullet List', 'Numbered List', 'CSV', 'HTML'];

      const overlay = document.createElement('div');
      overlay.id = 'lexai-copyAs-overlay';
      overlay.setAttribute('data-lexai', 'true');
      Object.assign(overlay.style, {
        position:   'fixed',
        inset:      '0',
        zIndex:     '2147483646',
        background: 'rgba(0,0,0,0.25)',
      });

      const picker = document.createElement('div');
      picker.setAttribute('data-lexai', 'true');
      Object.assign(picker.style, {
        position:    'fixed',
        top:         '50%',
        left:        '50%',
        transform:   'translate(-50%,-50%)',
        zIndex:      '2147483647',
        background:  'linear-gradient(135deg, #1e1e2e 0%, #181825 100%)',
        borderRadius:'12px',
        padding:     '16px',
        width:       '280px',
        boxShadow:   '0 12px 40px rgba(0,0,0,0.55)',
        border:      '1px solid rgba(205,214,244,0.15)',
        fontFamily:  '-apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif',
      });

      // Header row
      const header = document.createElement('div');
      Object.assign(header.style, {
        display:        'flex',
        justifyContent: 'space-between',
        alignItems:     'center',
        marginBottom:   '12px',
      });
      const title = document.createElement('div');
      title.innerHTML = '⚡ <strong style="color:#cdd6f4">Copy As</strong>';
      Object.assign(title.style, { fontSize: '12px', color: '#89b4fa' });

      const closeX = document.createElement('button');
      closeX.textContent = '✕';
      Object.assign(closeX.style, {
        background: 'none', border: 'none', color: '#6c7086',
        cursor: 'pointer', fontSize: '14px', padding: '0 2px', lineHeight: '1',
      });
      closeX.addEventListener('click', () => overlay.remove());
      header.appendChild(title);
      header.appendChild(closeX);
      picker.appendChild(header);

      // Spinner slot (shown while loading)
      const spinnerSlot = document.createElement('div');
      spinnerSlot.setAttribute('data-lexai', 'true');
      Object.assign(spinnerSlot.style, {
        display:        'none',
        justifyContent: 'center',
        alignItems:     'center',
        padding:        '12px 0',
        color:          '#a6adc8',
        fontSize:       '13px',
        gap:            '8px',
      });
      spinnerSlot.innerHTML = '<span style="font-size:18px;animation:lexai-spin 0.8s linear infinite;display:inline-block">⟳</span> Formatting…';
      picker.appendChild(spinnerSlot);

      // 2-column grid of format buttons
      const grid = document.createElement('div');
      Object.assign(grid.style, {
        display:             'grid',
        gridTemplateColumns: '1fr 1fr',
        gap:                 '6px',
      });

      formats.forEach((fmt) => {
        const btn = document.createElement('button');
        btn.textContent = fmt;
        btn.setAttribute('data-lexai', 'true');
        Object.assign(btn.style, {
          background:   'rgba(49,50,68,0.8)',
          color:        '#cdd6f4',
          border:       '1px solid rgba(205,214,244,0.12)',
          borderRadius: '8px',
          padding:      '8px 10px',
          fontSize:     '12px',
          fontWeight:   '600',
          cursor:       'pointer',
          transition:   'all 0.12s ease',
          textAlign:    'center',
        });
        btn.addEventListener('mouseenter', () => {
          btn.style.background  = 'rgba(69,71,90,0.9)';
          btn.style.borderColor = '#89b4fa50';
        });
        btn.addEventListener('mouseleave', () => {
          btn.style.background  = 'rgba(49,50,68,0.8)';
          btn.style.borderColor = 'rgba(205,214,244,0.12)';
        });
        btn.addEventListener('click', async () => {
          // Show spinner, hide grid
          grid.style.display       = 'none';
          spinnerSlot.style.display = 'flex';

          const response = await sendToBackground({
            type:   'COPY_AS',
            text,
            format: fmt,
          }) as { error?: string; result?: string } | null;

          overlay.remove();

          if (!response || response.error) {
            showToast(`❌ ${response?.error ?? 'Unknown error'}`, true);
            return;
          }

          try {
            await navigator.clipboard.writeText(response.result ?? '');
            showToast(`✓ Copied as ${fmt}!`);
          } catch {
            showToast('❌ Clipboard write failed', true);
          }
        });
        grid.appendChild(btn);
      });

      picker.appendChild(grid);
      overlay.appendChild(picker);

      // Close on overlay click (outside picker)
      overlay.addEventListener('click', (e) => {
        if (e.target === overlay) overlay.remove();
      });

      document.body.appendChild(overlay);
    }

    // ─── Safe chrome.runtime.sendMessage wrapper ──────────────────────────────
    // Shared implementation; on a stale extension context we toast and clean up.

    async function sendToBackground(payload: Record<string, unknown>): Promise<unknown> {
      return safeSendMessage(payload, () => {
        showErrorToast('LexAI was updated — please refresh this page.');
        hideToolbar();
      });
    }

    // ─── LLM call ─────────────────────────────────────────────────────────────

    async function runAction(action: string) {
      if (!selectedText) return;
      const textToProcess = selectedText;

      // SNAPSHOT selection state RIGHT NOW before anything async happens.
      // These are passed directly into showModal so the Replace handler has
      // guaranteed access — no closure-scope ambiguity.
      const snapStart   = storedStart;
      const snapEnd     = storedEnd;
      const snapElement = storedElement;
      const snapRange   = storedRange ? storedRange.cloneRange() : null;

      debugLog('[LexAI runAction] snapshot captured', {
        snapStart,
        snapEnd,
        snapElement: snapElement?.tagName,
        snapElementId: snapElement?.id || snapElement?.name || '(no id)',
        hasSnapRange: !!snapRange,
        selectedText: textToProcess.substring(0, 50),
      });

      selectionLocked = true; // lock stored selection — modal interaction won't reset it

      // Guard: extension may have been reloaded/updated
      if (!isExtensionValid()) {
        hideToolbar();
        showErrorToast('LexAI was updated. Please refresh the page.');
        return;
      }

      // Show loading in toolbar
      if (toolbar) {
        toolbar.innerHTML = '';
        const loading = document.createElement('span');
        loading.textContent = '⏳ LexAI thinking…';
        loading.setAttribute('data-lexai', 'true');
        Object.assign(loading.style, {
          color:    '#a6adc8',
          fontSize: '12px',
          padding:  '4px 10px',
        });
        toolbar.appendChild(loading);
      }

      try {
        const response = await sendToBackground({
          type:    'ANALYZE_TEXT',
          payload: { text: textToProcess, action, style: currentStyle },
        }) as { error?: string; result?: string } | null;

        hideToolbar();

        if (response === null) return; // sendToBackground already handled the error

        if (response?.error) {
          showModal(`❌ ${response.error}`, null, snapStart, snapEnd, snapElement, snapRange, action, textToProcess);
        } else {
          // Tone analysis is display-only — no Replace button (pass null for originalText)
          const isExplainAction = action === 'explain';
          showModal(
            response?.result ?? '(no result)',
            isExplainAction ? null : textToProcess,
            snapStart, snapEnd, snapElement, snapRange,
            action, textToProcess,
          );
        }
      } catch (err) {
        hideToolbar();
        if (String(err).includes('Extension context invalidated') ||
            String(err).includes('message channel closed')) {
          showErrorToast('LexAI was updated — please refresh this page.');
        } else {
          showModal(`❌ Error: ${String(err)}`, null, snapStart, snapEnd, snapElement, snapRange, action, textToProcess);
        }
      }
    }

    // ─── Result Modal ─────────────────────────────────────────────────────────
    // snapStart/snapEnd/snapElement/snapRange are passed explicitly from runAction
    // so the Replace handler has guaranteed in-scope access to them.

    function showModal(
      resultText:  string,
      originalText: string | null,
      snapStart:   number,
      snapEnd:     number,
      snapElement: HTMLTextAreaElement | HTMLInputElement | null,
      snapRange:   Range | null,
      action?:     string,
      textForRegenerate?: string,
    ) {
      if (modal) modal.remove();

      // Overlay
      const overlay = document.createElement('div');
      overlay.setAttribute('data-lexai', 'true');
      Object.assign(overlay.style, {
        position:       'fixed',
        inset:          '0',
        zIndex:         '2147483646',
        background:     'rgba(0,0,0,0.3)',
        backdropFilter: 'blur(2px)',
      });

      modal = document.createElement('div');
      modal.setAttribute('data-lexai', 'true');
      Object.assign(modal.style, {
        position:   'fixed',
        top:        '50%',
        left:       '50%',
        transform:  'translate(-50%,-50%)',
        zIndex:     '2147483647',
        background: 'linear-gradient(135deg, #1e1e2e 0%, #181825 100%)',
        borderRadius:'14px',
        padding:    '20px 22px',
        width:      '420px',
        maxWidth:   'min(90vw, 420px)',
        maxHeight:  '70vh',
        overflowY:  'auto',
        boxShadow:  '0 12px 48px rgba(0,0,0,0.6)',
        border:     '1px solid rgba(205,214,244,0.15)',
        fontFamily: '-apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif',
        color:      '#cdd6f4',
      });

      // Header
      const header = document.createElement('div');
      Object.assign(header.style, {
        display:       'flex',
        justifyContent:'space-between',
        alignItems:    'center',
        marginBottom:  '12px',
      });

      const title = document.createElement('div');
      title.innerHTML = '⚡ <strong>LexAI</strong> Suggestion';
      Object.assign(title.style, { fontSize: '13px', color: '#89b4fa' });

      const closeBtn = document.createElement('button');
      closeBtn.textContent = '✕';
      Object.assign(closeBtn.style, {
        background: 'none',
        border:     'none',
        color:      '#6c7086',
        cursor:     'pointer',
        fontSize:   '16px',
        padding:    '0 4px',
        lineHeight: '1',
      });
      closeBtn.addEventListener('click', () => {
        selectionLocked = false;
        overlay.remove();
        modal!.remove();
        modal = null;
      });

      header.appendChild(title);
      header.appendChild(closeBtn);
      modal.appendChild(header);

      // Result text
      const body = document.createElement('div');
      body.textContent = resultText;
      Object.assign(body.style, {
        fontSize:   '14px',
        lineHeight: '1.65',
        color:      '#cdd6f4',
        background: 'rgba(49,50,68,0.5)',
        borderRadius:'8px',
        padding:    '12px 14px',
        marginBottom:'14px',
        whiteSpace: 'pre-wrap',
        wordBreak:  'break-word',
      });
      modal.appendChild(body);

      // ─── Style selector row ───────────────────────────────────────────────
      const STYLES = ['Default', 'Formal', 'Casual', 'Academic', 'Creative', 'Concise'];

      const styleRow = document.createElement('div');
      Object.assign(styleRow.style, {
        display:     'flex',
        alignItems:  'center',
        gap:         '8px',
        marginBottom:'12px',
      });

      const styleLabel = document.createElement('span');
      styleLabel.textContent = 'Style:';
      Object.assign(styleLabel.style, {
        fontSize: '12px',
        color:    '#a6adc8',
        flexShrink: '0',
      });

      const styleSelect = document.createElement('select');
      Object.assign(styleSelect.style, {
        background:   'rgba(49,50,68,0.9)',
        color:        '#cdd6f4',
        border:       '1px solid rgba(205,214,244,0.2)',
        borderRadius: '6px',
        padding:      '4px 8px',
        fontSize:     '12px',
        cursor:       'pointer',
        flex:         '1',
        outline:      'none',
      });

      STYLES.forEach(s => {
        const opt = document.createElement('option');
        opt.value = s;
        opt.textContent = s;
        styleSelect.appendChild(opt);
      });

      // Load saved style and pre-select
      chrome.storage.local.get(['writingStyle'], (res) => {
        const saved = (res.writingStyle as string) || 'Default';
        styleSelect.value = saved;
        currentStyle = saved;
      });

      const regenBtn = document.createElement('button');
      regenBtn.textContent = '↺ Regenerate';
      regenBtn.setAttribute('data-lexai', 'true');
      Object.assign(regenBtn.style, {
        background:   'rgba(137,180,250,0.15)',
        color:        '#89b4fa',
        border:       '1px solid rgba(137,180,250,0.3)',
        borderRadius: '6px',
        padding:      '5px 10px',
        fontSize:     '12px',
        fontWeight:   '600',
        cursor:       'pointer',
        flexShrink:   '0',
        whiteSpace:   'nowrap',
      });

      async function doRegenerate() {
        if (!action || !textForRegenerate) return;
        const chosenStyle = styleSelect.value;
        // Save style
        currentStyle = chosenStyle;
        await chrome.storage.local.set({ writingStyle: chosenStyle });
        // Show spinner
        body.innerHTML = '<span style="font-size:18px;animation:lexai-spin 0.8s linear infinite;display:inline-block">⟳</span> Regenerating…';
        regenBtn.disabled = true;
        styleSelect.disabled = true;

        const response = await sendToBackground({
          type:    'ANALYZE_TEXT',
          payload: { text: textForRegenerate, action, style: chosenStyle },
        }) as { error?: string; result?: string } | null;

        regenBtn.disabled = false;
        styleSelect.disabled = false;

        if (!response || response.error) {
          body.textContent = `❌ ${response?.error ?? 'Unknown error'}`;
        } else {
          body.textContent = response.result ?? '(no result)';
        }
      }

      styleSelect.addEventListener('change', () => doRegenerate());
      regenBtn.addEventListener('click', () => doRegenerate());

      styleRow.appendChild(styleLabel);
      styleRow.appendChild(styleSelect);
      styleRow.appendChild(regenBtn);
      modal.appendChild(styleRow);

      // Buttons
      const btnRow = document.createElement('div');
      Object.assign(btnRow.style, { display: 'flex', gap: '8px' });

      if (originalText !== null) {
        const replaceBtn = document.createElement('button');
        replaceBtn.textContent = '↩ Replace';
        replaceBtn.setAttribute('data-lexai', 'true');
        Object.assign(replaceBtn.style, {
          background:  '#89b4fa',
          color:       '#1e1e2e',
          border:      'none',
          borderRadius:'8px',
          padding:     '8px 18px',
          fontSize:    '13px',
          fontWeight:  '700',
          cursor:      'pointer',
          flex:        '1',
        });
        replaceBtn.addEventListener('mousedown', (e) => e.preventDefault());
        replaceBtn.addEventListener('click', () => {
          // Debug log — confirms exactly what state we have at replace time
          debugLog('[LexAI Replace]', {
            snapStart,
            snapEnd,
            snapElement: snapElement?.tagName ?? 'null',
            snapElementValue: snapElement?.value?.substring(0, 50) ?? 'N/A',
            hasSnapRange: !!snapRange,
            resultText: resultText.substring(0, 50),
          });

          if (snapElement && snapStart >= 0 && snapEnd >= 0 && snapEnd > snapStart) {
            // textarea / input path — use the snapshot indices directly
            const before = snapElement.value.substring(0, snapStart);
            const after  = snapElement.value.substring(snapEnd);
            debugLog('[LexAI Replace] before:', JSON.stringify(before.substring(0, 30)));
            debugLog('[LexAI Replace] after:', JSON.stringify(after.substring(0, 30)));
            snapElement.value = before + resultText + after;
            snapElement.setSelectionRange(snapStart, snapStart + resultText.length);
            snapElement.dispatchEvent(new Event('input',  { bubbles: true }));
            snapElement.dispatchEvent(new Event('change', { bubbles: true }));
            snapElement.focus();
          } else if (snapRange) {
            // contenteditable / DOM range path
            // Strategy: restore the original selection, then just "type" the new text.
            // The browser automatically replaces selected text on insert — no manual
            // deleteContents() needed (and safer since stale range deletes can go wrong).
            debugLog('[LexAI Replace] using DOM range — restoring selection + inserting');
            try {
              const sel = window.getSelection();
              if (sel) {
                sel.removeAllRanges();
                sel.addRange(snapRange); // restore original highlighted selection
              }
              document.execCommand('insertText', false, resultText); // replaces selection
            } catch (err) {
              debugWarn('[LexAI Replace] execCommand failed:', err);
            }
          } else {
            debugWarn('[LexAI Replace] NO VALID SNAP STATE — snapStart:', snapStart, 'snapEnd:', snapEnd, 'snapElement:', snapElement, 'snapRange:', snapRange);
          }

          selectionLocked = false;
          overlay.remove();
          modal!.remove();
          modal = null;
        });
        btnRow.appendChild(replaceBtn);
      }

      const copyBtn = document.createElement('button');
      copyBtn.textContent = '⎘ Copy';
      Object.assign(copyBtn.style, {
        background:  'rgba(49,50,68,0.8)',
        color:       '#cdd6f4',
        border:      '1px solid rgba(205,214,244,0.1)',
        borderRadius:'8px',
        padding:     '8px 18px',
        fontSize:    '13px',
        fontWeight:  '600',
        cursor:      'pointer',
      });
      copyBtn.addEventListener('click', () => {
        navigator.clipboard.writeText(resultText).then(() => {
          copyBtn.textContent = '✓ Copied!';
          setTimeout(() => { copyBtn.textContent = '⎘ Copy'; }, 1500);
        });
      });
      btnRow.appendChild(copyBtn);

      const dismissBtn = document.createElement('button');
      dismissBtn.textContent = 'Dismiss';
      Object.assign(dismissBtn.style, {
        background:  'rgba(49,50,68,0.8)',
        color:       '#6c7086',
        border:      '1px solid rgba(205,214,244,0.1)',
        borderRadius:'8px',
        padding:     '8px 14px',
        fontSize:    '13px',
        cursor:      'pointer',
      });
      dismissBtn.addEventListener('click', () => {
        selectionLocked = false; // unlock on dismiss
        overlay.remove();
        modal!.remove();
        modal = null;
      });
      btnRow.appendChild(dismissBtn);

      modal.appendChild(btnRow);

      // Close overlay on click
      overlay.addEventListener('click', () => {
        selectionLocked = false;
        overlay.remove();
        modal!.remove();
        modal = null;
      });

      document.body.appendChild(overlay);
      document.body.appendChild(modal);
    }

    // ─── Global extension context error handlers ──────────────────────────────

    window.addEventListener('error', (e) => {
      if (e.message?.includes('Extension context invalidated')) {
        e.preventDefault(); // suppress console error
        hideToolbar();
        showErrorToast('LexAI was updated — please refresh this page.');
      }
    });

    window.addEventListener('unhandledrejection', (e) => {
      if (String(e.reason)?.includes('Extension context invalidated') ||
          String(e.reason)?.includes('message channel closed')) {
        e.preventDefault();
        hideToolbar();
        showErrorToast('LexAI was updated — please refresh this page.');
      }
    });

    // ─── Event listeners ──────────────────────────────────────────────────────

    document.addEventListener('mouseup', (e) => {
      const target = e.target as Element;

      // Don't trigger on our own UI
      if (target?.closest?.('[data-lexai="true"]')) return;

      // Capture selection state IMMEDIATELY — positions are valid right now.
      // The setTimeout below lets the browser finalise the selection before we read it.
      const captured = captureSelectionNow();

      setTimeout(() => {
        if (!captured || !selectedText || selectedText.length <= 10) {
          hideToolbar();
          return;
        }

        // Get bounding rect from the live selection (still valid inside setTimeout)
        let rect: DOMRect | null = null;
        const sel = window.getSelection();
        if (sel && sel.rangeCount > 0) {
          rect = sel.getRangeAt(0).getBoundingClientRect();
        } else if (storedElement) {
          // fallback: use the element's bounding rect for textarea/input
          rect = storedElement.getBoundingClientRect();
        }

        if (rect && rect.width > 0) {
          showToolbar(rect);
        } else {
          hideToolbar();
        }
      }, 10);
    });

    // Hide toolbar when clicking elsewhere (not on our UI)
    document.addEventListener('mousedown', (e) => {
      const target = e.target as Element;
      if (!target?.closest?.('[data-lexai="true"]')) {
        hideToolbar();
      }
    });

    // Hide toolbar on scroll
    document.addEventListener('scroll', () => {
      hideToolbar();
    }, { passive: true });

    // Keyboard: Escape closes both
    document.addEventListener('keydown', (e) => {
      if (e.key === 'Escape') {
        hideToolbar();
        if (modal) {
          selectionLocked = false;
          modal.remove();
          modal = null;
          document.querySelectorAll('[data-lexai="true"]').forEach(el => el.remove());
        }
      }
    });

    // ─── Context menu trigger from background ─────────────────────────────────
    chrome.runtime.onMessage.addListener((message) => {
      if (message.type === 'lexai-context-menu') {
        selectedText = message.text as string;
        currentStyle = (message.style as string) || 'Default';
        // For context menu, we have no DOM selection positions — zero them out
        storedStart   = -1;
        storedEnd     = -1;
        storedElement = null;
        storedRange   = null;
        runAction(message.action as string);
      }
    });
  },
});