api_key_dialog.py

from PySide6.QtWidgets import QDialog, QVBoxLayout, QHBoxLayout, QLabel, QLineEdit, QPushButton, QMessageBox
from PySide6.QtCore import Qt

class ApiKeyDialog(QDialog):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowTitle("Google API Key Required")
        self.setModal(True) # Block other windows until this is handled
        self.setMinimumWidth(400)

        self.api_key = None
        self.should_save_key = False

        layout = QVBoxLayout(self)

        instruction_text = (
            "No GOOGLE_API_KEY environment variable found, and no API key was found in the local config.\n"
            "Please enter your Google API Key to use Gemini features.\n"
            "You can obtain an API key from Google AI Studio."
        )
        # This text might be slightly adjusted later by the caller if a saved key was found but invalid.
        self.instruction_label = QLabel(instruction_text)
        self.instruction_label.setWordWrap(True)
        layout.addWidget(self.instruction_label)

        self.key_input = QLineEdit()
        self.key_input.setPlaceholderText("Enter your Google API Key here")
        self.key_input.setEchoMode(QLineEdit.Password)
        layout.addWidget(self.key_input)

        self.save_key_checkbox = QCheckBox("Save API Key for future sessions")
        self.save_key_checkbox.setToolTip(
            "If checked, the API key will be stored unencrypted in a local configuration file.\n"
            "This is convenient but less secure if others have access to your computer."
        )
        layout.addWidget(self.save_key_checkbox)

        # Buttons
        button_layout = QHBoxLayout()
        self.ok_button = QPushButton("OK")
        self.ok_button.clicked.connect(self.accept_input)

        self.cancel_button = QPushButton("Cancel")
        self.cancel_button.clicked.connect(self.reject) # QDialog's reject slot

        button_layout.addStretch()
        button_layout.addWidget(self.ok_button)
        button_layout.addWidget(self.cancel_button)
        layout.addLayout(button_layout)

        self.setLayout(layout)

    def accept_input(self):
        entered_key = self.key_input.text().strip()
        if not entered_key:
            QMessageBox.warning(self, "API Key Missing", "Please enter an API key or press Cancel.")
            return

        # Basic validation (presence check) - more robust validation happens when Gemini tries to use it.
        if len(entered_key) < 30: # Typical Google API keys are longer
             if QMessageBox.question(self, "Possible Invalid Key",
                                   "The entered key seems short. Are you sure it's correct?",
                                   QMessageBox.Yes | QMessageBox.No, QMessageBox.No) == QMessageBox.No:
                return

        self.api_key = entered_key
        self.should_save_key = self.save_key_checkbox.isChecked()

        if self.should_save_key:
            # Re-confirm with a more explicit warning if they chose to save
            confirm_save = QMessageBox.warning(
                self,
                "Security Warning",
                "Saving your API key locally in an unencrypted file can be a security risk "
                "if others have access to your computer or its files.\n\n"
                "Are you sure you want to save the API key for future sessions?",
                QMessageBox.Yes | QMessageBox.Cancel,
                QMessageBox.Cancel
            )
            if confirm_save == QMessageBox.Cancel:
                self.should_save_key = False # User cancelled saving
                # Do not proceed with accept() yet, let them uncheck or reconsider
                return

        self.accept() # QDialog's accept slot

    def get_api_key_and_save_preference(parent=None, current_key_invalid=False):
        # Static method to create, show the dialog, and return the key and save preference
        dialog = ApiKeyDialog(parent)
        if current_key_invalid:
            dialog.instruction_label.setText(
                "The previously used/saved API Key is invalid or caused an error.\n"
                "Please enter a valid Google API Key.\n"
                "You can obtain an API key from Google AI Studio."
            )
            dialog.key_input.setPlaceholderText("Enter new valid API Key")

        if dialog.exec() == QDialog.Accepted:
            return dialog.api_key, dialog.should_save_key
        return None, False

if __name__ == '__main__':
    # Test the dialog
    app = QApplication([])

    print("Simulating API key needed...")
    key = ApiKeyDialog.get_api_key()

    if key:
        print(f"API Key entered: {key}")
    else:
        print("API Key entry was cancelled.")

    # Example of direct instantiation if needed
    # dialog_instance = ApiKeyDialog()
    # result = dialog_instance.exec()
    # if result == QDialog.Accepted:
    #     print(f"API Key from instance: {dialog_instance.api_key}")
    # else:
    #     print("Dialog cancelled (instance).")

    sys.exit(0)