auto_toggle.py

import sys
import json
import logging
from PyQt5.QtWidgets import (QApplication, QSystemTrayIcon, QMenu, QMainWindow, 
                            QLabel, QStyle, QVBoxLayout, QWidget, QPushButton, QHBoxLayout)
from PyQt5.QtCore import Qt, QTimer
from PyQt5.QtGui import QIcon
import win32gui
import win32process
import psutil
import cohere
import ctypes
from ctypes import wintypes
from pathlib import Path
import keyboard
import time
import win32api
import win32con
import pyautogui
pyautogui.FAILSAFE = False  # Disable fail-safe
pyautogui.PAUSE = 0.01  # Set default delay between actions

import pyperclip

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class SuggestionWindow(QMainWindow):
    def __init__(self, config):
        super().__init__()
        self.setWindowFlags(Qt.FramelessWindowHint | Qt.WindowStaysOnTopHint | Qt.Tool)
        self.setAttribute(Qt.WA_TranslucentBackground)
        self.config = config  # Store config
        self.setWindowFlags(Qt.FramelessWindowHint | Qt.WindowStaysOnTopHint | Qt.Tool)
        
        self.setStyleSheet("""
            QMainWindow {
                background: transparent;
            }
        """)
        
        # Create central widget
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QVBoxLayout(central_widget)
        
        # Create suggestion container
        suggestion_container = QWidget()
        suggestion_layout = QVBoxLayout(suggestion_container)
        
        # Create suggestion label
        self.label = QLabel()
        self.label.setStyleSheet("""
            QLabel {
                background-color: rgba(32, 32, 32, 0.95);
                color: white;
                padding: 16px;
                border-radius: 12px;
                font-size: 14px;
                font-family: 'Segoe UI', sans-serif;
                min-width: 250px;
                max-width: 400px;
                margin: 8px;
                border: 1px solid rgba(255, 255, 255, 0.1);
            }
        """)
        suggestion_layout.addWidget(self.label)
        
        # Create buttons container
        buttons_container = QWidget()
        buttons_layout = QHBoxLayout(buttons_container)
        buttons_layout.setSpacing(8)
        
        # Create accept button
        self.accept_btn = QPushButton("✓")
        self.accept_btn.setFixedSize(32, 32)
        self.accept_btn.setStyleSheet("""
            QPushButton {
                background-color: rgba(40, 167, 69, 0.95);
                color: white;
                border: none;
                border-radius: 16px;
                font-family: 'Segoe UI', sans-serif;
                font-size: 16px;
                font-weight: bold;
            }
            QPushButton:hover {
                background-color: rgba(50, 177, 79, 0.95);
                transform: scale(1.1);
            }
        """)
        
        # Create reject button
        self.reject_btn = QPushButton("✕")
        self.reject_btn.setFixedSize(32, 32)
        self.reject_btn.setStyleSheet("""
            QPushButton {
                background-color: rgba(220, 53, 69, 0.95);
                color: white;
                border: none;
                border-radius: 16px;
                font-family: 'Segoe UI', sans-serif;
                font-size: 16px;
                font-weight: bold;
            }
            QPushButton:hover {
                background-color: rgba(230, 63, 79, 0.95);
                transform: scale(1.1);
            }
        """)
        
        # Add buttons to layout
        buttons_layout.addWidget(self.accept_btn)
        buttons_layout.addWidget(self.reject_btn)
        suggestion_layout.addWidget(buttons_container)
        
        # Add suggestion container to main layout
        main_layout.addWidget(suggestion_container)
        
        # Create shortcut hint label
        self.shortcut_hint = QLabel("Tab: Accept • Esc: Dismiss")
        self.shortcut_hint.setStyleSheet("""
            QLabel {
                color: rgba(255, 255, 255, 0.7);
                                         

                font-size: 11px;
                font-family: 'Segoe UI', sans-serif;
                padding: 4px;
            }
        """)
        main_layout.addWidget(self.shortcut_hint)
        
        # Connect button signals
        self.accept_btn.clicked.connect(self.accept_suggestion)
        self.reject_btn.clicked.connect(self.hide)
        
        # Initialize suggestion text
        self.current_suggestion = ""
        
    def show_suggestion(self, text):
        self.current_suggestion = text
        self.label.setText(text)
        self.adjustSize()
        
        # Position relative to cursor with offset
        cursor_pos = win32gui.GetCursorPos()
        screen = QApplication.primaryScreen().geometry()
        
        # Calculate position to ensure window stays on screen
        x = min(cursor_pos[0] + 20, screen.width() - self.width() - 20)
        y = min(cursor_pos[1] + 20, screen.height() - self.height() - 20)
        
        # Add fade-in effect
        self.setWindowOpacity(0)
        self.move(x, y)
        self.show()
        
        # Fade in animation
        for i in range(10):
            self.setWindowOpacity(i/10)
            QApplication.processEvents()
            time.sleep(0.01)
        
        # Auto-hide after 20 seconds (20000 milliseconds)
        QTimer.singleShot(20000, self.fade_out)

    def fade_out(self):
        # Fade out animation
        for i in range(10, -1, -1):
            self.setWindowOpacity(i/10)
            QApplication.processEvents()
            time.sleep(0.01)
        self.hide()
        

    def accept_suggestion(self):
        if not self.current_suggestion:
            return
        try:
            # Store current mouse position
            old_position = pyautogui.position()
            
            # Type the suggestion directly
            pyautogui.typewrite(self.current_suggestion)
            print(self.current_suggestion)
            print("1")
            pyautogui.press('enter')
            # Restore mouse position
            pyautogui.moveTo(old_position[0], old_position[1])
            
            logger.info(f"Successfully typed suggestion: {self.current_suggestion}")
            
        except Exception as e:
            logger.error(f"PyAutoGUI typing failed: {e}")
        finally:
            self.hide()


    def _insert_using_win32_sendmessage(self):
        """Method 1: Using SendMessage"""
        try:
            hwnd = win32gui.GetForegroundWindow()
            for char in self.current_suggestion:
                win32gui.SendMessage(hwnd, win32con.WM_CHAR, ord(char), 0)
                time.sleep(0.001)  # Tiny delay between characters
            return True
        except:
            return False



    def _insert_using_clipboard_fallback(self):
        """Method 3: Using pyperclip"""
        try:
            original = pyperclip.paste()
            pyperclip.copy(self.current_suggestion)
            time.sleep(0.05)
            
            keyboard.press_and_release('ctrl+v')
            time.sleep(0.1)
            
            pyperclip.copy(original)
            return True
        except:
            return False

    def _insert_using_keyboard_simulation(self):
        """Method 4: Direct keyboard simulation"""
        try:
            keyboard.write(self.current_suggestion)
            return True
        except:
            return False

    def keyPressEvent(self, event):
        if event.key() == Qt.Key_Tab:
            self.accept_suggestion()
            event.accept()  # Prevent Tab from moving focus
        elif event.key() == Qt.Key_Escape:
            self.hide()
            event.accept()
class KeyboardHook:
    def __init__(self):
        try:
            # Load DLL
            dll_path = Path(__file__).parent / "native" / "keyboard_hook.dll"
            if not dll_path.exists():
                raise FileNotFoundError(f"DLL not found: {dll_path}")
            
            self.dll = ctypes.WinDLL(str(dll_path))
            
            # Define function prototypes
            self.dll.Initialize.restype = wintypes.BOOL
            self.dll.Initialize.argtypes = []
            
            self.dll.Cleanup.restype = wintypes.BOOL
            self.dll.Cleanup.argtypes = []
            
            self.dll.GetNextKey.restype = wintypes.BOOL
            self.dll.GetNextKey.argtypes = [wintypes.LPWSTR]
            
            # Initialize the hook
            if not self.dll.Initialize():
                raise RuntimeError("Failed to initialize keyboard hook")
                
            logger.info("Keyboard hook initialized successfully")
            
        except Exception as e:
            logger.error(f"Failed to initialize keyboard hook: {e}")
            raise
    
    def get_next_event(self):
        try:
            buffer = ctypes.create_unicode_buffer(2)
            if self.dll.GetNextKey(buffer):
                return buffer.value
            return None
        except Exception as e:
            logger.error(f"Error getting keyboard event: {e}")
            return None
    
    def cleanup(self):
        try:
            if hasattr(self, 'dll'):
                self.dll.Cleanup()
                logger.info("Keyboard hook cleaned up")
        except Exception as e:
            logger.error(f"Error cleaning up keyboard hook: {e}")

class WritingAssistant(QMainWindow):
    def __init__(self):
        super().__init__()
        
        # Initialize basic attributes
        self.text_buffer = []
        self.is_active = True
        self.auto_suggest_enabled = True  # Can be toggled
        
        # Load configuration first
        try:
            with open('config.json', 'r') as f:
                self.config = json.load(f)
                logger.info("Configuration loaded")
                self.cohere_client = cohere.Client(self.config['cohere_api_key'])
        except Exception as e:
            logger.error(f"Failed to load config: {e}")
            sys.exit(1)
        
        # Initialize components after config is loaded
        try:
            self.keyboard_hook = KeyboardHook()
            self.suggestion_window = SuggestionWindow(config=self.config)
            
            # Setup system tray
            self.setup_tray()
            
            # Setup timers
            self.poll_timer = QTimer()
            self.poll_timer.timeout.connect(self.poll_keyboard)
            self.poll_timer.start(50)  # Poll every 50ms
            
            self.suggestion_timer = QTimer()
            self.suggestion_timer.setSingleShot(True)
            self.suggestion_timer.timeout.connect(self.process_buffer)
            
            logger.info("Writing Assistant initialized")
            
        except Exception as e:
            logger.error(f"Failed to initialize components: {e}")
            sys.exit(1)

    def setup_tray(self):
        self.tray = QSystemTrayIcon()
        icon = QApplication.style().standardIcon(QStyle.SP_ComputerIcon)
        self.tray.setIcon(icon)
        
        menu = QMenu()
        # Add auto-suggest toggle
        self.auto_suggest_action = menu.addAction("Auto-Suggest")
        self.auto_suggest_action.setCheckable(True)
        self.auto_suggest_action.setChecked(True)
        self.auto_suggest_action.triggered.connect(self.toggle_auto_suggest)
        
        toggle_action = menu.addAction("Enable/Disable")
        toggle_action.triggered.connect(self.toggle)
        exit_action = menu.addAction("Exit")
        exit_action.triggered.connect(self.quit)
        
        self.tray.setContextMenu(menu)
        self.tray.show()
        logger.info("System tray initialized")

    def toggle_auto_suggest(self):
        self.auto_suggest_enabled = not self.auto_suggest_enabled
        status = "enabled" if self.auto_suggest_enabled else "disabled"
        logger.info(f"Auto-suggest {status}")
        self.tray.showMessage("Writing Assistant", f"Auto-suggest {status}")

    def poll_keyboard(self):
        if not self.is_active:
            return
            
        char = self.keyboard_hook.get_next_event()
        if char and char.isprintable():
            self.text_buffer.append(char)
            self.suggestion_timer.start(
                int(self.config.get('suggestion_delay', 1.0) * 1000)
            )

    def process_buffer(self):
        if not self.text_buffer or not self.auto_suggest_enabled:
            return
            
        text = ''.join(self.text_buffer)
        
        # Get current active window for context
        hwnd = win32gui.GetForegroundWindow()
        window_title = win32gui.GetWindowText(hwnd)
        
        # Build context-aware prompt
        prompt = f"""Based on the context in {window_title}:
        Previous text: {text}
        
        Provide a single natural completion that:
        1. Matches the writing style
        2. Completes the current thought
        3. Is concise and relevant
        
        Complete this: {text}"""

        try:
            response = self.cohere_client.generate(
                model='command',
                prompt=prompt,
                max_tokens=30,
                temperature=0.7,
                k=1,
                stop_sequences=["\n", ".", "!", "?"]
            )
            
            suggestion = response.generations[0].text.strip()
            if suggestion:
                logger.info(f"Auto-inserting suggestion: {suggestion}")
                self._auto_insert_text(suggestion)
                
        except Exception as e:
            logger.error(f"Error getting suggestions: {e}")
        finally:
            # Clear buffer after certain size or time
            if len(self.text_buffer) > self.config.get('context_window', 100):
                self.text_buffer = self.text_buffer[-100:]

    def _auto_insert_text(self, text):
        """Automatically insert text at cursor position"""
        try:
            # Store current mouse position
            old_position = pyautogui.position()
            
            # Type the suggestion directly
            pyautogui.typewrite(text)
            
            # Restore mouse position
            pyautogui.moveTo(old_position[0], old_position[1])
            
            logger.info(f"Auto-inserted text: {text}")
            
        except Exception as e:
            logger.error(f"Auto-insertion failed: {e}")

    def toggle(self):
        self.is_active = not self.is_active
        status = "enabled" if self.is_active else "disabled"
        logger.info(f"Assistant {status}")
        self.tray.showMessage("Writing Assistant", f"Assistant {status}")

    def quit(self):
        logger.info("Shutting down...")
        self.keyboard_hook.cleanup()
        self.tray.hide()
        QApplication.quit()

def main():
    try:
        app = QApplication(sys.argv)
        assistant = WritingAssistant()
        assistant.hide()  # Hide main window
        return app.exec_()
    except Exception as e:
        logger.error(f"Application error: {e}")
        return 1

if __name__ == "__main__":
    sys.exit(main())