path: root/ui/main_window.py

"""
ui/main_window.py

Toolbar (left→right):
  [▶ RUN] [⬤ LOG] | [⊞ Devices] | [⚗ Signals] | [📐 Plot] | [⚙ Settings]
                                                               spacer | timer | [📁 File]
"""

from PyQt6.QtWidgets import (
    QMainWindow, QWidget, QHBoxLayout, QVBoxLayout,
    QSplitter, QStatusBar, QLabel, QPushButton,
    QToolBar, QSizePolicy, QApplication, QFrame,
    QAbstractSpinBox, QComboBox,
)
from PyQt6.QtCore import Qt, QTimer, QObject, QEvent, pyqtSlot
import os


class _WheelBlocker(QObject):
    """App-level event filter: blocks accidental scroll-wheel changes on
    spin boxes, combo boxes and sliders that don't have keyboard focus."""
    def eventFilter(self, obj, event):
        if event.type() == QEvent.Type.Wheel:
            if isinstance(obj, (QAbstractSpinBox, QComboBox)):
                if not obj.hasFocus():
                    event.ignore()
                    return True
        return super().eventFilter(obj, event)

from devices.arduino_device  import ArduinoDevice
from devices.nidaqmx_device  import NidaqmxDevice
from devices.serial_device   import SerialDevice
from devices.device_registry import DeviceRegistry
from core.acquisition        import AcquisitionEngine
from core.signal_processor   import SignalProcessor
from core.profile            import Profile, ProfileManager

from ui.control_panel              import ControlPanel
from ui.strip_chart                import StripChartWidget
from ui.profile_manager_ui         import ProfileButton
from ui.windows.devices_window     import DevicesWindow
from ui.windows.signals_window     import ChannelsWindow
from ui.windows.plot_window        import PlotWindow, build_default_layout
from ui.windows.settings_window    import SettingsWindow

from plugins.plugin_manager import PluginManager
from plugins.base_plugin    import PluginContext


_DARK_QSS  = os.path.join(os.path.dirname(os.path.abspath(__file__)), "style_dark.qss")
_LIGHT_QSS = os.path.join(os.path.dirname(os.path.abspath(__file__)), "style_light.qss")


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("LabDAQ")
        self.setMinimumSize(1000, 640)

        self.registry  = DeviceRegistry()
        self.engine    = AcquisitionEngine(poll_interval_ms=100)
        self.processor = SignalProcessor()
        self._settings = dict(SettingsWindow._defaults)

        self._elapsed          = 0
        self._win_devices      = None
        self._win_channels     = None
        self._win_plot         = None
        self._win_settings     = None

        _plugins_dir = os.path.join(os.path.dirname(os.path.dirname(
            os.path.abspath(__file__))), "plugins")
        self._plugin_mgr = PluginManager(_plugins_dir)
        self._plugin_mgr.discover()
        # {plugin_id: [QAction, ...]}  toolbar actions to remove on unload
        self._plugin_toolbar_actions: dict = {}

        self._wheel_blocker = _WheelBlocker(self)
        QApplication.instance().installEventFilter(self._wheel_blocker)

        self._init_demo_devices()
        self._build_ui()
        self._connect_signals()
        self._init_plugins()

    # ── Demo ──────────────────────────────────────────────────────────────

    def _init_demo_devices(self):
        for dev in [
            NidaqmxDevice(device_id="ni_0",  num_analog=4, num_di=2, num_do=4, simulate=True),
            ArduinoDevice(device_id="ard_0", num_analog=4, num_di=2, num_do=4, simulate=True),
            SerialDevice (device_id="ser_0", num_channels=3, simulate=True),
        ]:
            dev.connect()
            self.registry.add_instance(dev)
            self.engine.add_device(dev)

    # ── UI ────────────────────────────────────────────────────────────────

    def _build_ui(self):
        tb = QToolBar(); tb.setObjectName("mainToolbar"); tb.setMovable(False)
        self.addToolBar(tb)
        self._toolbar = tb

        def _sep():
            s = QFrame(); s.setFrameShape(QFrame.Shape.VLine)
            s.setObjectName("toolbarSep")
            s.setFixedWidth(6); return s

        # ── 📁 File button (profiles) ─────────────────────────────────────
        self._file_btn = ProfileButton(
            on_new=self._profile_new,
            get_profile=self._profile_capture,
            apply_profile=self._profile_apply,
        )
        tb.addWidget(self._file_btn)
        tb.addWidget(_sep())

        self._run_btn = QPushButton("▶  RUN")
        self._run_btn.setObjectName("runButton"); self._run_btn.setCheckable(True)
        self._run_btn.clicked.connect(self._toggle_run); tb.addWidget(self._run_btn)

        self._log_btn = QPushButton("⬤  LOG")
        self._log_btn.setObjectName("logButton"); self._log_btn.setCheckable(True)
        self._log_btn.setEnabled(False); self._log_btn.clicked.connect(self._toggle_log)
        tb.addWidget(self._log_btn)

        self._clear_btn = QPushButton("⌫  Clear")
        self._clear_btn.setObjectName("toolbarSectionBtn")
        self._clear_btn.setToolTip("Clear plot history")
        self._clear_btn.clicked.connect(self._clear_history)
        tb.addWidget(self._clear_btn)

        tb.addWidget(_sep())

        dev_btn = QPushButton("⊞  Devices")
        dev_btn.setObjectName("toolbarSectionBtn"); dev_btn.setCheckable(True)
        dev_btn.clicked.connect(lambda c: self._toggle_win("devices", c, dev_btn))
        tb.addWidget(dev_btn); self._btn_devices = dev_btn

        sig_btn = QPushButton("⚗  Channels")
        sig_btn.setObjectName("toolbarSectionBtn"); sig_btn.setCheckable(True)
        sig_btn.clicked.connect(lambda c: self._toggle_win("channels", c, sig_btn))
        tb.addWidget(sig_btn); self._btn_channels = sig_btn

        plot_btn = QPushButton("📐  Plot")
        plot_btn.setObjectName("toolbarSectionBtn"); plot_btn.setCheckable(True)
        plot_btn.clicked.connect(lambda c: self._toggle_win("plot", c, plot_btn))
        tb.addWidget(plot_btn); self._btn_plot = plot_btn

        tb.addWidget(_sep())

        # Plugin buttons are inserted here at runtime (between this sep and spacer)
        self._plugin_sep_action = tb.addWidget(_sep())
        self._plugin_sep_action.setVisible(False)

        # Spacer + clock — plugin buttons insert before this action
        spacer = QWidget()
        spacer.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Preferred)
        self._spacer_action = tb.addWidget(spacer)

        self._time_lbl = QLabel("00:00:00"); self._time_lbl.setObjectName("timeLabel")
        tb.addWidget(self._time_lbl)

        tb.addWidget(_sep())

        set_btn = QPushButton("⚙  Settings")
        set_btn.setObjectName("toolbarSectionBtn"); set_btn.setCheckable(True)
        set_btn.clicked.connect(lambda c: self._toggle_win("settings", c, set_btn))
        tb.addWidget(set_btn); self._btn_settings = set_btn


        # ── Central ───────────────────────────────────────────────────────
        central = QWidget(); self.setCentralWidget(central)
        root = QHBoxLayout(central); root.setContentsMargins(0,0,0,0); root.setSpacing(0)

        hsplit = QSplitter(Qt.Orientation.Horizontal); hsplit.setHandleWidth(3)

        self._ctrl = ControlPanel(self.registry)
        self._ctrl.setMinimumWidth(200); self._ctrl.setMaximumWidth(340)
        self._ctrl._add_demo_widgets()
        hsplit.addWidget(self._ctrl)

        self._chart = StripChartWidget(self.engine, self.registry, self.processor)
        hsplit.addWidget(self._chart)
        hsplit.setSizes([260, 1000])

        root.addWidget(hsplit)

        sb = QStatusBar(); self.setStatusBar(sb)
        self._status = QLabel("Ready"); sb.addWidget(self._status)
        self._log_lbl = QLabel(""); sb.addPermanentWidget(self._log_lbl)

        self._clock = QTimer(self); self._clock.setInterval(1000)
        self._clock.timeout.connect(self._tick)

    def _connect_signals(self):
        self.engine.new_data.connect(self.processor.on_raw_data)
        self.processor.processed_data.connect(self._chart.on_new_data)
        self.engine.log_started.connect(lambda p: self._log_lbl.setText(f"● {p}"))
        self.engine.log_stopped.connect(lambda p: self._log_lbl.setText(f"✓ {p}"))

    # ── Plugin lifecycle ──────────────────────────────────────────────────

    def _make_plugin_context(self) -> PluginContext:
        return PluginContext(
            registry=self.registry,
            engine=self.engine,
            processor=self.processor,
            main_window=self,
        )

    def _init_plugins(self):
        ctx = self._make_plugin_context()
        for plugin in self._plugin_mgr.load_enabled(ctx):
            self._install_plugin(plugin)

    def _install_plugin(self, plugin):
        """Register a loaded plugin's devices, filters, and toolbar buttons."""
        from core.signal_processor import FILTER_CLASSES

        # Devices
        for dev in plugin.get_devices():
            try:
                dev.connect()
                self.registry.add_instance(dev)
                self.engine.add_device(dev)
            except Exception as exc:
                print(f"[Plugin:{plugin.plugin_id}] Device error: {exc}")

        # Custom filter classes
        for name, cls in plugin.get_filter_classes().items():
            FILTER_CLASSES[name] = cls

        # Toolbar buttons
        actions = plugin.get_toolbar_actions()
        tb_actions = []
        if actions:
            self._plugin_sep_action.setVisible(True)
            for pa in actions:
                label = f"{pa.icon}  {pa.label}" if pa.icon else pa.label
                btn = QPushButton(label)
                btn.setObjectName("toolbarSectionBtn")
                btn.setCheckable(pa.checkable)
                if pa.tooltip:
                    btn.setToolTip(pa.tooltip)
                btn.clicked.connect(pa.callback)
                if pa.button_ref_callback:
                    pa.button_ref_callback(btn)
                action = self._toolbar.insertWidget(self._spacer_action, btn)
                tb_actions.append(action)
        self._plugin_toolbar_actions[plugin.plugin_id] = tb_actions

        self._chart.refresh()
        if self._win_plot:
            self._win_plot.refresh_channels()
        self._status.setText(f"Plugin enabled: {plugin.name}")

    def _uninstall_plugin(self, plugin_id: str):
        """Remove a plugin's toolbar buttons, devices, and filter classes."""
        from core.signal_processor import FILTER_CLASSES

        plugin = self._plugin_mgr.get_plugin(plugin_id)

        # Remove toolbar buttons first (before unload)
        for action in self._plugin_toolbar_actions.pop(plugin_id, []):
            self._toolbar.removeAction(action)

        # Hide plugin separator if no plugins remain
        if not any(acts for acts in self._plugin_toolbar_actions.values()):
            self._plugin_sep_action.setVisible(False)

        if plugin is None:
            return

        # Remove custom filter classes
        for name in plugin.get_filter_classes():
            FILTER_CLASSES.pop(name, None)

        # Remove devices contributed by this plugin
        for dev in plugin.get_devices():
            dev_id = dev.info.device_id
            try:
                dev.disconnect()
            except Exception:
                pass
            self.registry.remove_instance(dev_id)
            self.engine.remove_device(dev_id)

        self._chart.refresh()
        if self._win_plot:
            self._win_plot.refresh_channels()
        self._status.setText(f"Plugin disabled: {plugin.name}")

    def plugin_enable(self, plugin_id: str):
        """Called by SettingsWindow when user enables a plugin."""
        ctx = self._make_plugin_context()
        plugin = self._plugin_mgr.enable(plugin_id, ctx)
        if plugin:
            self._install_plugin(plugin)

    def plugin_disable(self, plugin_id: str):
        """Called by SettingsWindow when user disables a plugin."""
        self._uninstall_plugin(plugin_id)
        self._plugin_mgr.disable(plugin_id)

    # ── Window management ─────────────────────────────────────────────────

    def _toggle_win(self, name: str, checked: bool, btn: QPushButton):
        creators = {
            "devices":  self._open_devices,
            "channels": self._open_channels,
            "plot":     self._open_plot,
            "settings": self._open_settings,
        }
        wins = {
            "devices":  "_win_devices",
            "channels": "_win_channels",
            "plot":     "_win_plot",
            "settings": "_win_settings",
        }
        if checked:
            creators[name]()
        else:
            win = getattr(self, wins[name], None)
            if win: win.hide()

    def _open_devices(self):
        if self._win_devices is None:
            self._win_devices = DevicesWindow(self.registry, self.engine, self)
            self._win_devices.device_added.connect(self._on_device_added)
            self._win_devices.device_removed.connect(self._on_device_removed)
            self._win_devices.device_reconfigured.connect(self._on_device_reconfigured)
            self._win_devices.channel_visibility_changed.connect(self._on_channel_visibility_changed)
            self._win_devices.channel_name_changed.connect(self._on_channel_name_changed)
            self._win_devices.channel_unit_changed.connect(self._on_channel_unit_changed)
            self._win_devices.closed.connect(lambda: self._btn_devices.setChecked(False))
        self._show_win(self._win_devices, "right")

    def _open_channels(self):
        if self._win_channels is None:
            self._win_channels = ChannelsWindow(self.registry, self.processor, self)
            self._win_channels.derived_changed.connect(self._on_derived_changed)
            self._win_channels.closed.connect(lambda: self._btn_channels.setChecked(False))
        self._show_win(self._win_channels, "right")

    def _open_plot(self):
        if self._win_plot is None:
            self._win_plot = PlotWindow(self.registry, self.processor,
                                        self._chart._cfg, self)
            self._win_plot.layout_applied.connect(self._chart.apply_layout)
            self._win_plot.closed.connect(lambda: self._btn_plot.setChecked(False))
        else:
            self._win_plot.cfg = self._chart._cfg
            self._win_plot._populate()
        self._show_win(self._win_plot, "below")

    def _open_settings(self):
        if self._win_settings is None:
            self._win_settings = SettingsWindow(self.registry, self.engine,
                                                 self._settings, self,
                                                 plugin_manager=self._plugin_mgr)
            self._win_settings.theme_changed.connect(self._apply_theme)
            self._win_settings.settings_changed.connect(self._on_settings)
            self._win_settings.plugin_enable_requested.connect(self.plugin_enable)
            self._win_settings.plugin_disable_requested.connect(self.plugin_disable)
            self._win_settings.closed.connect(lambda: self._btn_settings.setChecked(False))
        self._show_win(self._win_settings, "right")

    def _show_win(self, win: QWidget, position: str = "right"):
        if not win.isVisible():
            screen = QApplication.screenAt(self.geometry().center())
            if screen is None:
                screen = QApplication.primaryScreen()
            avail = screen.availableGeometry()
            # Use normalGeometry so maximized windows don't push child off-screen
            geo = self.normalGeometry()
            win.adjustSize()
            w, h = win.width(), win.height()
            if position == "right":
                x = geo.right() + 8
                y = geo.top() + 40
            else:
                x = geo.left()
                y = geo.bottom() + 8
            # Clamp to available screen area
            x = max(avail.left(), min(x, avail.right()  - w))
            y = max(avail.top(),  min(y, avail.bottom() - h))
            win.move(x, y)
        win.show(); win.raise_(); win.activateWindow()

    # ── Profile callbacks ─────────────────────────────────────────────────

    def _profile_new(self):
        """Reset to a blank slate."""
        # Clear devices
        for dev in list(self.registry.all_instances()):
            try:
                dev.disconnect()
            except Exception:
                pass
            self.registry.remove_instance(dev.info.device_id)
            self.engine.remove_device(dev.info.device_id)

        # Clear signal pipelines
        self.processor._pipelines.clear()

        # Clear derived channels
        for dc in list(self.processor.get_derived()):
            self.processor.remove_derived(dc.channel_id)

        # Clear controls
        self._ctrl.clear_widgets()

        # Refresh open windows
        if self._win_devices:
            self._win_devices.refresh()
        if self._win_channels:
            self._win_channels.refresh_derived()
        if self._win_plot:
            self._win_plot.refresh_channels()

        self._chart.refresh()
        self._status.setText("New profile — blank slate.")

    def _profile_capture(self, name: str = "Profile") -> Profile:
        """Serialise current state into a Profile object."""
        return ProfileManager.capture(
            registry=self.registry,
            processor=self.processor,
            plot_cfg=self._chart._cfg,
            control_specs=self._ctrl.get_specs(),
            settings=self._settings,
            profile_name=name,
            plugin_manager=self._plugin_mgr,
        )

    def _profile_apply(self, profile: Profile):
        """Restore state from a Profile object."""
        # Reconcile plugin enabled state before the rest of apply runs,
        # so plugin devices are present when channels/pipelines are restored.
        if profile.plugins_enabled is not None:
            wanted  = set(profile.plugins_enabled)
            current = set(self._plugin_mgr.get_enabled_ids())
            for pid in current - wanted:
                self.plugin_disable(pid)
            for pid in wanted - current:
                self.plugin_enable(pid)

        plot_cfg = ProfileManager.apply(
            profile=profile,
            registry=self.registry,
            processor=self.processor,
            control_panel=self._ctrl,
            settings_ref=self._settings,
            engine=self.engine,
            plugin_manager=self._plugin_mgr,
        )
        if plot_cfg:
            self._chart.apply_layout(plot_cfg)
        else:
            self._chart.refresh()
        # Refresh open windows
        if self._win_devices:
            self._win_devices.refresh()
        if self._win_channels:
            self._win_channels.refresh_derived()
        if self._win_plot:
            self._win_plot.refresh_channels()
        self._status.setText(f"Profile loaded: {profile.name}")

    # ── Device / signal events ────────────────────────────────────────────

    def _on_device_added(self):
        self._chart.refresh()
        if self._win_plot: self._win_plot.refresh_channels()
        if self._win_channels: self._win_channels.on_device_added()
        self._status.setText("Device added.")

    def _on_device_removed(self, dev_id: str):
        self._chart.refresh()
        if self._win_channels: self._win_channels.on_device_removed(dev_id)
        self._status.setText(f"Device '{dev_id}' removed.")

    def _on_device_reconfigured(self, dev_id: str):
        self._chart.refresh()
        if self._win_plot: self._win_plot.refresh_channels()
        if self._win_channels: self._win_channels.on_device_reconfigured(dev_id)
        self._status.setText(f"Device '{dev_id}' reconfigured.")

    def _on_channel_visibility_changed(self, dev_id: str, ch_id: str, enabled: bool):
        self._chart.on_channel_enabled_changed(dev_id, ch_id, enabled)
        if self._win_channels:
            self._win_channels.on_channel_enabled_changed(dev_id, ch_id, enabled)
        if self._win_plot:
            self._win_plot.refresh_channels()

    def _on_channel_name_changed(self, dev_id: str, ch_id: str, name: str):
        if self._win_channels:
            self._win_channels.on_channel_name_changed(dev_id, ch_id, name)

    def _on_channel_unit_changed(self, dev_id: str, ch_id: str, unit: str):
        if self._win_channels:
            self._win_channels.on_channel_unit_changed(dev_id, ch_id, unit)

    def _on_derived_changed(self):
        self._chart.refresh()
        if self._win_plot: self._win_plot.refresh_channels()

    # ── Run / Log ─────────────────────────────────────────────────────────

    def _toggle_run(self, c: bool):
        if c:
            self.engine.start()
            self._run_btn.setText("⏹  STOP"); self._log_btn.setEnabled(True)
            self._clock.start(); self._status.setText("Acquiring…")
        else:
            self.engine.stop()
            self._run_btn.setText("▶  RUN")
            if self._log_btn.isChecked(): self._log_btn.setChecked(False)
            self._log_btn.setEnabled(False); self._clock.stop()
            self._status.setText("Stopped")

    def _clear_history(self):
        self.engine.clear_history()
        self.processor.clear_history()
        self._chart.refresh()
        self._status.setText("History cleared.")

    def _toggle_log(self, c: bool):
        if c:
            p = self.engine.start_logging(
                os.path.join(self._settings.get("log_dir", "logs"), ""))
            self._log_btn.setText("⏹  LOGGING")
            self._status.setText(f"Logging → {p}")
        else:
            self.engine.stop_logging(); self._log_btn.setText("⬤  LOG")

    # ── Theme / settings ──────────────────────────────────────────────────

    def _apply_theme(self, theme: str):
        qss_file = _DARK_QSS if theme == "dark" else _LIGHT_QSS
        if os.path.exists(qss_file):
            with open(qss_file) as f:
                QApplication.instance().setStyleSheet(f.read())
        self._chart.set_theme(theme)

    def _on_settings(self, cfg: dict):
        self._settings.update(cfg)
        self.engine._interval = cfg.get("poll_ms", 100) / 1000.0

    def _tick(self):
        self._elapsed += 1
        h=self._elapsed//3600; m=(self._elapsed%3600)//60; s=self._elapsed%60
        self._time_lbl.setText(f"{h:02d}:{m:02d}:{s:02d}")

    def closeEvent(self, event):
        for w in (self._win_devices, self._win_channels,
                  self._win_plot, self._win_settings):
            if w: w.close()
        for plugin in list(self._plugin_mgr.get_loaded()):
            try:
                plugin.on_unload()
            except Exception:
                pass
        self.engine.stop(); event.accept()