"""
ui/strip_chart.py

Live scrolling strip chart.  One pyqtgraph plot per device,
all time-axes linked.  Per-channel colored traces with legend.
"""

import numpy as np
from typing import Dict

from PyQt6.QtWidgets import (
    QWidget, QVBoxLayout, QHBoxLayout, QLabel,
    QDoubleSpinBox, QPushButton, QSizePolicy, QComboBox,
)
from PyQt6.QtCore import Qt, pyqtSlot

try:
    import pyqtgraph as pg
    pg.setConfigOptions(antialias=True, background="#0b0e13", foreground="#475569")
    _HAS_PG = True
except ImportError:
    _HAS_PG = False

from devices.device_registry import DeviceRegistry
from core.acquisition import AcquisitionEngine


class StripChartWidget(QWidget):
    def __init__(self, engine: AcquisitionEngine, registry: DeviceRegistry):
        super().__init__()
        self.engine   = engine
        self.registry = registry
        self._window  = 30.0
        self._paused  = False
        self._plots: Dict[str, Dict] = {}  # dev_id -> {ch_id -> {curve, plot}}
        self._build()
        self.refresh()

    # ── Build ────────────────────────────────────────────────────────────

    def _build(self):
        layout = QVBoxLayout(self)
        layout.setContentsMargins(6, 6, 6, 4)
        layout.setSpacing(4)

        # Control bar
        ctrl = QHBoxLayout()
        ctrl.addWidget(QLabel("Window:"))

        self._win_spin = QDoubleSpinBox()
        self._win_spin.setRange(1.0, 3600.0)
        self._win_spin.setValue(self._window)
        self._win_spin.setSuffix(" s")
        self._win_spin.valueChanged.connect(self._on_window_changed)
        ctrl.addWidget(self._win_spin)
        ctrl.addSpacing(16)

        ctrl.addWidget(QLabel("Y-Scale:"))
        self._scale_cb = QComboBox()
        self._scale_cb.addItems(["Auto", "Fixed"])
        self._scale_cb.currentTextChanged.connect(self._on_scale_changed)
        ctrl.addWidget(self._scale_cb)

        ctrl.addStretch()

        self._pause_btn = QPushButton("⏸  Pause")
        self._pause_btn.setCheckable(True)
        self._pause_btn.setObjectName("pauseButton")
        self._pause_btn.toggled.connect(self._on_pause)
        ctrl.addWidget(self._pause_btn)

        layout.addLayout(ctrl)

        if _HAS_PG:
            self._gw = pg.GraphicsLayoutWidget()
            self._gw.setSizePolicy(
                QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding
            )
            layout.addWidget(self._gw)
        else:
            layout.addWidget(QLabel(
                "⚠  pyqtgraph not installed.\n\npip install pyqtgraph\n\nData is still acquired & logged.",
                alignment=Qt.AlignmentFlag.AlignCenter,
            ))

    # ── Refresh (rebuild plots after device list changes) ─────────────────

    def refresh(self):
        if not _HAS_PG:
            return
        self._gw.clear()
        self._plots.clear()

        devs   = self.registry.all_instances()
        n_devs = len(devs)
        if n_devs == 0:
            return

        ref_plot = None  # for x-axis linking

        for row, dev in enumerate(devs):
            plot = self._gw.addPlot(row=row, col=0)
            plot.setLabel("left", f"{dev.info.icon} {dev.info.name}")
            plot.showGrid(x=True, y=True, alpha=0.12)
            plot.getAxis("left").setStyle(tickFont=self._mono_font())
            plot.getAxis("bottom").setStyle(tickFont=self._mono_font())

            if row < n_devs - 1:
                plot.getAxis("bottom").setStyle(showValues=False)
                plot.getAxis("bottom").setHeight(0)
            else:
                plot.setLabel("bottom", "Elapsed (s)")

            if ref_plot is not None:
                plot.setXLink(ref_plot)
            ref_plot = plot

            legend = plot.addLegend(
                offset=(5, 5),
                labelTextColor="#94a3b8",
                brush=pg.mkBrush("#161d2e"),
                pen=pg.mkPen("#2a3558"),
            )

            self._plots[dev.info.device_id] = {}
            for ch in dev.info.channels:
                if not ch.enabled:
                    continue
                pen   = pg.mkPen(color=ch.color, width=1.8)
                curve = plot.plot([], [], pen=pen, name=ch.name)
                self._plots[dev.info.device_id][ch.channel_id] = {
                    "curve": curve,
                    "plot":  plot,
                }

    @staticmethod
    def _mono_font():
        from PyQt6.QtGui import QFont
        f = QFont("IBM Plex Mono", 8)
        return f

    # ── Slots ────────────────────────────────────────────────────────────

    @pyqtSlot(str, str, float, float)
    def on_new_data(self, device_id: str, channel_id: str, timestamp: float, value: float):
        if self._paused or not _HAS_PG:
            return
        dev_plots = self._plots.get(device_id)
        if dev_plots is None:
            return
        entry = dev_plots.get(channel_id)
        if entry is None:
            return

        buf = self.engine.get_buffer(device_id, channel_id)
        if buf is None or len(buf) < 2:
            return

        ts, vs = buf.window(self._window)
        if not ts:
            return
        ts_arr = np.array(ts, dtype=np.float64)
        vs_arr = np.array(vs, dtype=np.float64)

        entry["curve"].setData(ts_arr, vs_arr)
        t_max = ts_arr[-1]
        t_min = t_max - self._window
        entry["plot"].setXRange(t_min, t_max, padding=0)

        if self._scale_cb.currentText() == "Auto":
            entry["plot"].enableAutoRange(axis="y")

    def _on_window_changed(self, v: float):
        self._window = v

    def _on_pause(self, checked: bool):
        self._paused = checked
        self._pause_btn.setText("▶  Resume" if checked else "⏸  Pause")

    def _on_scale_changed(self, text: str):
        if not _HAS_PG:
            return
        if text == "Auto":
            for dev_plots in self._plots.values():
                for entry in dev_plots.values():
                    entry["plot"].enableAutoRange(axis="y")