quick_ask_overlay.py

import sys
from PySide6.QtWidgets import QWidget, QVBoxLayout, QLineEdit, QLabel, QApplication, QTextEdit, QPushButton, QHBoxLayout, QScrollArea, QMessageBox
from PySide6.QtCore import Qt, Signal, Slot, QTimer, QRect
from PySide6.QtGui import QKeySequence, QShortcut, QPainter, QColor, QBrush, QPen, QFontMetrics, QPixmap, QTextCursor, QIcon
# from context_analyzer import analyze_text_for_contextual_items # No longer used here
from typing import List, Dict, Optional

# Import RegionSelectorWidget locally if it's in the same directory, or adjust import path
try:
    from region_selector_widget import RegionSelectorWidget
except ImportError:
    # Fallback or error handling if region_selector_widget is not found
    print("Warning: region_selector_widget.py not found, region selection will not work.")
    RegionSelectorWidget = None


class InsightOverlay(QWidget):
    query_submitted = Signal(str)
    request_close = Signal()
    action_button_clicked_signal = Signal(str, object) # Changed object for value
    region_selection_initiated = Signal()


    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowFlags(Qt.FramelessWindowHint | Qt.WindowStaysOnTopHint | Qt.Tool)
        self.setAttribute(Qt.WA_TranslucentBackground)
        self.setStyleSheet("""
            InsightOverlay {
                background-color: rgba(40, 42, 54, 0.95);
                border-radius: 18px; color: #f8f8f2;
            }
            QPushButton#ActionButton {
                background-color: #44475a; color: #f8f8f2;
                border: 1px solid #6272a4; border-radius: 5px;
                padding: 3px 8px; font-size: 9pt; margin-right: 4px;
            } QPushButton#ActionButton:hover { background-color: #6272a4; }
        """)

        self.main_layout = QVBoxLayout(self)
        self.main_layout.setContentsMargins(12, 12, 12, 12); self.main_layout.setSpacing(8)

        self.context_preview_widget = QWidget()
        self.context_preview_layout = QHBoxLayout(self.context_preview_widget)
        self.context_preview_layout.setContentsMargins(0,0,0,0); self.context_preview_layout.setSpacing(5)
        self.context_image_preview_label = QLabel(alignment=Qt.AlignCenter, styleSheet="border: 1px solid #6272a4; background-color: #282a36;")
        self.context_image_preview_label.setFixedSize(50, 50)
        self.context_ocr_snippet_label = QLabel("OCR: N/A", styleSheet="color: #bd93f9; font-size: 9pt;", wordWrap=True)
        self.context_preview_layout.addWidget(self.context_image_preview_label)
        self.context_preview_layout.addWidget(self.context_ocr_snippet_label, 1)
        self.context_preview_widget.hide()
        self.main_layout.addWidget(self.context_preview_widget)

        self.action_buttons_scroll_area = QScrollArea()
        self.action_buttons_scroll_area.setWidgetResizable(True); self.action_buttons_scroll_area.setFixedHeight(40)
        self.action_buttons_scroll_area.setHorizontalScrollBarPolicy(Qt.ScrollBarAsNeeded)
        self.action_buttons_scroll_area.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
        self.action_buttons_scroll_area.setStyleSheet("QScrollArea { border: none; background-color: transparent; }")
        self.action_buttons_widget_container = QWidget()
        self.action_buttons_layout = QHBoxLayout(self.action_buttons_widget_container)
        self.action_buttons_layout.setContentsMargins(0, 0, 0, 0); self.action_buttons_layout.setSpacing(6)
        self.action_buttons_layout.setAlignment(Qt.AlignLeft)
        self.action_buttons_scroll_area.setWidget(self.action_buttons_widget_container)
        self.action_buttons_scroll_area.hide()
        self.main_layout.addWidget(self.action_buttons_scroll_area)

        self.chat_display = QTextEdit(readOnly=True, placeholderText="Ask Gemini anything...")
        self.chat_display.setStyleSheet("QTextEdit {background-color: rgba(68,71,90,0.7); border: 1px solid #6272a4; border-radius:10px; padding:8px; color:#f8f8f2; font-size:11pt;}")
        self.chat_display.setMinimumHeight(80); self.chat_display.setMaximumHeight(300) # Adjusted max height
        self.main_layout.addWidget(self.chat_display)

        input_area_widget = QWidget(); input_area_layout = QHBoxLayout(input_area_widget)
        input_area_layout.setContentsMargins(0,0,0,0); input_area_layout.setSpacing(6)

        self.select_region_button = QPushButton("Region")
        self.select_region_button.setStyleSheet("QPushButton {background-color:#6272a4; color:#f8f8f2; border-radius:10px; padding:5px 10px; font-size:9pt;} QPushButton:hover {background-color:#7082b6;}")
        self.select_region_button.setToolTip("Select a screen region for new context (Ctrl+Shift+R - if enabled).")
        self.select_region_button.clicked.connect(self.initiate_region_selection_from_overlay)
        input_area_layout.addWidget(self.select_region_button)

        self.query_input = QLineEdit(placeholderText="Type your message...")
        self.query_input.setStyleSheet("QLineEdit {background-color:#282a36; border:1px solid #6272a4; border-radius:15px; padding:10px 15px; font-size:11pt; color:#f8f8f2;} QLineEdit:focus {border:1px solid #50fa7b;}")
        self.query_input.returnPressed.connect(self.submit_current_query)
        input_area_layout.addWidget(self.query_input, 1)
        self.send_button = QPushButton("Send", styleSheet="QPushButton {background-color:#50fa7b; color:#282a36; border-radius:15px; padding:10px 15px; font-size:11pt; font-weight:bold;} QPushButton:hover {background-color:#47e070;} QPushButton:pressed {background-color:#3fc963;}")
        self.send_button.clicked.connect(self.submit_current_query)
        input_area_layout.addWidget(self.send_button)
        self.main_layout.addWidget(input_area_widget)

        self.setLayout(self.main_layout)
        self.setMinimumWidth(450); self.setMaximumWidth(650); self.resize(550, 280) # Default size
        self.escape_shortcut = QShortcut(QKeySequence(Qt.Key_Escape), self)
        self.escape_shortcut.activated.connect(self.hide_overlay)

        self.conversation_history = []; self.context_image_path = None
        self.context_ocr_text = None; self.suggested_actions: List[Dict] = []
        # self.region_selector = None # RegionSelectorWidget is now handled by MainApp

    def set_initial_context(self, image_path: Optional[str], ocr_text: Optional[str]):
        self.context_image_path = image_path; self.context_ocr_text = ocr_text
        self._clear_action_buttons()
        self.conversation_history = []
        self.chat_display.clear()

        if image_path:
            pixmap = QPixmap(image_path)
            if not pixmap.isNull(): self.context_image_preview_label.setPixmap(pixmap.scaled(self.context_image_preview_label.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation))
            else: self.context_image_preview_label.setText("Err") # Error loading image
            self.context_image_preview_label.show()
        else: self.context_image_preview_label.hide(); self.context_image_preview_label.clear()

        ocr_display_text = "OCR: N/A"
        if ocr_text and ocr_text.strip():
            snippet = ocr_text.strip()[:80] + "..." if len(ocr_text.strip()) > 80 else ocr_text.strip()
            ocr_display_text = f"OCR: {snippet}"
        self.context_ocr_snippet_label.setText(ocr_display_text)

        if image_path or (ocr_text and ocr_text.strip()): self.context_preview_widget.show()
        else: self.context_preview_widget.hide()

        # This message is shown while MainApp fetches actions from Gemini
        self.add_message_to_display("System", "<i>Analyzing screen content for actions...</i>")

    def update_suggested_actions(self, actions_data: List[Dict]):
        # This method is called by MainApp AFTER Gemini responds with actions.
        # The "Analyzing..." message should be removed.
        self.remove_last_message_from_display() # Removes "Analyzing..."
        self.suggested_actions = actions_data
        self._populate_action_buttons()
        if not actions_data :
             self.add_message_to_display("System", "<i>No specific actions suggested. Ask a question or try capturing a new region.</i>")


    def _clear_action_buttons(self):
        while self.action_buttons_layout.count():
            child = self.action_buttons_layout.takeAt(0)
            if child.widget(): child.widget().deleteLater()
        self.action_buttons_scroll_area.hide()

    def _populate_action_buttons(self):
        self._clear_action_buttons()
        buttons_added = False
        if not self.suggested_actions: self.action_buttons_scroll_area.hide(); return

        for action_item in self.suggested_actions:
            label = action_item.get("label", "Action")
            action_type = action_item.get("type", "unknown")
            action_value = action_item.get("value", "") # Value can be string, list, etc.
            button_text = label[:25] + "..." if len(label) > 25 else label
            button = QPushButton(button_text)
            # Store potentially complex action_value. QPushButton properties are limited.
            # A common way is to use a lambda with functools.partial or store index/lookup.
            # For simplicity, if value is simple (str, int), property is fine.
            # If value is list/dict, it might not store/retrieve correctly.
            # For now, assume value is simple enough or MainApp handles complex lookup.
            button.setProperty("action_type", action_type)
            button.setProperty("action_value", action_value)
            button.clicked.connect(self.handle_action_button_clicked)
            button.setObjectName("ActionButton")
            tooltip_value_str = str(action_value)
            if len(tooltip_value_str) > 100: tooltip_value_str = tooltip_value_str[:100] + "..."
            button.setToolTip(f"Type: {action_type.capitalize()}\nValue: {tooltip_value_str}\nLabel: {label}")
            self.action_buttons_layout.addWidget(button)
            buttons_added = True

        if buttons_added: self.action_buttons_scroll_area.show()
        else: self.action_buttons_scroll_area.hide()

    @Slot()
    def handle_action_button_clicked(self):
        sender_button = self.sender()
        if sender_button:
            action_type = sender_button.property("action_type")
            action_value = sender_button.property("action_value") # This retrieves it as QVariant, convert as needed
            print(f"[InsightOverlay] Action button clicked: Type='{action_type}', Value='{action_value}' (type: {type(action_value)})")
            self.action_button_clicked_signal.emit(action_type, action_value) # action_value is QVariant here

    @Slot()
    def initiate_region_selection_from_overlay(self):
        if RegionSelectorWidget is None:
            # Use QMessageBox for user feedback if component is missing
            msg_box = QMessageBox(self)
            msg_box.setIcon(QMessageBox.Warning)
            msg_box.setText("Region selection feature is not available.")
            msg_box.setInformativeText("The necessary component (RegionSelectorWidget) could not be loaded.")
            msg_box.setWindowTitle("Feature Unavailable")
            msg_box.setStandardButtons(QMessageBox.Ok)
            msg_box.exec_()
            return
        self.region_selection_initiated.emit()

    def submit_current_query(self):
        query_text = self.query_input.text().strip()
        if query_text:
            self.add_message_to_display("You", query_text)
            self.query_submitted.emit(query_text)
            self.query_input.clear()

    def add_message_to_display(self, sender: str, message: str):
        # Basic HTML escaping for messages to prevent accidental HTML rendering issues
        # More robust escaping might be needed if messages can contain complex HTML-like structures.
        escaped_message = message.replace('&', '&amp;').replace('<', '&lt;').replace('>', '&gt;')

        # Add bold tags for sender, and allow basic HTML like <i> for system messages
        if sender.lower() == "system": # Allow italics for system messages
            self.chat_display.append(f"<b>{sender}:</b> {message}") # Assume system messages are safe or pre-formatted
        else:
            self.chat_display.append(f"<b>{sender}:</b> {escaped_message}")

        self.conversation_history.append({"sender": sender, "message": message}) # Store original message
        self.chat_display.verticalScrollBar().setValue(self.chat_display.verticalScrollBar().maximum()) # Auto-scroll


    def receive_gemini_response(self, response_text: str):
        self.add_message_to_display("Gemini", response_text) # Gemini responses can contain markdown/HTML for citations

    def remove_last_message_from_display(self):
        # This is a bit fragile. It assumes the message to remove is the very last block.
        cursor = self.chat_display.textCursor()
        cursor.movePosition(QTextCursor.End)
        cursor.select(QTextCursor.BlockUnderCursor) # Selects the last line/block

        # Check if the selection is not empty before removing
        if cursor.hasSelection() and cursor.selectedText().strip(): # Check if the selected block has content
            cursor.removeSelectedText()
            # After removing, check if the new last line is empty and remove it too (often an extra newline)
            cursor.movePosition(QTextCursor.End)
            cursor.select(QTextCursor.BlockUnderCursor)
            if not cursor.selectedText().strip(): # If the new last block is empty
                 cursor.removeSelectedText()

        # Ensure cursor is at the end after modification
        cursor.movePosition(QTextCursor.End)
        self.chat_display.setTextCursor(cursor)


    def show_overlay(self):
        # Context and chat are now set by set_initial_context, called by MainApp
        self.center_on_screen();
        self.query_input.setFocus();
        self.query_input.clear()
        self.show(); self.activateWindow()

    def hide_overlay(self):
        self.hide()
        # Clear context and history when hiding, to ensure fresh state next time
        self.context_image_path = None; self.context_ocr_text = None; self.suggested_actions = []
        if self.context_image_preview_label: self.context_image_preview_label.clear()
        if self.context_ocr_snippet_label: self.context_ocr_snippet_label.setText("OCR: N/A")
        if self.context_preview_widget: self.context_preview_widget.hide()
        self._clear_action_buttons()
        if self.chat_display: self.chat_display.clear()
        self.conversation_history = []
        self.request_close.emit() # Signal MainApp that it was closed

    def center_on_screen(self):
        try:
            app_instance = QApplication.instance()
            if app_instance:
                screen = app_instance.primaryScreen()
                if screen:
                    screen_geo = screen.geometry()
                    # Position it a bit higher than center, more like top-third
                    self.move(int((screen_geo.width()-self.width())/2), int((screen_geo.height()-self.height())/3.5))
                    return
            print("Warning: Could not get screen geometry for centering (primaryScreen not available).")
        except AttributeError:
            print("Warning: Could not get screen geometry for centering (AttributeError). Might be running headless.")

if __name__ == '__main__':
    app = QApplication(sys.argv)
    # Dummy main window for testing the overlay standalone
    main_win_for_testing = QWidget(styleSheet="background-color: #222;") # Dark background
    main_win_for_testing.setFixedSize(1000,700)

    overlay_instance = InsightOverlay()

    def test_query_submission(query_text_from_overlay):
        print(f"Overlay query submitted: {query_text_from_overlay}")
        overlay_instance.add_message_to_display("Gemini", "<i>Thinking about your query...</i>")
        # Simulate a delay and then a response
        QTimer.singleShot(1200, lambda: (
            overlay_instance.remove_last_message_from_display(), # Remove "Thinking..."
            overlay_instance.receive_gemini_response(f"This is a simulated response to: '{query_text_from_overlay}'.")
        ))

    def test_action_button_handler(action_type, action_value):
        print(f"Action received from overlay: Type='{action_type}', Value='{action_value}' (Value Python type: {type(action_value)})")
        overlay_instance.add_message_to_display("System", f"Simulating action: {action_type} with value: {action_value}")

    def test_region_selection_request_handler():
        print("InsightOverlay requested region selection!")
        # Simulate a region selection and then update actions
        QTimer.singleShot(100, lambda: overlay_instance.update_suggested_actions([
            {"label":"Region Test Action", "type":"test_region", "value":"Simulated Region Data"}
        ]))

    overlay_instance.query_submitted.connect(test_query_submission)
    overlay_instance.action_button_clicked_signal.connect(test_action_button_handler)
    overlay_instance.region_selection_initiated.connect(test_region_selection_request_handler)

    # Simulate initial context being set (as MainApp would do)
    # For this test, let's provide some example OCR text
    example_ocr = "Example document text. Contact support at help@example.com or visit https://www.example.com for more info. Call 123-456-7890."
    overlay_instance.set_initial_context(None, example_ocr) # No image, just OCR for this test

    # Simulate Gemini providing actions (as MainApp would do after analysis)
    # This call would happen in MainApp's handle_gemini_finished after Gemini responds to context analysis.
    QTimer.singleShot(500, lambda: overlay_instance.update_suggested_actions([
        {"label": "Email Help", "type": "email", "value": "help@example.com"},
        {"label": "Open Example.com", "type": "url", "value": "https://www.example.com"},
        {"label": "Copy Phone", "type": "copy_text", "value": "123-456-7890"},
        {"label": "Search OCR", "type": "search_web_for_text", "value": example_ocr[:50]}
    ]))

    # Button in the dummy main window to show the overlay
    show_overlay_button = QPushButton("Show Insight Overlay", main_win_for_testing)
    show_overlay_button.clicked.connect(overlay_instance.show_overlay)

    dummy_layout = QHBoxLayout(main_win_for_testing)
    dummy_layout.addStretch()
    dummy_layout.addWidget(show_overlay_button)
    dummy_layout.addStretch()
    main_win_for_testing.setLayout(dummy_layout)

    main_win_for_testing.show()
    sys.exit(app.exec())