forwarder.py

import logging
from logging.handlers import RotatingFileHandler
import os
import ipaddress
import webbrowser
import socket
import sys
import time


from PyQt5.QtCore import QTimer, QObject, pyqtSignal
from PyQt5.QtWidgets import (
    QApplication,
    QSystemTrayIcon,
    QMenu,
    QAction,
    QMessageBox,
    QProgressDialog)
from polyhost._version import __version__
from polyhost.services import problem_report
from polyhost.gui.get_icon import get_icon
from polyhost.services import log_bundle
from polyhost.gui.theme import apply_dark_palette
from polyhost.gui.update_ui import UpdateProgressController
from polyhost.gui.icon_state_manager import IconStateManager
from polyhost.gui.qt_crash import install_qt_message_handler
from polyhost.gui.tray_wait import TrayVisibilityWaiter
from polyhost.gui.log_viewer import LogViewerDialog
from polyhost.handler.remote_window import TCP_PORT
from polyhost.handler.browser_url_source import BrowserUrlSource


IS_PLASMA = os.getenv("XDG_CURRENT_DESKTOP") == "KDE"
_IS_WAYLAND = os.getenv("XDG_SESSION_TYPE") == "wayland"

if IS_PLASMA:
    import polyhost.handler.kde_win_reporter as pwc
elif _IS_WAYLAND:
    # pywinctl can't see native Wayland windows; use the GNOME Shell extension
    # reporter (untested — needs the 'Window Calls' extension). X11 is unaffected.
    import polyhost.handler.gnome_wayland_reporter as pwc
else:
    import pywinctl as pwc

UPDATE_CYCLE_MSEC = 250
NEW_WINDOW_ACCEPT_TIME_MSEC = 1000
HEARTBEAT_MSEC = 15000  # resend current window state periodically so the host can catch up

from polyhost.util.log_util import DEBUG_DETAILED, make_stream_handler, make_collapse_handler  # noqa: F401  (registers debug_detailed on import)
from polyhost.handler.active_window import log_env_info

class _UpdateBridge(QObject):
    """Marshals the updater threads' callbacks (which fire off the Qt thread) back
    onto the Qt main thread via queued signals — the forwarder has no WorkerBridge."""
    available = pyqtSignal(object)
    no_update = pyqtSignal()
    error = pyqtSignal(str)
    progress = pyqtSignal(int, str)
    finished_ok = pyqtSignal()
    relay_needed = pyqtSignal(str)
    failed = pyqtSignal(str)


class PolyForwarder(QApplication):
    def __init__(self, log_level, host=None, host_file=None,
                 report_rpc=False, report_port=None, report_authkey_file=None):
        super().__init__(sys.argv)
        # Tray-only app: keep it out of the macOS Dock (no-op elsewhere).
        from polyhost.util.macos_ui import hide_dock_icon
        hide_dock_icon()
        self.host = host
        self.host_file = os.path.expanduser(host_file) if host_file else None

        # H4d: optionally push the active window over the authenticated network
        # window-report endpoint instead of the plaintext TCP relay. ⚠️ The RPC
        # transport is unit-tested but UNTESTED on hardware / cross-machine; the
        # default (report_rpc False) keeps the proven TCP path untouched.
        self._report_rpc = report_rpc
        self._report_port = report_port
        self._report_session = None
        self._report_authkey = None

        # This machine's OS, forwarded alongside the active window so the keyboard
        # reflects the OS of the computer you are working on (not just the one it is
        # plugged into). Constant for the process lifetime. An OsType value int.
        from polyhost.input.unicode_input import get_host_os
        self._os_value = get_host_os().value

        fmt = "[%(asctime)s] %(levelname)-7s {%(filename)s:%(lineno)d} - %(message)s"
        file_handler = RotatingFileHandler(
            filename="forwarder_log.txt",
            maxBytes=10 * 1024 * 1024,
            backupCount=3,
            encoding="utf-8"
        )
        file_handler.setFormatter(logging.Formatter(fmt))
        logging.basicConfig(level=log_level, handlers=[
            make_collapse_handler(file_handler),
            make_collapse_handler(make_stream_handler(fmt)),
        ])
        self.log = logging.getLogger("PolyForwarder")
        install_qt_message_handler(self.log)
        log_env_info(self.log)

        if self._report_rpc:
            from polyhost.server import protocol as _proto
            if report_authkey_file:
                try:
                    with open(report_authkey_file, "rb") as f:
                        self._report_authkey = f.read().strip()
                except OSError as e:
                    self.log.error("Could not read report authkey file %s: %s",
                                   report_authkey_file, e)
            if not self._report_authkey:
                self._report_authkey = _proto.load_or_create_authkey(
                    _proto.window_report_authkey_path())
                self.log.warning(
                    "Using this machine's local window-report authkey; for a "
                    "different keyboard machine pass --report-authkey-file with "
                    "its polykybd-winreport.authkey.")
            from polyhost.server.window_report_client import WindowReportSession
            self._report_session = WindowReportSession(
                self._report_port, self._report_authkey)
            self.log.info("Forwarder using the authenticated window-report endpoint (H4d).")

        # Browser-URL feed for THIS machine. The extension is already willing to
        # report here — it POSTs to 127.0.0.1 on whatever machine it runs on —
        # but in forwarder mode nothing was listening, so its /ping failed and a
        # forwarded browser could only ever match on its window title.
        self._url_dirty = False
        self._url_source = BrowserUrlSource(
            self.log, on_change=self._on_browser_url_changed)
        if self._url_source.start():
            self.log.info("Browser-URL reporting enabled for forwarded windows.")

        # Create the tray
        self.tray = QSystemTrayIcon(parent=self)
        self.icon_manager = IconStateManager(self, False, f"({__version__}) Forwarding to {host}")
        self._tray_waiter = TrayVisibilityWaiter(
            show=lambda: self.tray.setVisible(True),
            is_available=QSystemTrayIcon.isSystemTrayAvailable,
            log=self.log)
        self._tray_waiter.start()
        
        self.setQuitOnLastWindowClosed(False)
        self.win = None
        self.prev_win = None
        self.is_closing = False
        # Set when a self-update lands: main_app re-execs into the new version
        # *after* exec_() returns (from a clean, fully-unwound event loop),
        # mirroring how the headless daemon restarts — rather than calling
        # os.execv from inside a live Qt slot with the tray/timer still up.
        self.wants_restart = False
        self.title = None
        self.last_update_msec = 0

        self.heartbeat_msec = 0

        self._tray_waiter.start()
        self.set_style()

        self.menu = QMenu()

        self.exit = QAction(get_icon("power.svg"), "Quit", parent=self)
        # noinspection PyUnresolvedReferences
        self.exit.triggered.connect(self.quit_app)
        self.support = QAction(get_icon("support.svg"), "Get Support", parent=self)
        # noinspection PyUnresolvedReferences
        self.support.triggered.connect(self.open_support)
        self.about = QAction(get_icon("info.svg"), "About", parent=self)
        # noinspection PyUnresolvedReferences
        self.about.triggered.connect(self.open_about)

        self.log_dialog = QAction(get_icon("log.svg"), "Log file...", parent=self)
        # noinspection PyUnresolvedReferences
        self.log_dialog.triggered.connect(self.open_log)
        self.log_viewer = None

        # The forwarder runs on a DIFFERENT machine from the keyboard, so its
        # logs can never appear in a bundle collected on the host side — and its
        # failure modes (which window backend this desktop selects, the report
        # transport, the authkey) are exactly the log-diagnosable kind. Both
        # entries are therefore worth as much here as in the tray app.
        self.report_problem_action = QAction(get_icon("feedback.svg"),
                                             "Report a Problem...", parent=self)
        # noinspection PyUnresolvedReferences
        self.report_problem_action.triggered.connect(self.open_report_problem)
        self.report_problem_dialog = None

        self.collect_logs_action = QAction(get_icon("archive.svg"),
                                           "Collect logs...", parent=self)
        # noinspection PyUnresolvedReferences
        self.collect_logs_action.triggered.connect(self.open_log_bundle)
        self.log_bundle_dialog = None

        self.update_action = QAction(get_icon("browser_updated.svg"), "Check for updates...", parent=self)
        # noinspection PyUnresolvedReferences
        self.update_action.triggered.connect(self._on_update_clicked)

        # Update plumbing (host-app update only — the forwarder has no device).
        self._update_bridge = _UpdateBridge()
        self._update_bridge.available.connect(self._on_update_available)
        self._update_bridge.no_update.connect(self._on_no_update)
        self._update_bridge.error.connect(self._on_check_error)
        self._update_bridge.progress.connect(self._on_update_progress)
        self._update_bridge.finished_ok.connect(self._on_update_done)
        self._update_bridge.relay_needed.connect(self._on_relay_needed)
        self._update_bridge.failed.connect(self._on_update_failed)
        self._update_checker = None
        self._update_installer = None
        self._update_ui = UpdateProgressController(self.log)

        self.menu.addAction(self.log_dialog)
        self.menu.addAction(self.report_problem_action)
        self.menu.addAction(self.collect_logs_action)
        self.menu.addAction(self.update_action)
        self.menu.addAction(self.support)
        self.menu.addAction(self.about)
        self.menu.addAction(self.exit)
        
        self.tray.setContextMenu(self.menu)

        self.icon_manager.set_connected()
        
        QTimer.singleShot(1000, self.active_window_reporter)

    def set_style(self):
        """Dark Fusion theme — shared with PolyHost (gui/theme.py)."""
        apply_dark_palette(self)

    def _on_browser_url_changed(self):
        """A tab switch / SPA navigation changed the URL. The window-change test
        below is handle+title, which such a change does not move, so flag it and
        let the next 250 ms tick re-send instead of waiting for the heartbeat.

        Called from the receiver's HTTP thread, so it deliberately does nothing
        but set a bool — no Qt objects, no send from off the main thread."""
        self._url_dirty = True

    def _resolve_host(self):
        """Return the target host string (from --host or the host-file), or None."""
        host = self.host
        if self.host_file:
            try:
                with open(self.host_file) as f:
                    host = f.read().strip()
            except OSError:
                return None  # file absent means no active session
        return host or None

    def _send_via_rpc(self, handle, title, name, url=None):
        """Push the active window over the authenticated window-report endpoint
        (H4d). `WindowReportSession` keeps the connection, reconnecting when it
        fails *or when the resolved host changes* — with `--host-file` the
        target can be rewritten between two reports."""
        host = self._resolve_host()
        if not host:
            # The host-file is gone, i.e. no active session. Drop the
            # connection: when the file returns it may name a different
            # machine, and a stale connection would keep serving the old one.
            if self._report_session is not None:
                self._report_session.close()
            return False
        try:
            self._report_session.report(
                host, handle, name, title, os=self._os_value, url=url)
            return True
        except Exception as e:
            self.log.error("Window-report RPC to %s failed: %s", host, e)
            return False

    def send_to_host(self, handle, title, name, url=None):
        # ⚠️ `url` rides the authenticated RPC path ONLY. The legacy relay's
        # framing is positional `handle;name;title;os` with the free-text field
        # in the middle, so a title containing ';' already truncates the title
        # and kills the os field — a fifth field would deepen a live bug on a
        # transport that is off by default.
        if self._report_rpc:
            return self._send_via_rpc(handle, title, name, url=url)
        host = self._resolve_host()
        if not host:
            return False
        try:
            ip = ipaddress.ip_address(host)
        except ValueError:
            ip = socket.gethostbyname(host)
        except OSError as err:
            self.log.error("Could not resolve %s: %s", host, err)
            return False
        try:
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            s.settimeout(3.0)
            s.connect((str(ip), TCP_PORT))
            # 4th field = this machine's OS (an OsType value int); older daemons
            # split on ';' and read the first three fields, ignoring the extra one.
            s.send(f"{handle};{name};{title};{self._os_value}".encode("utf-8"))
            s.close()
            return True
        except socket.timeout as err:
            self.log.error("Connection timed out: %s",err)
        except ConnectionRefusedError as err:
            self.log.error("Connection refused: %s", err)
        except ConnectionAbortedError as err:
            self.log.error("Connection aborted: %s", err)
        except ConnectionResetError as err:
            self.log.error("Connection reset: %s", err)
        except ConnectionError as err:
            self.log.error("Connection error: %s", err)
        return False

    def _diagnostics_text(self) -> str:
        """Diagnostics for a forwarder report.

        Composition lives in the Qt-free `problem_report` module: this file
        imports pywinctl at module load, so anything left here cannot be tested
        in the documented environment (the forwarder smoke mode skips without
        it). This method only supplies the state.
        """
        return problem_report.forwarder_diagnostics(
            __version__, host=self.host, host_file=self.host_file,
            report_rpc=self._report_rpc, report_port=self._report_port)

    def open_report_problem(self):
        """Guided problem report (retained instance — see PolyHost.open_report_problem)."""
        from polyhost.gui.report_problem_dialog import ReportProblemDialog
        if self.report_problem_dialog is None:
            self.report_problem_dialog = ReportProblemDialog(
                parent=None, diagnostics_cb=self._diagnostics_text)
        self.report_problem_dialog.show()
        self.report_problem_dialog.raise_()
        self.report_problem_dialog.activateWindow()

    def open_log_bundle(self):
        """Log-collection dialog (retained instance — see PolyHost.open_log_bundle)."""
        from polyhost.gui.log_bundle_dialog import LogBundleDialog
        if self.log_bundle_dialog is None:
            self.log_bundle_dialog = LogBundleDialog(
                parent=None, diagnostics_cb=self._diagnostics_text)
        self.log_bundle_dialog.show()
        self.log_bundle_dialog.raise_()
        self.log_bundle_dialog.activateWindow()

    def open_log(self):
        # assignment is needed otherwise the dialog would go away immediately
        delta = time.perf_counter()
        # See host.py: one declaration in log_bundle.LOG_SOURCES feeds the
        # bundle, the clipboard text and both viewers' tabs.
        log_files = log_bundle.viewer_files(always=("forwarder",))
        self.log_viewer = LogViewerDialog(log_files, collect_cb=self.open_log_bundle)
        self.log_viewer.show()
        delta = time.perf_counter() - delta
        self.log.info("Opened log dialog in '%f' sec", delta)
        
    @staticmethod
    def open_support():
        webbrowser.open("https://discord.gg/gW8JescH7M", new=0, autoraise=True)

    @staticmethod
    def open_about():
        webbrowser.open("https://ko-fi.com/polykb", new=0, autoraise=True)
        
    # ------------------------------------------------------------------
    # Host-app update (mirrors the tray app's flow, host-only — no firmware,
    # since the forwarder has no device). Threads marshal back via _update_bridge.
    # ------------------------------------------------------------------
    def _on_update_clicked(self):
        from polyhost.services.updater import UpdateChecker
        if self._update_installer is not None and self._update_installer.is_alive():
            return
        if self._update_checker is not None and self._update_checker.is_alive():
            return
        self.update_action.setText("Checking for updates...")
        ub = self._update_bridge
        self._update_checker = UpdateChecker(
            current_fw_version=None,   # host-only: the forwarder owns no keyboard
            on_update_available=lambda rel: ub.available.emit(rel),
            on_host_no_update=lambda: ub.no_update.emit(),
            on_error=lambda msg: ub.error.emit(msg),
        )
        self._update_checker.start()

    def _on_update_available(self, release):
        from polyhost.gui.update_dialog import confirm_update
        self.update_action.setText("Check for updates...")
        message = f"Version {release.version} is available."
        if not confirm_update("Update PolyKybdHost", message,
                              notes=getattr(release, "notes", ""),
                              html_url=getattr(release, "html_url", ""),
                              release_name=getattr(release, "name", ""),
                              question="Download, install, and restart the forwarder now?"):
            return
        self._run_update_installer(release)

    def _on_no_update(self):
        self.update_action.setText("Check for updates...")
        QMessageBox.information(
            None, "PolyKybdHost Update",
            f"You are running the latest version (v{__version__}).")

    def _on_check_error(self, msg):
        self.update_action.setText("Check for updates...")
        QMessageBox.warning(
            None, "PolyKybdHost Update",
            f"Could not check for updates:\n\n{msg}\n\nRun with --dev 1 for details.")

    def _run_update_installer(self, release):
        from polyhost.services.updater import UpdateInstaller
        if self._update_installer is not None and self._update_installer.is_alive():
            return
        self.update_action.setEnabled(False)
        dlg = QProgressDialog(f"Downloading v{release.version}…", "", 0, 100)
        dlg.setWindowTitle("PolyKybdHost Update")
        dlg.setCancelButton(None)
        dlg.setMinimumDuration(0)
        dlg.show()
        self._update_ui.attach(dlg)
        ub = self._update_bridge
        self._update_installer = UpdateInstaller(
            release,
            on_progress=lambda pct, m: ub.progress.emit(pct, m),
            on_finished_ok=lambda: ub.finished_ok.emit(),
            on_relay_needed=lambda path: ub.relay_needed.emit(path),
            on_failed=lambda m: ub.failed.emit(m),
        )
        self._update_installer.start()

    def _on_update_progress(self, percent, message):
        self._update_ui.on_progress(percent, message)

    def _on_update_done(self):
        self._update_ui.close()
        self.log.info("Update applied, restarting forwarder...")
        # Re-exec from a clean state after the event loop exits (see
        # `wants_restart` and main_app), not with os.execv from inside this
        # slot while the tray/timer are still live — exactly like the daemon.
        self.wants_restart = True
        self.quit_app()

    def _on_relay_needed(self, relay_path):
        """Windows: some files were locked; a relay script finishes the copy
        once we exit. Spawned through the shared controller so it gets the
        normalised (windowless) interpreter and updater.spawn_detached — this
        used to be a bare Popen on sys.executable, which left the restarted
        forwarder owning a console window and dying with any job object."""
        if not self._update_ui.stage_relay(relay_path):
            # Nothing will finish the locked-file copy if we exit now, and the
            # tree is already partially rewritten — surface it and stay up.
            self._on_update_failed(
                "Could not start the update relay; the update is incomplete. "
                "See the log for details.")
            return
        # Brief pause so the user sees the "Restarting" label before we vanish.
        QTimer.singleShot(1200, self.quit)

    def _on_update_failed(self, message):
        self._update_ui.close()
        self.update_action.setEnabled(True)
        self.log.error("Update failed: %s", message)
        QMessageBox.warning(
            None, "Update failed", f"Could not apply the update:\n\n{message}")

    def quit_app(self):
        self.icon_manager.set_disconnected()
        self.is_closing = True
        if self._report_session is not None:
            self._report_session.close()
        self._url_source.close()
        # Tear the tray icon down before the loop exits so a re-exec (update
        # restart) doesn't leave a stale icon behind / a doubled tray.
        try:
            self.tray.hide()
        except Exception:  # noqa: BLE001 — tray teardown must never block quit
            pass
        self.quit()

    def active_window_reporter(self):
        self.last_update_msec += UPDATE_CYCLE_MSEC
        self.heartbeat_msec += UPDATE_CYCLE_MSEC
        win = pwc.getActiveWindow()
        if win:
            try:
                if self.prev_win != win:
                    self.prev_win = win
                    self.last_update_msec = 0
                if self.last_update_msec > NEW_WINDOW_ACCEPT_TIME_MSEC:
                    #just to limit the time value:
                    self.last_update_msec = NEW_WINDOW_ACCEPT_TIME_MSEC * 2
                    app_name = win.getAppName()
                    # None for every non-browser app, and for a browser whose
                    # extension report is stale/unfocused — so a URL can never
                    # linger onto the wrong window.
                    url = self._url_source.current_url(app_name)
                    changed = (
                        self.win is None
                        or win.getHandle() != self.win.getHandle()
                        or win.title != self.title
                        or self._url_dirty
                    )
                    if changed or self.heartbeat_msec >= HEARTBEAT_MSEC:
                        self.win = win
                        self.title = win.title
                        self._url_dirty = False
                        handle = win.getHandle()
                        self.send_to_host(handle, self.title, app_name, url=url)
                        if changed:
                            self.log.info("Active App: '%s' %s %d", self.title, app_name, handle)
                        else:
                            self.log.debug("Heartbeat: '%s' %s %d", self.title, app_name, handle)
                        self.heartbeat_msec = 0
            except Exception as e:
                self.log.warning("Exception in window reporter: %s", e)
        elif self.win:
            self.log.info("No active window")
            self.win = None
            self.title = None
            self.heartbeat_msec = 0
            self.send_to_host(0, "", "")

        if not self.is_closing:
            QTimer.singleShot(UPDATE_CYCLE_MSEC, self.active_window_reporter)
        else:
            self.log.info("No more active window reporting.")