pyGUI/src/pygui/backend/controllers/device_controller.py

"""QML-exposed controller for wafer device communication."""

from __future__ import annotations

import logging
import threading
from datetime import datetime
from typing import Any, Optional

from PySide6.QtCore import Property, QObject, Qt, Signal, Slot

from pygui.backend.controllers.session_controller import SessionController
from pygui.backend.data.local_settings import LocalSettings
from pygui.backend.models.data_model import TemperatureTableModel
from pygui.backend.utils import slot_error_boundary
from pygui.backend.visualization.graph_view import GraphView
from pygui.serialcomm.data_parser import (
    convert_to_temperatures,
    parse_binary_data,
    remove_trailing_zeros,
    save_to_csv,
)
from pygui.serialcomm.device_service import DeviceService
from pygui.serialcomm.serial_port import WaferInfo

# import pygui.backend.wafer_map_item

stream_controller = SessionController()
# engine.rootContext().setContextProperty("streamController", stream_controller)

log = logging.getLogger(__name__)


class DeviceController(QObject):
    """Controls serial communication with the temperature-sensing wafer.

    Exposed to QML as ``deviceController`` context property.
    """

    # ---- public signals ----
    portsUpdated = Signal(list)
    detectResult = Signal(object)  # WaferInfo dict or None
    readResult = Signal(dict)  # {"success": bool, "bytes": int} or {"error": str}
    eraseResult = Signal(dict)  # {"success": bool}
    debugResult = Signal(dict)  # {"success": bool, "sensor_bytes": int, "debug_bytes": int}
    parsedDataReady = Signal(dict)  # {"success": bool, "rows": int, "cols": int, "data": ..., "csv_path": str}
    statusRestored = Signal()  # Emitted when status is restored from previous session
    logMessage = Signal(str)
    activityLogUpdated = Signal(str)

    # ---- private signals (marshal worker-thread results back to main thread) ----
    _detectFinished = Signal(object, object)     # (port_or_None, WaferInfo_or_None)
    _readFinished = Signal(str, str, object)     # (port, family_code, bytes_or_None)

    def __init__(self, settings: LocalSettings, data_dir: str, parent: QObject | None = None) -> None:
        super().__init__(parent)
        self._settings = settings
        self._data_dir = data_dir
        self._service = DeviceService(settings)

        # Always start fresh — never restore connection status or logs from a prior session
        self._connection_status = "Disconnected"
        self._operation_in_progress = False
        self._activity_log: list[str] = []
        self._raw_bytes: Optional[bytes] = None
        from pathlib import Path
        self._save_data_dir: str = getattr(settings, 'save_data_dir', "") or str(Path(self._data_dir) / "csv")
        self._last_wafer_info: dict[str, Any] = {}
        self._last_csv_path: str = ""
        self._selected_port: str = ""
        self._data_row_count: int = 0
        self._data_col_count: int = 0
        self._data_model = TemperatureTableModel(self)
        self._graph_view = GraphView(self)

        # Reset connection status on fresh start
        self._connection_status = "Disconnected"

        # Marshal worker-thread results onto the Qt main thread.
        self._detectFinished.connect(self._handle_detect_finished, Qt.ConnectionType.QueuedConnection)
        self._readFinished.connect(self._handle_read_finished, Qt.ConnectionType.QueuedConnection)

        root_logger = logging.getLogger()
        root_logger.addHandler(QmlActivityLogHandler(self))
        root_logger.setLevel(logging.INFO)

    # ---- Properties ----
    @Property(list, notify=portsUpdated)
    def availablePorts(self) -> list[str]:
        """Currently available serial port names."""
        return self._service.enumerate_ports()

    def _set_connection_status(self, status: str) -> None:
        """Update connection status and notify QML."""
        if self._connection_status != status:
            self._connection_status = status
            self.portsUpdated.emit(self.availablePorts)

    @Property(str, notify=portsUpdated)
    def connectionStatus(self) -> str:
        """Current connection status string for display."""
        return self._connection_status

    @Property(bool, notify=portsUpdated)
    def operationInProgress(self) -> bool:
        """Whether a hardware operation is currently running."""
        return self._operation_in_progress

    @Property(str, notify=activityLogUpdated)
    def saveDataDir(self) -> str:
        """Directory where parsed CSV data files are saved."""
        return self._save_data_dir

    @Slot(str)
    @slot_error_boundary
    def setSaveDataDir(self, path: str) -> None:
        """Set the directory for saving parsed CSV data files."""
        self._save_data_dir = path
        self._append_log(f"Save data dir set to: {path}")
        self._save_status()

    @Property(str, notify=portsUpdated)
    def selectedPort(self) -> str:
        """Currently selected serial port shared across the UI."""
        return self._selected_port

    @Slot(str)
    @slot_error_boundary
    def setSelectedPort(self, port: str) -> None:
        """Set the active serial port from the port selector."""
        self._selected_port = port
        self._save_status()

    @Property(int, notify=parsedDataReady)
    def dataRowCount(self) -> int:
        """Number of rows in the parsed temperature dataset."""
        return self._data_row_count

    @Property(int, notify=parsedDataReady)
    def dataColCount(self) -> int:
        """Number of sensor columns in the parsed temperature dataset."""
        return self._data_col_count

    @Property(list, notify=detectResult)
    def lastWaferInfo(self) -> list:
        """Last detected wafer info as a flat list for QML bindings.

        Returns [familyCode, serialNumber, sensorCount, mfgDateHex, runtime, cycleCount].
        """
        info = self._last_wafer_info
        family_code = info.get("familyCode", "")
        sensor_count = info.get("sensorCount", 0)
        from pygui.serialcomm.data_parser import csv_column_count
        mapped_count = csv_column_count(family_code)
        if mapped_count > 0:
            sensor_count = mapped_count

        return [
            family_code,
            info.get("serialNumber", ""),
            sensor_count,
            info.get("mfgDateHex", ""),
            info.get("runtime", 0),
            info.get("cycleCount", 0),
        ]

    @Property(object, notify=detectResult)
    def dataModel(self) -> TemperatureTableModel:
        """QAbstractTableModel for the parsed temperature data table."""
        return self._data_model

    @Property(object, notify=debugResult)
    def graphView(self) -> GraphView:
        """GraphView for rendering sensor temperature line charts."""
        return self._graph_view

    @Slot(result=object)
    @slot_error_boundary
    def getChartData(self) -> dict[str, Any]:
        """Extract chart-ready data from the current model.

        Returns a dict with sensor names and per-sensor value lists
        suitable for GraphView.updateChart().
        """
        if not self._data_model or self._data_model.rowCount() == 0:
            return {"success": False}

        num_sensors = self._data_model.columnCount() - 1  # Exclude row index
        sensor_names = [f"Sensor{i+1}" for i in range(num_sensors)]

        # Extract per-sensor values from the model
        series_data = []
        for col in range(1, self._data_model.columnCount()):
            sensor_values = []
            for row in range(self._data_model.rowCount()):
                idx = self._data_model.index(row, col)
                val = self._data_model.data(idx)
                sensor_values.append(val if val else "0")
            series_data.append(sensor_values)

        return {
            "success": True,
            "sensor_names": sensor_names,
            "series": series_data,
        }

    def _append_log(self, message: str) -> None:
        """Append a message to the activity log and emit updated log."""
        ts = datetime.now().strftime("%H:%M:%S")
        self._activity_log.append(f"[{ts}] {message}")
        # Keep only last 200 lines
        if len(self._activity_log) > 200:
            self._activity_log = self._activity_log[-200:]
        self.activityLogUpdated.emit("\n".join(self._activity_log))

    @Slot()
    @slot_error_boundary
    def clearActivityLog(self) -> None:
        """Clear the activity log."""
        self._activity_log.clear()
        self.activityLogUpdated.emit("")
        self._save_status()

    @Slot()
    @slot_error_boundary
    def clearSession(self) -> None:
        """Clear the current session state, allowing a clean detect."""
        self._set_connection_status("Disconnected")
        self._selected_port = ""
        self._last_wafer_info = {}
        self._data_row_count = 0
        self._data_col_count = 0
        self._raw_bytes = None
        self._operation_in_progress = False
        self._data_model.reset()
        self._graph_view.resetChart()
        self._activity_log.clear()
        self.detectResult.emit(None)
        self.parsedDataReady.emit({"success": False})
        self.portsUpdated.emit(self.availablePorts)
        self.activityLogUpdated.emit("")
        self._save_status()


    def _set_operation_progress(self, in_progress: bool) -> None:
        """Set operation progress state and notify QML.

        portsUpdated is the shared notify signal for availablePorts,
        connectionStatus, and operationInProgress — emitting it forces
        QML to re-evaluate all three bindings.
        """
        self._operation_in_progress = in_progress
        self.portsUpdated.emit(self._service.enumerate_ports())

    # ---- Public Slots ----

    @Slot()
    @slot_error_boundary
    def refreshPorts(self) -> None:
        """Scan for available serial ports and emit updated list."""
        ports = self._service.enumerate_ports()
        self._set_connection_status("Disconnected")
        self.portsUpdated.emit(ports)
        self._append_log(f"Scanned {len(ports)} serial port(s)")
        self._save_status()

    @Slot()
    @slot_error_boundary
    def detectWafer(self) -> None:
        """Scan all serial ports for a wafer (runs in background thread).

        Mirrors C# chkConnection(): tries every available port with s0.
        Stores the found port in _selected_port for subsequent read/erase.
        Emits detectResult with WaferInfo dict on success, or None.
        """
        if self._operation_in_progress:
            self._append_log("Already busy — ignoring detect request")
            return

        # Clear previous session data & logs for a fresh session
        self._selected_port = ""
        self._last_wafer_info = {}
        self._data_row_count = 0
        self._data_col_count = 0
        self._raw_bytes = None
        self._data_model.reset()
        self._graph_view.resetChart()
        self._activity_log.clear()
        self.detectResult.emit(None)
        self.parsedDataReady.emit({"success": False})

        self._set_operation_progress(True)
        self._set_connection_status("Detecting...")
        self.portsUpdated.emit(self._service.enumerate_ports())
        self._append_log("Scanning all ports for wafer ...")

        # Spawn a daemon worker — UI stays responsive during the scan.
        threading.Thread(target=self._detect_worker, daemon=True).start()

    def _detect_worker(self) -> None:
        """Background-thread body for detectWafer."""
        try:
            port, info = self._service.detect_all_ports()
        except Exception as exc:
            log.exception("Detect worker crashed: %s", exc)
            port, info = None, None
        self._detectFinished.emit(port, info)

    @Slot(object, object)
    @slot_error_boundary
    def _handle_detect_finished(
        self, port: Optional[str], info: Optional[WaferInfo]
    ) -> None:
        """Main-thread completion handler for detectWafer."""
        if info is not None:
            self._selected_port = port or ""
            self._set_connection_status("Connected")
            wafer_dict = self._wafer_info_to_dict(info)
            self._last_wafer_info = wafer_dict
            self.detectResult.emit(wafer_dict)
            self._save_status()
        else:
            self._selected_port = ""
            self._set_connection_status("Disconnected")
            self._append_log("No wafer detected on any port")
            self._last_wafer_info = {}
            self.detectResult.emit(None)
        self._set_operation_progress(False)

    @Slot()
    @slot_error_boundary
    def readMemoryAsync(self) -> None:
        """Read wafer memory in a background thread.

        Uses the family code and port from the last detect.
        Emits readResult on completion; QML prompts for the save directory
        before calling parseAndSaveData.
        """
        if self._operation_in_progress:
            self._append_log("Already busy — ignoring read request")
            return

        if not self._selected_port:
            self._append_log("No wafer detected — run Detect Wafer first")
            self.readResult.emit({"error": "No port — detect wafer first"})
            return

        port = self._selected_port
        family_code = self._last_wafer_info.get("familyCode", "")

        self._set_operation_progress(True)
        self._set_connection_status("Reading...")
        self.portsUpdated.emit(self._service.enumerate_ports())
        self._append_log(
            f"Reading memory on {port} (family={family_code or 'auto'}) ..."
        )

        threading.Thread(
            target=self._read_worker,
            args=(port, family_code),
            daemon=True,
        ).start()

    def _read_worker(self, port: str, family_code: str) -> None:
        """Background-thread body for readMemoryAsync."""
        try:
            data = self._service.read_wafer_data(port, family_code)
        except Exception as exc:
            log.exception("Read worker crashed: %s", exc)
            data = None
        self._readFinished.emit(port, family_code, data)

    @Slot(str, str, object)
    @slot_error_boundary
    def _handle_read_finished(
        self, port: str, family_code: str, data: Optional[bytes]
    ) -> None:
        """Main-thread completion handler for readMemoryAsync."""
        if data is not None:
            self._set_connection_status("Connected")
            self._append_log(f"Read {len(data)} bytes")
            self._raw_bytes = data
            self.readResult.emit({"success": True, "bytes": len(data)})
            self._set_operation_progress(False)
            self._append_log("Memory read complete - choose save directory to save CSV")
            self._save_status()
            return
        self._set_connection_status("Disconnected")
        self._append_log("Read returned no data")
        self._raw_bytes = None
        self.readResult.emit({"error": "Read returned no data"})
        self._set_operation_progress(False)
        self._save_status()

    @Slot(str)
    @slot_error_boundary
    def eraseMemory(self, port: str) -> None:
        """Send p1 erase command."""
        self._set_operation_progress(True)
        self._set_connection_status("Erasing...")
        self.portsUpdated.emit(self._service.enumerate_ports())
        self._append_log(f"Erase command sent to {port} (wafer takes ~15s) ...")

        ok = self._service.erase_wafer(port)
        if ok:
            self._set_connection_status("Connected")
            self._append_log("Erase command accepted — wait ~15 seconds")
        else:
            self._set_connection_status("Disconnected")
            self._append_log("Erase command failed")
        self.eraseResult.emit({"success": ok})
        self._set_operation_progress(False)

    @Slot(str)
    @slot_error_boundary
    def readDebug(self, port: str) -> None:
        """Read both D1 (sensor) and F1 (debug/cold-junction) data.

        Emits debugResult with counts on success.
        """
        self._set_operation_progress(True)
        self._set_connection_status("Reading debug...")
        self.portsUpdated.emit(self._service.enumerate_ports())
        self._append_log(f"Reading debug data on {port} ...")

        # Step 1: D1 sensor data
        sensor_data = self._service.read_wafer_data(port, family_code="")
        if sensor_data is None:
            self._set_connection_status("Disconnected")
            self._append_log("D1 read failed — aborting debug read")
            self.debugResult.emit({"error": "D1 read failed"})
            self._set_operation_progress(False)
            return

        self._append_log(f"D1 read: {len(sensor_data)} bytes")

        # Step 2: F1 debug data
        debug_data = self._service.read_wafer_data(port, family_code="", cmd="F1")
        if debug_data is None:
            self._set_connection_status("Disconnected")
            self._append_log("F1 read failed")
            self.debugResult.emit({"error": "F1 read failed"})
            self._set_operation_progress(False)
            return

        self._set_connection_status("Connected")
        self._append_log(
            f"Debug read complete: {len(sensor_data)} sensor bytes, "
            f"{len(debug_data)} debug bytes"
        )
        self.debugResult.emit({
            "success": True,
            "sensor_bytes": len(sensor_data),
            "debug_bytes": len(debug_data),
        })
        self._set_operation_progress(False)

    @Slot(str, str)
    @slot_error_boundary
    def parseAndSaveData(self, family_code: str = "", port: str = "") -> None:
        """Parse raw bytes into temperatures and save to CSV.

        Uses the bytes from the most recent readMemoryAsync call.
        Emits parsedDataReady with the parsed result.
        """
        if not self._raw_bytes:
            self._append_log("No raw bytes to parse (read first)")
            self.parsedDataReady.emit({
                "success": False,
                "error": "No raw bytes available",
            })
            return

        if not self._save_data_dir:
            from pathlib import Path
            self._save_data_dir = str(Path(self._data_dir) / "csv")
            self._append_log(f"Auto-set save directory to: {self._save_data_dir}")
            self._save_status()

        fc = family_code or (
            self._last_wafer_info.get("familyCode", "")
            if self._last_wafer_info
            else ""
        )
        serial = (
            self._last_wafer_info.get("serialNumber", "UNKNOWN")
            if self._last_wafer_info
            else "UNKNOWN"
        )

        self._append_log(f"Parsing {len(self._raw_bytes)} bytes (family={fc or 'auto'}) ...")

        # Step 1: Parse binary → hex strings
        hex_data = parse_binary_data(self._raw_bytes, fc)
        if hex_data is None:
            self._append_log("Binary parse failed")
            self.parsedDataReady.emit({
                "success": False,
                "error": "Binary parse failed",
            })
            return

        rows = len(hex_data)
        cols = len(hex_data[0]) if hex_data else 0
        self._append_log(f"Parsed: {rows} rows × {cols} sensors")

        # Step 2: Convert hex → temperatures
        temp_data = convert_to_temperatures(hex_data, fc)

        # Step 3: Remove trailing zero rows
        remove_trailing_zeros(temp_data)
        self._append_log(f"After trim: {len(temp_data)} rows")

        # Step 4: Save to CSV
        csv_path = save_to_csv(temp_data, fc, serial, self._save_data_dir)
        if csv_path is None:
            self._append_log("CSV save failed")
            self.parsedDataReady.emit({
                "success": False,
                "error": "CSV save failed",
            })
            return

        self._append_log(f"Saved CSV: {csv_path}")
        self._last_csv_path = csv_path

        # Load data into the QAbstractTableModel
        self._data_model.load_data(temp_data, cols)
        self._data_row_count = len(temp_data)
        self._data_col_count = cols

        # Emit parsed data to QML
        self.parsedDataReady.emit({
            "success": True,
            "rows": len(temp_data),
            "cols": cols,
            "data": temp_data,
            "csv_path": csv_path,
            "familyCode": fc,
            "serialNumber": serial,
        })
        self._save_status()

    # ---- Status Persistence ----

    def _save_status(self) -> None:
        """Persist current operational status to settings."""
        # Update settings with current status values
        self._settings.connection_status = self._connection_status
        self._settings.selected_port = self._selected_port
        self._settings.last_wafer_info = self._last_wafer_info.copy()
        self._settings.save_data_dir = self._save_data_dir
        self._settings.activity_log = self._activity_log.copy()
        self._settings.data_row_count = self._data_row_count
        self._settings.data_col_count = self._data_col_count
        self._settings.last_csv_path = self._last_csv_path

        # Save to disk
        LocalSettings.save_settings(self._data_dir, self._settings)

    # ---- Helpers ----

    @staticmethod
    def _wafer_info_to_dict(info: WaferInfo) -> dict[str, Any]:
        """Convert WaferInfo dataclass to JSON-serialisable dict for QML."""
        from pygui.serialcomm.data_parser import csv_column_count
        sensor_count = info.sensor_count
        mapped_count = csv_column_count(info.family_code)
        if mapped_count > 0:
            sensor_count = mapped_count

        return {
            "familyCode": info.family_code,
            "serialNumber": info.serial_number,
            "sensorCount": sensor_count,
            "mfgDateHex": info.mfg_date_hex,
            "runtime": info.runtime,
            "cycleCount": info.cycle_count,
        }

class QmlActivityLogHandler(logging.Handler):
    def __init__(self, controller: DeviceController) -> None:
        super().__init__()
        self._controller = controller
    def emit(self, record: logging.LogRecord) -> None:
        message = record.getMessage()
        if "found wafer on" in message:
            return
        level = record.levelname
        msg = f"{level}: {message}"
        self._controller._append_log(msg)