From 652fca9d8727bef2233a9b689d11eb97d0185bd0 Mon Sep 17 00:00:00 2001
From: Beingpax <prakashjoshipaxtek@gmail.com>
Date: Thu, 29 May 2025 18:09:48 +0545
Subject: [PATCH] Improve screen capture timing and logic for power mode
 switching

---
 VoiceInk/PowerMode/ActiveWindowService.swift | 19 +++++++++++++++++++
 VoiceInk/Services/AIEnhancementService.swift |  6 ------
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/VoiceInk/PowerMode/ActiveWindowService.swift b/VoiceInk/PowerMode/ActiveWindowService.swift
index b4fdce9..4e32256 100644
--- a/VoiceInk/PowerMode/ActiveWindowService.swift
+++ b/VoiceInk/PowerMode/ActiveWindowService.swift
@@ -72,6 +72,13 @@ class ActiveWindowService: ObservableObject {
     func applyConfiguration(_ config: PowerModeConfig) async {
         guard let enhancementService = enhancementService else { return }
         
+        // Capture current state before making changes
+        let wasScreenCaptureEnabled = await MainActor.run { 
+            enhancementService.useScreenCaptureContext 
+        }
+        let wasEnhancementEnabled = await MainActor.run { 
+            enhancementService.isEnhancementEnabled 
+        }
 
         await MainActor.run {
             enhancementService.isEnhancementEnabled = config.isAIEnhancementEnabled
@@ -126,5 +133,17 @@ class ActiveWindowService: ObservableObject {
                 }
             }
         }
+        
+        // Capture screen context at the end after a brief delay to ensure all changes are settled
+        // and either enhancement is newly enabled or screen capture is newly enabled
+        if config.isAIEnhancementEnabled && config.useScreenCapture {
+            let shouldCaptureScreen = !wasEnhancementEnabled || !wasScreenCaptureEnabled
+            
+            if shouldCaptureScreen {
+                // Wait a moment for UI changes and content loading to complete
+                try? await Task.sleep(nanoseconds: 1_500_000_000) // 1.5 seconds
+                await enhancementService.captureScreenContext()
+            }
+        }
     }
 } 
diff --git a/VoiceInk/Services/AIEnhancementService.swift b/VoiceInk/Services/AIEnhancementService.swift
index cded27b..33470c7 100644
--- a/VoiceInk/Services/AIEnhancementService.swift
+++ b/VoiceInk/Services/AIEnhancementService.swift
@@ -20,12 +20,6 @@ class AIEnhancementService: ObservableObject {
             if isEnhancementEnabled && selectedPromptId == nil {
                 selectedPromptId = customPrompts.first?.id
             }
-            
-            if isEnhancementEnabled && useScreenCaptureContext {
-                Task {
-                    await captureScreenContext()
-                }
-            }
         }
     }        
     @Published var useClipboardContext: Bool {