fix: capture selection on toolbar button mousedown - definitive replace fix

Root cause (dual bugs):
1. SCOPE BUG: snapStart/snapEnd/snapElement/snapRange were local vars inside
   runAction() but showModal() is a sibling function — those snap vars were
   never in scope inside the Replace button handler. The condition always
   silently fell through, doing nothing.

2. TIMING BUG: Selection was only captured on document mouseup, which could
   race against focus changes when clicking the toolbar.

Fix:
- showModal() now accepts snapStart/snapEnd/snapElement/snapRange as explicit
  parameters — guaranteed in-scope, no closure ambiguity
- Each toolbar button mousedown calls captureForButton() BEFORE any focus
  shift can occur (belt-and-suspenders, works even if e.preventDefault fails)
- Added console.log debug lines to Replace handler and runAction snapshot
- selectionLocked reset on Escape and overlay click (was missing)

entrypoints/content.ts

@@ -60,6 +60,38 @@ export default defineContentScript({
       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);
+          console.log('[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();
+          console.log('[LexAI captureForButton] DOM range captured', { text: selectedText.substring(0, 40) });
+        }
+      }
+    }
+
     // ─── Replace using stored state ───────────────────────────────────────────
 
     function replaceText(newText: string) {
@@ -185,7 +217,10 @@ export default defineContentScript({
           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('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);
@@ -251,6 +286,24 @@ export default defineContentScript({
     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;
+
+      console.log('[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
@@ -285,9 +338,9 @@ export default defineContentScript({
         if (response === null) return; // safeSendMessage already handled the error
 
         if (response?.error) {
-          showModal(`❌ ${response.error}`, null);
+          showModal(`❌ ${response.error}`, null, snapStart, snapEnd, snapElement, snapRange);
         } else {
-          showModal(response?.result ?? '(no result)', textToProcess);
+          showModal(response?.result ?? '(no result)', textToProcess, snapStart, snapEnd, snapElement, snapRange);
         }
       } catch (err) {
         hideToolbar();
@@ -295,14 +348,23 @@ export default defineContentScript({
             String(err).includes('message channel closed')) {
           showErrorToast('LexAI was updated — please refresh this page.');
         } else {
-          showModal(`❌ Error: ${String(err)}`, null);
+          showModal(`❌ Error: ${String(err)}`, null, snapStart, snapEnd, snapElement, snapRange);
         }
       }
     }
 
     // ─── 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) {
+    function showModal(
+      resultText:  string,
+      originalText: string | null,
+      snapStart:   number,
+      snapEnd:     number,
+      snapElement: HTMLTextAreaElement | HTMLInputElement | null,
+      snapRange:   Range | null,
+    ) {
       if (modal) modal.remove();
 
       // Overlay
@@ -362,6 +424,7 @@ export default defineContentScript({
         lineHeight: '1',
       });
       closeBtn.addEventListener('click', () => {
+        selectionLocked = false;
         overlay.remove();
         modal!.remove();
         modal = null;
@@ -406,10 +469,52 @@ export default defineContentScript({
           cursor:      'pointer',
           flex:        '1',
         });
-        replaceBtn.addEventListener('mousedown', (e) => e.preventDefault()); // keep stored selection intact
+        replaceBtn.addEventListener('mousedown', (e) => e.preventDefault());
         replaceBtn.addEventListener('click', () => {
-          replaceText(resultText);
-          selectionLocked = false; // unlock after replace
+          // Debug log — confirms exactly what state we have at replace time
+          console.log('[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);
+            console.log('[LexAI Replace] before:', JSON.stringify(before.substring(0, 30)));
+            console.log('[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
+            console.log('[LexAI Replace] using DOM range');
+            try {
+              snapRange.deleteContents();
+              const textNode = document.createTextNode(resultText);
+              snapRange.insertNode(textNode);
+              const sel = window.getSelection();
+              if (sel) {
+                sel.removeAllRanges();
+                const newRange = document.createRange();
+                newRange.setStartAfter(textNode);
+                newRange.collapse(true);
+                sel.addRange(newRange);
+              }
+            } catch (err) {
+              console.warn('[LexAI Replace] range replace failed:', err);
+            }
+          } else {
+            console.warn('[LexAI Replace] NO VALID SNAP STATE — snapStart:', snapStart, 'snapEnd:', snapEnd, 'snapElement:', snapElement, 'snapRange:', snapRange);
+          }
+
+          selectionLocked = false;
           overlay.remove();
           modal!.remove();
           modal = null;
@@ -460,6 +565,7 @@ export default defineContentScript({
 
       // Close overlay on click
       overlay.addEventListener('click', () => {
+        selectionLocked = false;
         overlay.remove();
         modal!.remove();
         modal = null;
@@ -542,6 +648,7 @@ export default defineContentScript({
       if (e.key === 'Escape') {
         hideToolbar();
         if (modal) {
+          selectionLocked = false;
           modal.remove();
           modal = null;
           document.querySelectorAll('[data-lexai="true"]').forEach(el => el.remove());