"""
Genesis SLM MX 532 Laser Core Hardware Control Module
======================================================

.. warning::
   QUARANTINED — do not modify semantics or dedupe against
   ``tools/genesis_laser_gui.py`` until the bench checklist in
   ``docs/genesis_verification.md`` has been run.  This module was
   extracted from that GUI but diverges from it in ways only hardware can
   adjudicate: ADS7828 command byte (0x84 vs 0x8C), LDD enable polarity
   (inverted vs not), shutter semantics (manual vs bit-controlled),
   dropped median-of-3 ADC filtering, dropped Amps/Watts scaling, dropped
   temperature reads and pre-flight check.

This module provides low-level hardware control for the Genesis SLM MX 532 laser
using NXP I2C-over-serial protocol. It contains reusable classes for serial
communication, I2C protocol handling, device control, and laser operations.

This module is GUI-independent and can be used by both command-line and GUI applications.

Protocol Overview:
-----------------
The laser uses NXP I2C-over-serial protocol with packet format:
[0x53] [I2C_ADDR] [LENGTH] [COMMAND_BYTES] [DATA_BYTES] [0x50]

For reads:
[0x53] [ADDR_WRITE] [CMD_LEN] [CMD] [0x53] [ADDR_READ] [DATA_LEN] [0x50]

I2C Devices:
-----------
- X9119 digital potentiometer at 0x52 - laser current control
- PCA9555 I/O expander at 0x4a - digital I/O
- AD5254 digital potentiometer at 0x58 - limits control
- ADS7828 ADC at 0x90 - sensor readings
- M24C64 EEPROM at 0xa4 - configuration storage

Author: Claude
Date: 2026-01-24
"""

import time
from datetime import datetime
from typing import Optional, List
from enum import IntEnum

import serial


# ============================================================================
# Constants - I2C Addresses
# ============================================================================

class I2CAddress(IntEnum):
    """I2C device addresses (7-bit addresses shifted left by 1)"""
    # X9119 Digital Potentiometer (current control)
    X9119_CURRENT_WRITE = 0x52
    X9119_CURRENT_READ = 0x53

    # PCA9555 I/O Expanders
    PCA9555_DIO_WRITE = 0x4a      # Main control I/O at 0x14a
    PCA9555_DIO_READ = 0x4b
    PCA9555_PSGLUE_WRITE = 0x48   # Power supply glue at 0x148
    PCA9555_PSGLUE_READ = 0x49
    PCA9555_HEAD_WRITE = 0x44     # Head DIO at 0x144
    PCA9555_HEAD_READ = 0x45
    PCA9555_LDD_WRITE = 0x40      # LDD control at 0x140
    PCA9555_LDD_READ = 0x41

    # AD5254 Digital Potentiometer (limits)
    AD5254_WRITE = 0x58

    # ADS7828 ADC (sensor readings)
    ADS7828_WRITE = 0x90
    ADS7828_READ = 0x91

    # M24C64 EEPROM (configuration)
    EEPROM_WRITE = 0xa4
    EEPROM_READ = 0xa5


# ============================================================================
# Constants - PCA9555 Register Addresses
# ============================================================================

class PCA9555Register(IntEnum):
    """PCA9555 I/O expander register addresses"""
    INPUT_PORT_0 = 0x00
    INPUT_PORT_1 = 0x01
    OUTPUT_PORT_0 = 0x02
    OUTPUT_PORT_1 = 0x03
    POLARITY_INV_0 = 0x04
    POLARITY_INV_1 = 0x05
    CONFIG_PORT_0 = 0x06
    CONFIG_PORT_1 = 0x07


# ============================================================================
# Constants - Control Bitmasks
# ============================================================================

class ControlBitmask(IntEnum):
    """Bitmasks for PCA9555 port 0 control bits"""
    SHUTTER = 0x01        # Bit 0
    CURRENT_MODE = 0x04   # Bit 2
    REMOTE_ENABLE = 0x08  # Bit 3
    ANALOG_ENABLE = 0x10  # Bit 4
    KEYSWITCH = 0x20      # Bit 5
    INTERLOCK = 0x01      # Bit 0 (on different PCA9555)
    LDD_ENABLE = 0x01     # Bit 0 (on LDD PCA9555)


# ============================================================================
# Serial Communication Layer
# ============================================================================

class SerialComm:
    """Handles low-level serial port communication"""

    def __init__(self):
        self.port: Optional[serial.Serial] = None
        self.port_name: str = "/dev/ttyUSB0"
        self.baudrate: int = 9600
        self.timeout: float = 1.0
        self.packet_log: List[str] = []

    def connect(self, port_name: str, baudrate: int = 9600) -> bool:
        """
        Connect to serial port

        Args:
            port_name: Serial port device name
            baudrate: Baud rate (default 9600)

        Returns:
            True if successful, False otherwise
        """
        try:
            self.port_name = port_name
            self.baudrate = baudrate
            self.port = serial.Serial(
                port=port_name,
                baudrate=baudrate,
                bytesize=serial.EIGHTBITS,
                parity=serial.PARITY_NONE,
                stopbits=serial.STOPBITS_ONE,
                timeout=self.timeout
            )
            time.sleep(0.1)  # Allow port to stabilize
            return True
        except Exception as e:
            print(f"Serial connection error: {e}")
            return False

    def disconnect(self):
        """Disconnect from serial port"""
        if self.port and self.port.is_open:
            self.port.close()
            self.port = None

    def is_connected(self) -> bool:
        """Check if serial port is connected"""
        return self.port is not None and self.port.is_open

    def write(self, data: bytes) -> bool:
        """
        Write data to serial port

        Args:
            data: Bytes to write

        Returns:
            True if successful, False otherwise
        """
        if not self.is_connected():
            return False

        try:
            self.port.write(data)
            timestamp = datetime.now().strftime("%H:%M:%S.%f")[:-3]
            hex_str = " ".join(f"{b:02x}" for b in data)
            self.packet_log.append(f"[{timestamp}] TX: {hex_str}")
            return True
        except Exception as e:
            print(f"Serial write error: {e}")
            return False

    def read(self, size: int) -> Optional[bytes]:
        """
        Read data from serial port

        Args:
            size: Number of bytes to read

        Returns:
            Bytes read or None on error
        """
        if not self.is_connected():
            return None

        try:
            data = self.port.read(size)
            if data:
                timestamp = datetime.now().strftime("%H:%M:%S.%f")[:-3]
                hex_str = " ".join(f"{b:02x}" for b in data)
                self.packet_log.append(f"[{timestamp}] RX: {hex_str}")
            return data
        except Exception as e:
            print(f"Serial read error: {e}")
            return None

    def flush(self):
        """Flush serial port buffers"""
        if self.is_connected():
            self.port.reset_input_buffer()
            self.port.reset_output_buffer()

    def get_packet_log(self, last_n: int = 100) -> List[str]:
        """Get last N entries from packet log"""
        return self.packet_log[-last_n:]

    def clear_packet_log(self):
        """Clear packet log"""
        self.packet_log.clear()


# ============================================================================
# I2C Protocol Layer
# ============================================================================

class I2CProtocol:
    """Handles NXP I2C-over-serial protocol packet construction"""

    NXP_START = 0x53
    NXP_STOP = 0x50

    def __init__(self, serial_comm: SerialComm):
        self.serial = serial_comm

    def write(self, i2c_addr_write: int, cmd: bytes, data: bytes = b'') -> bool:
        """
        Build and send NXP I2C write packet

        Packet format: [0x53] [I2C_ADDR] [LENGTH] [COMMAND] [DATA] [0x50]

        Args:
            i2c_addr_write: I2C write address
            cmd: Command bytes
            data: Data bytes (optional)

        Returns:
            True if successful, False otherwise
        """
        if not isinstance(cmd, bytes):
            cmd = bytes([cmd])
        if not isinstance(data, bytes):
            data = bytes(data)

        length = len(cmd) + len(data)
        packet = bytes([self.NXP_START, i2c_addr_write, length]) + cmd + data + bytes([self.NXP_STOP])

        return self.serial.write(packet)

    def read(self, i2c_addr_write: int, cmd: bytes, cmd_len: int, data_len: int) -> Optional[bytes]:
        """
        Build and send NXP I2C read packet, return data

        Packet format for read:
        [0x53] [ADDR_WRITE] [CMD_LEN] [CMD] [0x53] [ADDR_READ] [DATA_LEN] [0x50]

        Args:
            i2c_addr_write: I2C write address
            cmd: Command bytes
            cmd_len: Length of command
            data_len: Expected data length to read

        Returns:
            Data bytes read or None on error
        """
        if not isinstance(cmd, bytes):
            cmd = bytes([cmd])

        i2c_addr_read = i2c_addr_write | 0x01  # Set read bit

        # First part: write command
        packet_write = bytes([self.NXP_START, i2c_addr_write, cmd_len]) + cmd

        # Second part: read data
        packet_read = bytes([self.NXP_START, i2c_addr_read, data_len, self.NXP_STOP])

        packet = packet_write + packet_read

        if not self.serial.write(packet):
            return None

        # Read response
        # Expected response: [data_bytes]
        time.sleep(0.05)  # Give device time to respond
        response = self.serial.read(data_len)

        return response if response and len(response) == data_len else None


# ============================================================================
# I2C Device Functions
# ============================================================================

class I2CDevices:
    """High-level I2C device control functions"""

    def __init__(self, protocol: I2CProtocol):
        self.protocol = protocol

    # ------------------------------------------------------------------------
    # PCA9555 I/O Expander Functions
    # ------------------------------------------------------------------------

    def pca9555_read_port(self, addr_write: int, port: int) -> Optional[int]:
        """
        Read PCA9555 port register

        Args:
            addr_write: I2C write address
            port: Register address (0-7)

        Returns:
            Port value (0-255) or None on error
        """
        data = self.protocol.read(addr_write, bytes([port]), 1, 1)
        return data[0] if data else None

    def pca9555_write_port(self, addr_write: int, port: int, value: int) -> bool:
        """
        Write PCA9555 port register

        Args:
            addr_write: I2C write address
            port: Register address (0-7)
            value: Value to write (0-255)

        Returns:
            True if successful, False otherwise
        """
        return self.protocol.write(addr_write, bytes([port]), bytes([value]))

    def pca9555_set_bit(self, addr_write: int, port: int, bitmask: int, state: bool) -> bool:
        """
        Set or clear specific bit on PCA9555

        Args:
            addr_write: I2C write address
            port: Register address
            bitmask: Bit mask (e.g., 0x01 for bit 0)
            state: True to set bit, False to clear

        Returns:
            True if successful, False otherwise
        """
        # Read current value
        current = self.pca9555_read_port(addr_write, port)
        if current is None:
            return False

        # Modify bit
        if state:
            new_value = current | bitmask
        else:
            new_value = current & ~bitmask

        # Write back
        return self.pca9555_write_port(addr_write, port, new_value)

    # ------------------------------------------------------------------------
    # X9119 Digital Potentiometer Functions
    # ------------------------------------------------------------------------

    def x9119_write_wiper(self, addr_write: int, value: int) -> bool:
        """
        Set X9119 wiper position

        Args:
            addr_write: I2C write address
            value: Wiper position (0-1023, 10-bit)

        Returns:
            True if successful, False otherwise
        """
        # X9119 write wiper command: 0xa0 followed by 2 bytes (10-bit value)
        # Value is split into two bytes: MSB contains upper 2 bits, LSB contains lower 8 bits
        value = max(0, min(1023, value))  # Clamp to valid range

        msb = (value >> 8) & 0x03  # Upper 2 bits
        lsb = value & 0xFF          # Lower 8 bits

        cmd = bytes([0xa0])
        data = bytes([msb, lsb])

        return self.protocol.write(addr_write, cmd, data)

    # ------------------------------------------------------------------------
    # ADS7828 ADC Functions
    # ------------------------------------------------------------------------

    def ads7828_read(self, addr_write: int, channel: int) -> Optional[int]:
        """
        Read ADS7828 ADC channel

        Args:
            addr_write: I2C write address
            channel: Channel to read (0-7)

        Returns:
            12-bit ADC value (0-4095) or None on error
        """
        # ADS7828 command byte format:
        # Bit 7: SD (single-ended=1)
        # Bits 6-4: Channel select
        # Bit 3: PD1 (power-down mode)
        # Bit 2: PD0 (power-down mode)
        # Bits 1-0: Don't care

        # For single-ended read with internal reference on
        cmd_byte = 0x80 | ((channel & 0x07) << 4) | 0x0c

        data = self.protocol.read(addr_write, bytes([cmd_byte]), 1, 2)
        if not data or len(data) != 2:
            return None

        # Combine two bytes into 12-bit value
        value = (data[0] << 8) | data[1]
        value = (value >> 4) & 0x0FFF  # Extract 12-bit value

        return value


# ============================================================================
# Laser Control Layer
# ============================================================================

class LaserControl:
    """High-level laser control functions"""

    def __init__(self, devices: I2CDevices):
        self.devices = devices
        self._safe_state_active = False

    # ------------------------------------------------------------------------
    # SET Commands
    # ------------------------------------------------------------------------

    def set_current(self, value: int) -> bool:
        """
        Set laser current command (CCMD=)

        Args:
            value: Current value (0-1023)

        Returns:
            True if successful, False otherwise
        """
        return self.devices.x9119_write_wiper(I2CAddress.X9119_CURRENT_WRITE, value)

    def set_power_cmd(self, value: int) -> bool:
        """
        Set power command (PCMD=)

        Note: Uses same device as current control, different channel

        Args:
            value: Power value (0-1023)

        Returns:
            True if successful, False otherwise
        """
        return self.devices.x9119_write_wiper(I2CAddress.X9119_CURRENT_WRITE, value)

    def set_shutter(self, state: bool) -> bool:
        """
        Set shutter state (SHCMD=)

        Args:
            state: True for open, False for closed

        Returns:
            True if successful, False otherwise
        """
        return self.devices.pca9555_set_bit(
            I2CAddress.PCA9555_DIO_WRITE,
            PCA9555Register.OUTPUT_PORT_0,
            ControlBitmask.SHUTTER,
            state
        )

    def set_keyswitch(self, state: bool) -> bool:
        """
        Set keyswitch state (KSWCMD=)

        Args:
            state: True for on, False for off

        Returns:
            True if successful, False otherwise
        """
        return self.devices.pca9555_set_bit(
            I2CAddress.PCA9555_DIO_WRITE,
            PCA9555Register.OUTPUT_PORT_0,
            ControlBitmask.KEYSWITCH,
            state
        )

    def set_current_mode(self, state: bool) -> bool:
        """
        Set current mode (CMODECMD=)

        Args:
            state: True for on, False for off

        Returns:
            True if successful, False otherwise
        """
        return self.devices.pca9555_set_bit(
            I2CAddress.PCA9555_DIO_WRITE,
            PCA9555Register.OUTPUT_PORT_0,
            ControlBitmask.CURRENT_MODE,
            state
        )

    def set_analog_enable(self, state: bool) -> bool:
        """
        Set analog input enable (ANACMD=)

        Args:
            state: True for enabled, False for disabled

        Returns:
            True if successful, False otherwise
        """
        return self.devices.pca9555_set_bit(
            I2CAddress.PCA9555_DIO_WRITE,
            PCA9555Register.OUTPUT_PORT_0,
            ControlBitmask.ANALOG_ENABLE,
            state
        )

    def set_remote_enable(self, state: bool) -> bool:
        """
        Set remote enable (REM=)

        Args:
            state: True for enabled, False for disabled

        Returns:
            True if successful, False otherwise
        """
        return self.devices.pca9555_set_bit(
            I2CAddress.PCA9555_DIO_WRITE,
            PCA9555Register.OUTPUT_PORT_0,
            ControlBitmask.REMOTE_ENABLE,
            state
        )

    # ------------------------------------------------------------------------
    # GET/Query Commands
    # ------------------------------------------------------------------------

    def get_current_actual(self) -> Optional[float]:
        """
        Get actual current reading

        Returns:
            Current in arbitrary units (0-4095) or None on error
        """
        # Read from ADS7828, channel based on command 0x84
        # Command 0x84 suggests channel 0
        value = self.devices.ads7828_read(I2CAddress.ADS7828_WRITE, 0)
        return value if value is not None else None

    def get_interlock_status(self) -> Optional[bool]:
        """
        Get interlock status

        Returns:
            True if interlock OK, False if fault, None on error
        """
        value = self.devices.pca9555_read_port(
            I2CAddress.PCA9555_PSGLUE_WRITE,
            PCA9555Register.INPUT_PORT_0
        )
        if value is None:
            return None

        # Bit 0 indicates interlock status
        return bool(value & ControlBitmask.INTERLOCK)

    def get_ldd_enable_status(self) -> Optional[bool]:
        """
        Get laser diode driver enable status

        Returns:
            True if enabled, False if disabled, None on error
        """
        value = self.devices.pca9555_read_port(
            I2CAddress.PCA9555_LDD_WRITE,
            PCA9555Register.INPUT_PORT_0
        )
        if value is None:
            return None

        return bool(value & ControlBitmask.LDD_ENABLE)

    def get_psglue_in_status(self) -> Optional[int]:
        """
        Get power supply glue input status

        Returns:
            Port value or None on error
        """
        return self.devices.pca9555_read_port(
            I2CAddress.PCA9555_PSGLUE_WRITE,
            PCA9555Register.INPUT_PORT_0
        )

    def get_psglue_out_status(self) -> Optional[int]:
        """
        Get power supply glue output status

        Returns:
            Port value or None on error
        """
        return self.devices.pca9555_read_port(
            I2CAddress.PCA9555_DIO_WRITE,
            PCA9555Register.OUTPUT_PORT_0
        )

    def get_head_dio_status(self) -> Optional[int]:
        """
        Get head DIO status

        Returns:
            Port value or None on error
        """
        return self.devices.pca9555_read_port(
            I2CAddress.PCA9555_HEAD_WRITE,
            PCA9555Register.INPUT_PORT_0
        )

    # ------------------------------------------------------------------------
    # Safety Functions
    # ------------------------------------------------------------------------

    def emergency_stop(self) -> bool:
        """
        Emergency stop: close shutter, set current to 0, disable keyswitch

        Returns:
            True if all operations successful, False otherwise
        """
        success = True
        success &= self.set_shutter(False)
        success &= self.set_current(0)
        success &= self.set_keyswitch(False)
        self._safe_state_active = True
        return success

    def enter_safe_state(self) -> bool:
        """
        Enter safe state (similar to emergency stop but also disables remote)

        Returns:
            True if successful, False otherwise
        """
        success = True
        success &= self.set_shutter(False)
        success &= self.set_current(0)
        success &= self.set_power_cmd(0)
        success &= self.set_keyswitch(False)
        success &= self.set_remote_enable(False)
        self._safe_state_active = True
        return success