"""
Coherent HOPS Laser I2C Control Library

This module provides a Python interface to control Coherent HOPS laser systems
via I2C protocol through an FTDI FT2232C USB interface.

Dependencies:
    pip install pyftdi

Usage:
    from coherent_hops_laser import CoherentHOPSLaser

    laser = CoherentHOPSLaser()
    laser.connect()

    # Query system information
    model = laser.get_laser_model()
    wavelength = laser.get_wavelength()

    # Monitor temperatures
    main_temp = laser.get_temperature_main()

    # Control power
    laser.set_power_command(100.0)  # Set power in mW or W

    laser.disconnect()
"""

from typing import Optional, Union, List
from dataclasses import dataclass
from enum import Enum
import time
import logging

try:
    from pyftdi.i2c import I2cController, I2cNackError
except ImportError:
    raise ImportError(
        "pyftdi library is required. Install with: pip install pyftdi"
    )


# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class I2CMode(Enum):
    """I2C communication modes"""
    STANDARD = 100000  # 100 kHz
    FAST = 400000      # 400 kHz


class ControlMode(Enum):
    """Laser control modes"""
    POWER = "POWER"
    CURRENT = "CURRENT"


@dataclass
class LaserInfo:
    """Laser system information"""
    hardware_id: Optional[str] = None
    head_type: Optional[str] = None
    head_board_revision: Optional[str] = None
    laser_model: Optional[str] = None
    power_units: Optional[str] = None
    wavelength: Optional[str] = None


@dataclass
class TemperatureStatus:
    """Temperature monitoring data"""
    main: Optional[float] = None
    brf: Optional[float] = None
    shg: Optional[float] = None
    thg: Optional[float] = None
    eta: Optional[float] = None


class CoherentHOPSLaser:
    """
    Main interface class for Coherent HOPS laser control via I2C.

    This class provides high-level methods for all documented laser commands
    including system queries, temperature control, power/current management,
    and digital I/O operations.
    """

    # Default I2C slave address for NXP microcontroller
    DEFAULT_SLAVE_ADDRESS = 0x50

    # FTDI USB VID/PID for FT2232C
    FTDI_VID = 0x0403
    FTDI_PID = 0x6010

    def __init__(
        self,
        slave_address: int = DEFAULT_SLAVE_ADDRESS,
        i2c_mode: I2CMode = I2CMode.STANDARD,
        timeout: float = 1.0
    ):
        """
        Initialize the laser controller.

        Args:
            slave_address: I2C slave address of the NXP microcontroller
            i2c_mode: I2C communication speed mode
            timeout: Command timeout in seconds
        """
        self.slave_address = slave_address
        self.i2c_mode = i2c_mode
        self.timeout = timeout

        self._i2c_controller = I2cController()
        self._i2c_slave = None
        self._connected = False

    def connect(self, url: str = 'ftdi://ftdi:2232/1') -> None:
        """
        Connect to the FTDI I2C device.

        Args:
            url: FTDI device URL (default: first FT2232C device, channel 1)
                Examples:
                - 'ftdi://ftdi:2232/1' - First FT2232 device, channel 1
                - 'ftdi://ftdi:2232:SERIAL/1' - Device with specific serial number

        Raises:
            IOError: If connection fails
        """
        try:
            # Configure I2C controller
            self._i2c_controller.configure(url, frequency=self.i2c_mode.value)

            # Get I2C slave interface
            self._i2c_slave = self._i2c_controller.get_port(self.slave_address)

            self._connected = True
            logger.info(f"Connected to laser at I2C address 0x{self.slave_address:02X}")

        except Exception as e:
            logger.error(f"Failed to connect to I2C device: {e}")
            raise IOError(f"I2C connection failed: {e}")

    def disconnect(self) -> None:
        """Disconnect from the I2C device."""
        if self._connected:
            self._i2c_controller.terminate()
            self._connected = False
            logger.info("Disconnected from laser")

    def _ensure_connected(self) -> None:
        """Verify device is connected before operations."""
        if not self._connected:
            raise RuntimeError("Not connected. Call connect() first.")

    def _send_command(self, command: str, value: Optional[str] = None) -> str:
        """
        Send a command to the laser and read response.

        Args:
            command: Command string (e.g., 'HID', 'LASERMODEL')
            value: Optional value for set commands

        Returns:
            Response string from the laser

        Raises:
            I2cNackError: If I2C communication fails
            TimeoutError: If response timeout occurs
        """
        self._ensure_connected()

        # Format command: query = ?COMMAND, set = COMMAND=VALUE
        if value is not None:
            cmd_str = f"{command}={value}"
        else:
            cmd_str = f"?{command}"

        cmd_bytes = cmd_str.encode('ascii')

        try:
            # Write command
            self._i2c_slave.write(cmd_bytes)

            # Small delay for laser to process
            time.sleep(0.01)

            # Read response (max 256 bytes)
            response = self._i2c_slave.read(256)

            # Decode and strip null bytes and whitespace
            result = response.decode('ascii', errors='ignore').rstrip('\x00').strip()

            logger.debug(f"Command: {cmd_str} -> Response: {result}")
            return result

        except I2cNackError as e:
            logger.error(f"I2C NACK error for command {cmd_str}: {e}")
            raise
        except Exception as e:
            logger.error(f"Communication error for command {cmd_str}: {e}")
            raise

    # ==========================================
    # System Information Commands
    # ==========================================

    def get_hardware_id(self) -> str:
        """Query Hardware ID."""
        return self._send_command('HID')

    def get_head_type(self) -> str:
        """Query Head Type."""
        return self._send_command('HTYPE')

    def get_head_board_revision(self) -> str:
        """Query Head Board Revision."""
        return self._send_command('HBDREV')

    def get_head_digital_io(self) -> str:
        """Query Head Digital I/O configuration."""
        return self._send_command('HEADDIO')

    def get_laser_model(self) -> str:
        """
        Query Laser Model.

        Returns:
            Model name (e.g., 'G532', 'Tina', 'Mini00', 'MiniX')
        """
        return self._send_command('LASERMODEL')

    def get_power_units(self) -> str:
        """
        Query Power Units.

        Returns:
            Power units (e.g., 'mW', 'W')
        """
        return self._send_command('POWERUNITS')

    def get_wavelength(self) -> str:
        """
        Query Wavelength.

        Returns:
            Wavelength (e.g., '532nm')
        """
        return self._send_command('WAVELENGTH')

    def get_system_info(self) -> LaserInfo:
        """
        Query all system information.

        Returns:
            LaserInfo dataclass with all system parameters
        """
        return LaserInfo(
            hardware_id=self.get_hardware_id(),
            head_type=self.get_head_type(),
            head_board_revision=self.get_head_board_revision(),
            laser_model=self.get_laser_model(),
            power_units=self.get_power_units(),
            wavelength=self.get_wavelength()
        )

    # ==========================================
    # Temperature Monitoring
    # ==========================================

    def get_temperature_main(self) -> float:
        """
        Get Main Heatsink Temperature.

        Returns:
            Temperature in degrees Celsius
        """
        response = self._send_command('TMAIN')
        return float(response)

    def get_temperature_brf(self) -> float:
        """
        Get BRF (Birefringent Filter) Temperature.

        Returns:
            Temperature in degrees Celsius
        """
        response = self._send_command('TBRF')
        return float(response)

    def get_temperature_shg(self) -> float:
        """
        Get SHG (Second Harmonic Generator) Temperature.

        Returns:
            Temperature in degrees Celsius
        """
        response = self._send_command('TSHG')
        return float(response)

    def get_temperature_thg(self) -> float:
        """
        Get THG (Third Harmonic Generator) Temperature.

        Returns:
            Temperature in degrees Celsius
        """
        response = self._send_command('TTHG')
        return float(response)

    def get_temperature_eta(self) -> float:
        """
        Get ETA Temperature.

        Returns:
            Temperature in degrees Celsius
        """
        response = self._send_command('TETA')
        return float(response)

    def get_all_temperatures(self) -> TemperatureStatus:
        """
        Query all temperature sensors.

        Returns:
            TemperatureStatus dataclass with all temperature readings
        """
        return TemperatureStatus(
            main=self.get_temperature_main(),
            brf=self.get_temperature_brf(),
            shg=self.get_temperature_shg(),
            thg=self.get_temperature_thg(),
            eta=self.get_temperature_eta()
        )

    # ==========================================
    # Temperature Control (Setpoints)
    # ==========================================

    def get_temperature_setpoint_main(self) -> float:
        """Get Main Temperature Setpoint."""
        response = self._send_command('TMAINCMD')
        return float(response)

    def set_temperature_setpoint_main(self, temperature: float) -> None:
        """Set Main Temperature Setpoint."""
        self._send_command('TMAINCMD', str(temperature))

    def get_temperature_setpoint_brf(self) -> float:
        """Get BRF Temperature Setpoint."""
        response = self._send_command('TBRFCMD')
        return float(response)

    def set_temperature_setpoint_brf(self, temperature: float) -> None:
        """Set BRF Temperature Setpoint."""
        self._send_command('TBRFCMD', str(temperature))

    def get_temperature_setpoint_shg(self) -> float:
        """Get SHG Temperature Setpoint."""
        response = self._send_command('TSHGCMD')
        return float(response)

    def set_temperature_setpoint_shg(self, temperature: float) -> None:
        """Set SHG Temperature Setpoint."""
        self._send_command('TSHGCMD', str(temperature))

    def get_temperature_setpoint_thg(self) -> float:
        """Get THG Temperature Setpoint."""
        response = self._send_command('TTHGCMD')
        return float(response)

    def set_temperature_setpoint_thg(self, temperature: float) -> None:
        """Set THG Temperature Setpoint."""
        self._send_command('TTHGCMD', str(temperature))

    def get_temperature_setpoint_eta(self) -> float:
        """Get ETA Temperature Setpoint."""
        response = self._send_command('TETACMD')
        return float(response)

    def set_temperature_setpoint_eta(self, temperature: float) -> None:
        """Set ETA Temperature Setpoint."""
        self._send_command('TETACMD', str(temperature))

    # ==========================================
    # Temperature Data
    # ==========================================

    def get_temperature_data_main(self) -> str:
        """Get Main Temperature Data."""
        return self._send_command('MAIND')

    def get_temperature_data_brf(self) -> str:
        """Get BRF Temperature Data."""
        return self._send_command('BRFD')

    def get_temperature_data_shg(self) -> str:
        """Get SHG Temperature Data."""
        return self._send_command('SHGD')

    def get_temperature_data_thg(self) -> str:
        """Get THG Temperature Data."""
        return self._send_command('THGD')

    def get_temperature_data_eta(self) -> str:
        """Get ETA Temperature Data."""
        return self._send_command('ETAD')

    # ==========================================
    # Power Control
    # ==========================================

    def get_power_command(self) -> float:
        """
        Get Power Command value.

        Returns:
            Power value in units from get_power_units()
        """
        response = self._send_command('PCMD')
        return float(response)

    def set_power_command(self, power: float) -> None:
        """
        Set Power Command value.

        Args:
            power: Power value in units from get_power_units()
        """
        self._send_command('PCMD', str(power))

    def get_power_memory(self) -> str:
        """Query Power Memory (stored settings)."""
        return self._send_command('PMEM')

    def get_power_limits(self) -> str:
        """Query Power Limits."""
        return self._send_command('PLIM')

    # ==========================================
    # Current Control
    # ==========================================

    def get_current_command(self) -> float:
        """
        Get Current Command value.

        Returns:
            Current value in Amperes
        """
        response = self._send_command('CCMD')
        return float(response)

    def set_current_command(self, current: float) -> None:
        """
        Set Current Command value.

        Args:
            current: Current value in Amperes
        """
        self._send_command('CCMD', str(current))

    def get_current_limits(self) -> str:
        """Query Current Limits."""
        return self._send_command('CLIM')

    def get_control_mode(self) -> str:
        """
        Get Control Mode.

        Returns:
            Control mode (e.g., 'POWER' or 'CURRENT')
        """
        return self._send_command('CMODE')

    def set_control_mode(self, mode: Union[str, ControlMode]) -> None:
        """
        Set Control Mode.

        Args:
            mode: Control mode ('POWER' or 'CURRENT', or ControlMode enum)
        """
        if isinstance(mode, ControlMode):
            mode = mode.value
        self._send_command('CMODE', mode)

    def get_control_mode_command(self) -> str:
        """Get Control Mode Command."""
        return self._send_command('CMODECMD')

    def set_control_mode_command(self, mode: str) -> None:
        """Set Control Mode Command."""
        self._send_command('CMODECMD', mode)

    # ==========================================
    # Digital I/O
    # ==========================================

    def get_ps_digital_io(self) -> str:
        """Get Power Supply Digital I/O status."""
        return self._send_command('PSDIO')

    def get_ps_glue_input(self) -> str:
        """Get Power Supply Glue Logic Input status."""
        return self._send_command('PSGLUEIN')

    def get_ps_glue_output(self) -> str:
        """Get Power Supply Glue Logic Output status."""
        return self._send_command('PSGLUEOUT')

    def set_ps_glue_output(self, value: str) -> None:
        """Set Power Supply Glue Logic Output."""
        self._send_command('PSGLUEOUT', value)

    # ==========================================
    # Monitoring & Status
    # ==========================================

    def get_analog_values(self) -> str:
        """Query Analog Values."""
        return self._send_command('ANA')

    def get_analog_command(self) -> str:
        """Get Analog Command."""
        return self._send_command('ANACMD')

    def set_analog_command(self, value: str) -> None:
        """Set Analog Command."""
        self._send_command('ANACMD', value)

    def get_key_switch_status(self) -> str:
        """Get Key Switch Status."""
        return self._send_command('KSW')

    def get_key_switch_command(self) -> str:
        """Get Key Switch Command."""
        return self._send_command('KSWCMD')

    def set_key_switch_command(self, value: str) -> None:
        """Set Key Switch Command."""
        self._send_command('KSWCMD', value)

    def get_fan_status(self) -> str:
        """Get Fan Status/Control."""
        return self._send_command('FAN')

    def set_fan_control(self, value: str) -> None:
        """Set Fan Control."""
        self._send_command('FAN', value)

    def get_interlock_status(self) -> str:
        """Get Interlock Status."""
        return self._send_command('INT')

    def get_remote_control_status(self) -> str:
        """Get Remote Control Status."""
        return self._send_command('REM')

    def get_eeprom_header(self) -> str:
        """Get EEPROM Header."""
        return self._send_command('EEH')

    # ==========================================
    # Configuration Registers
    # ==========================================

    def get_config_register_0(self) -> str:
        """Get Configuration Register 0."""
        return self._send_command('CFG0')

    def set_config_register_0(self, value: str) -> None:
        """Set Configuration Register 0."""
        self._send_command('CFG0', value)

    def get_config_register_1(self) -> str:
        """Get Configuration Register 1."""
        return self._send_command('CFG1')

    def set_config_register_1(self, value: str) -> None:
        """Set Configuration Register 1."""
        self._send_command('CFG1', value)

    def get_config_register_2(self) -> str:
        """Get Configuration Register 2."""
        return self._send_command('CFG2')

    def set_config_register_2(self, value: str) -> None:
        """Set Configuration Register 2."""
        self._send_command('CFG2', value)

    def get_config_register_3(self) -> str:
        """Get Configuration Register 3."""
        return self._send_command('CFG3')

    def set_config_register_3(self, value: str) -> None:
        """Set Configuration Register 3."""
        self._send_command('CFG3', value)

    # ==========================================
    # Context Manager Support
    # ==========================================

    def __enter__(self):
        """Context manager entry."""
        if not self._connected:
            self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        self.disconnect()
        return False


class DummyLaser(CoherentHOPSLaser):
    """
    Simulated laser for testing without hardware.

    This class provides dummy responses for all commands to enable
    software development and testing without physical hardware.
    """

    def __init__(self):
        """Initialize dummy laser (no I2C connection needed)."""
        super().__init__()
        self._connected = True  # Simulate connection

        # Simulated state
        self._power_cmd = 50.0
        self._current_cmd = 1.5
        self._control_mode = 'POWER'
        self._temps = {
            'main': 25.0,
            'brf': 30.0,
            'shg': 35.0,
            'thg': 32.0,
            'eta': 28.0
        }
        self._temp_setpoints = {
            'main': 25.0,
            'brf': 30.0,
            'shg': 35.0,
            'thg': 32.0,
            'eta': 28.0
        }

    def connect(self, url: str = 'dummy') -> None:
        """Dummy connection (always succeeds)."""
        self._connected = True
        logger.info("Connected to DummyLaser (simulation mode)")

    def _send_command(self, command: str, value: Optional[str] = None) -> str:
        """Simulate command responses."""
        logger.debug(f"DummyLaser command: {command}, value: {value}")

        # Handle set commands
        if value is not None:
            if command == 'PCMD':
                self._power_cmd = float(value)
                return value
            elif command == 'CCMD':
                self._current_cmd = float(value)
                return value
            elif command == 'CMODE':
                self._control_mode = value
                return value
            elif 'CMD' in command and 'T' in command:
                # Temperature setpoint
                key = command.replace('CMD', '').replace('T', '').lower()
                if key in self._temp_setpoints:
                    self._temp_setpoints[key] = float(value)
                return value
            return 'OK'

        # Handle query commands
        responses = {
            'HID': 'HOPS-12345',
            'HTYPE': 'Standard',
            'HBDREV': 'Rev 2.1',
            'HEADDIO': '0xFF',
            'LASERMODEL': 'G532',
            'POWERUNITS': 'mW',
            'WAVELENGTH': '532nm',
            'TMAIN': str(self._temps['main']),
            'TBRF': str(self._temps['brf']),
            'TSHG': str(self._temps['shg']),
            'TTHG': str(self._temps['thg']),
            'TETA': str(self._temps['eta']),
            'TMAINCMD': str(self._temp_setpoints['main']),
            'TBRFCMD': str(self._temp_setpoints['brf']),
            'TSHGCMD': str(self._temp_setpoints['shg']),
            'TTHGCMD': str(self._temp_setpoints['thg']),
            'TETACMD': str(self._temp_setpoints['eta']),
            'MAIND': 'MainTempData',
            'BRFD': 'BRFTempData',
            'SHGD': 'SHGTempData',
            'THGD': 'THGTempData',
            'ETAD': 'ETATempData',
            'PCMD': str(self._power_cmd),
            'PMEM': '100',
            'PLIM': '0-200',
            'CCMD': str(self._current_cmd),
            'CLIM': '0-5',
            'CMODE': self._control_mode,
            'CMODECMD': self._control_mode,
            'PSDIO': '0x00',
            'PSGLUEIN': '0x00',
            'PSGLUEOUT': '0x00',
            'ANA': '0,0,0,0',
            'ANACMD': '0',
            'KSW': 'ON',
            'KSWCMD': 'ON',
            'FAN': 'AUTO',
            'INT': 'OK',
            'REM': 'ENABLED',
            'EEH': 'EEPROM_V1',
            'CFG0': '0x00',
            'CFG1': '0x00',
            'CFG2': '0x00',
            'CFG3': '0x00',
        }

        return responses.get(command, 'UNKNOWN')


if __name__ == '__main__':
    """Example usage and testing."""

    # Test with dummy laser
    print("=== Testing with DummyLaser ===\n")

    with DummyLaser() as laser:
        # System info
        print("System Information:")
        info = laser.get_system_info()
        print(f"  Model: {info.laser_model}")
        print(f"  Wavelength: {info.wavelength}")
        print(f"  Power Units: {info.power_units}")
        print(f"  Hardware ID: {info.hardware_id}")
        print()

        # Temperature monitoring
        print("Temperature Status:")
        temps = laser.get_all_temperatures()
        print(f"  Main: {temps.main}°C")
        print(f"  BRF: {temps.brf}°C")
        print(f"  SHG: {temps.shg}°C")
        print(f"  THG: {temps.thg}°C")
        print(f"  ETA: {temps.eta}°C")
        print()

        # Power control
        print("Power Control:")
        current_power = laser.get_power_command()
        print(f"  Current Power: {current_power} {info.power_units}")
        laser.set_power_command(75.0)
        new_power = laser.get_power_command()
        print(f"  New Power: {new_power} {info.power_units}")
        print()

        # Control mode
        print("Control Mode:")
        mode = laser.get_control_mode()
        print(f"  Current Mode: {mode}")
        print()

    print("\n=== For real hardware, use: ===")
    print("laser = CoherentHOPSLaser()")
    print("laser.connect('ftdi://ftdi:2232/1')")
    print("# ... perform operations ...")
    print("laser.disconnect()")