LexAI/entrypoints/content.ts

import { defineContentScript } from 'wxt/utils/define-content-script';

export default defineContentScript({
  matches: ['<all_urls>'],
  main() {
    console.log('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;

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

    function captureSelectionNow(): boolean {
      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;
    }

    // ─── 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 };
    }

    function showToolbar(rect: DOMRect) {
      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' },
      ];

      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()); // don't lose selection
        btn.addEventListener('click', (e) => {
          e.stopPropagation();
          runAction(action);
        });

        toolbar!.appendChild(btn);
      });

      document.body.appendChild(toolbar);
    }

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

    // ─── Extension context guard ──────────────────────────────────────────────

    function isExtensionValid(): boolean {
      try {
        return typeof chrome !== 'undefined' && !!chrome.runtime?.id;
      } catch {
        return false;
      }
    }

    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);
    }

    // ─── Safe chrome.runtime.sendMessage wrapper ──────────────────────────────

    async function safeSendMessage(payload: Record<string, unknown>): Promise<unknown> {
      if (!isExtensionValid()) {
        showErrorToast('LexAI was updated — please refresh this page.');
        return null;
      }
      try {
        return await chrome.runtime.sendMessage(payload);
      } catch (err) {
        if (String(err).includes('Extension context invalidated') ||
            String(err).includes('message channel closed')) {
          showErrorToast('LexAI was updated — please refresh this page.');
          hideToolbar();
        }
        return null;
      }
    }

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

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

      // 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 safeSendMessage({
          type:    'ANALYZE_TEXT',
          payload: { text: textToProcess, action },
        }) as { error?: string; result?: string } | null;

        hideToolbar();

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

        if (response?.error) {
          showModal(`❌ ${response.error}`, null);
        } else {
          showModal(response?.result ?? '(no result)', 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);
        }
      }
    }

    // ─── Result Modal ─────────────────────────────────────────────────────────

    function showModal(resultText: string, originalText: string | null) {
      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', () => {
        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);

      // 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()); // keep stored selection intact
        replaceBtn.addEventListener('click', () => {
          replaceText(resultText);
          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', () => {
        overlay.remove();
        modal!.remove();
        modal = null;
      });
      btnRow.appendChild(dismissBtn);

      modal.appendChild(btnRow);

      // Close overlay on click
      overlay.addEventListener('click', () => {
        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) {
          modal.remove();
          modal = null;
          document.querySelectorAll('[data-lexai="true"]').forEach(el => el.remove());
        }
      }
    });
  },
});