"""
ui/strip_chart.py

Live scrolling strip chart using pyqtgraph.
Each channel gets its own colored trace. Window duration is adjustable.
"""

from typing import Dict, List
import numpy as np

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

try:
    import pyqtgraph as pg
    pg.setConfigOptions(antialias=True, background="#0d1117", foreground="#c9d1d9")
    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._plots: Dict[str, Dict] = {}  # device_id -> {channel_id -> curve}
        self._window_s = 30.0
        self._paused = False
        self._build_ui()
        self.refresh()

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

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

        self.window_spin = QDoubleSpinBox()
        self.window_spin.setRange(1, 600)
        self.window_spin.setValue(self._window_s)
        self.window_spin.setSuffix(" s")
        self.window_spin.valueChanged.connect(lambda v: setattr(self, "_window_s", v))
        ctrl.addWidget(self.window_spin)

        ctrl.addStretch()

        self.pause_btn = QPushButton("⏸ Pause")
        self.pause_btn.setCheckable(True)
        self.pause_btn.setObjectName("pauseButton")
        self.pause_btn.toggled.connect(lambda c: setattr(self, "_paused", c))
        self.pause_btn.toggled.connect(lambda c: self.pause_btn.setText("▶ Resume" if c else "⏸ Pause"))
        ctrl.addWidget(self.pause_btn)

        layout.addLayout(ctrl)

        if HAS_PG:
            self._build_pg_chart(layout)
        else:
            layout.addWidget(QLabel(
                "pyqtgraph not installed.\nRun: pip install pyqtgraph\n\nData is still being acquired.",
                alignment=Qt.AlignmentFlag.AlignCenter
            ))

    def _build_pg_chart(self, parent_layout):
        self.plot_widget = pg.GraphicsLayoutWidget()
        self.plot_widget.setSizePolicy(
            QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding
        )
        parent_layout.addWidget(self.plot_widget)

    def refresh(self):
        if not HAS_PG:
            return
        self.plot_widget.clear()
        self._plots.clear()

        devices = self.registry.all_instances()
        num = len(devices)
        if num == 0:
            return

        for row_idx, dev in enumerate(devices):
            plot = self.plot_widget.addPlot(row=row_idx, col=0)
            plot.setLabel("left", f"{dev.info.icon} {dev.info.name}")
            plot.showGrid(x=True, y=True, alpha=0.15)
            plot.getAxis("bottom").setStyle(showValues=(row_idx == num - 1))

            if row_idx < num - 1:
                plot.getAxis("bottom").setHeight(0)
            else:
                plot.setLabel("bottom", "Time (s)")

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

    @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, {})
        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.all()
        ts = np.array(ts, dtype=np.float64)
        vs = np.array(vs, dtype=np.float64)

        t_max = ts[-1]
        t_min = t_max - self._window_s
        mask = ts >= t_min
        ts = ts[mask]
        vs = vs[mask]

        entry["curve"].setData(ts, vs)
        entry["plot"].setXRange(t_min, t_max, padding=0)