"""
Helios Laser System Driver
Basic implementation for controlling the Helios pulsed laser.
"""

import time
import logging
from typing import Optional, List
from enum import Enum

from hardware.serial_util import open_8n1

logger = logging.getLogger(__name__)


class PulseMode(Enum):
    """Helios pulse mode enumeration"""
    SINGLE_PULSE = 0
    CONTINUOUS_GATING = 1
    CONTINUOUS_PULSING = 2


class HeliosLaser:
    """
    Driver for Helios pulsed laser system.

    Communication: RS-232, 9600 baud, 8N1
    Commands are ASCII strings terminated with CR
    """

    def __init__(self, port: str = None, timeout: float = 1.0):
        """
        Initialize Helios laser driver.

        Args:
            port: Serial port (e.g., '/dev/ttyUSB0' or 'COM5')
            timeout: Serial timeout in seconds
        """
        self.port = port
        self.timeout = timeout
        self.serial = None
        self.is_connected = False

    @staticmethod
    def list_available_ports() -> List[str]:
        """List available serial ports, likeliest devices first."""
        from hardware.serial_util import list_port_devices
        return list_port_devices()

    def connect(self, port: str = None) -> bool:
        """
        Connect to the Helios laser.

        Args:
            port: Serial port (uses stored port if None)

        Returns:
            True if connection successful
        """
        if port:
            self.port = port

        if not self.port:
            logger.error("No port specified")
            return False

        try:
            self.serial = open_8n1(self.port, baudrate=9600, timeout=self.timeout)
            time.sleep(0.1)  # Allow time for connection to stabilize
            self.is_connected = True
            logger.info(f"Connected to Helios laser on {self.port}")
            return True

        except Exception as e:
            logger.error(f"Failed to connect to Helios laser: {e}")
            self.is_connected = False
            return False

    def disconnect(self):
        """Disconnect from the laser"""
        if self.serial and self.serial.is_open:
            try:
                # Disable laser before disconnecting
                self.set_laser_enable(False)
                self.serial.close()
                logger.info("Disconnected from Helios laser")
            except Exception as e:
                logger.error(f"Error during disconnect: {e}")

        self.is_connected = False
        self.serial = None

    def _send_command(self, command: str) -> bool:
        """
        Send a command to the laser.

        Args:
            command: ASCII command string (without CR)

        Returns:
            True if sent successfully
        """
        if not self.is_connected or not self.serial:
            logger.error("Not connected to laser")
            return False

        try:
            cmd_bytes = (command + '\r').encode('ascii')
            self.serial.write(cmd_bytes)
            logger.debug(f"Sent command: {command}")
            return True

        except Exception as e:
            logger.error(f"Failed to send command '{command}': {e}")
            return False

    def _query(self, command: str) -> Optional[str]:
        """
        Send a query and read response.

        Args:
            command: ASCII query command (without CR or ?)

        Returns:
            Response value string or None if error
        """
        try:
            # Clear any pending data in the buffer
            self.serial.reset_input_buffer()
            time.sleep(0.05)

            if not self._send_command(command):
                return None

            time.sleep(0.2)  # Give device time to respond

            response = self.serial.read_until(b'\r').decode('ascii', errors='replace').strip()
            logger.debug(f"Query '{command}' response: {response}")

            # Helios format: "COMMAND = VALUE UNIT"
            # Extract just the value part
            if '=' in response:
                parts = response.split('=')
                if len(parts) >= 2:
                    value_part = parts[1].strip()
                    # Remove unit suffix if present (e.g., "ns", "mA", "mW")
                    value = value_part.split()[0]
                    return value

            return response

        except Exception as e:
            logger.error(f"Failed to read response for '{command}': {e}")
            return None

    def set_frequency_hz(self, frequency: int) -> bool:
        """
        Set laser pulse frequency in Hz.

        Args:
            frequency: Frequency in Hz (16700 - 125000)

        Returns:
            True if successful
        """
        if not (16700 <= frequency <= 125000):
            logger.error(f"Frequency {frequency} Hz out of range (16700-125000)")
            return False

        # Convert frequency to period in nanoseconds
        period_ns = int(1e9 / frequency)

        # Clamp to valid range (8000-60000 ns)
        if not (8000 <= period_ns <= 60000):
            logger.error(f"Period {period_ns} ns out of range (8000-60000)")
            return False

        command = f"LDF {period_ns}"
        return self._send_command(command)

    def set_current_ma(self, current: int) -> bool:
        """
        Set pump diode current in mA.

        Args:
            current: Current in mA (0 - 2000 for this model)

        Returns:
            True if successful
        """
        if not (0 <= current <= 2000):
            logger.error(f"Current {current} mA out of range (0-2000)")
            return False

        command = f"LDS {current}"
        return self._send_command(command)

    def set_pulse_mode(self, mode: PulseMode) -> bool:
        """
        Set pulse mode.

        Args:
            mode: PulseMode enumeration value

        Returns:
            True if successful
        """
        command = f"LDG {mode.value}"
        return self._send_command(command)

    def set_laser_enable(self, enable: bool) -> bool:
        """
        Enable or disable laser emission.

        Args:
            enable: True to enable, False to disable

        Returns:
            True if successful
        """
        command = f"LDO {1 if enable else 0}"
        success = self._send_command(command)

        if success:
            state = "enabled" if enable else "disabled"
            logger.info(f"Laser {state}")

        return success

    def is_laser_enabled(self) -> bool:
        """
        Check if laser is currently enabled.

        Returns:
            True if laser is enabled
        """
        response = self._query("LDO")
        if response:
            try:
                return int(response) == 1
            except ValueError:
                logger.error(f"Invalid response for LDO: {response}")
        return False

    def get_frequency_hz(self) -> Optional[int]:
        """
        Get current laser frequency in Hz.

        Returns:
            Frequency in Hz or None if error
        """
        response = self._query("LDF")
        if response:
            try:
                period_ns = int(response)
                return int(1e9 / period_ns)
            except (ValueError, ZeroDivisionError):
                logger.error(f"Invalid response for LDF: {response}")
        return None

    def get_current_ma(self) -> Optional[int]:
        """
        Get current pump diode current in mA.

        Returns:
            Current in mA or None if error
        """
        response = self._query("LDS")
        if response:
            try:
                return int(response)
            except ValueError:
                logger.error(f"Invalid response for LDS: {response}")
        return None

    def _query_millicelsius(self, command: str) -> Optional[float]:
        """Query a temperature register (returns milli-°C) and convert to °C."""
        response = self._query(command)
        if response:
            try:
                return int(response) / 1000.0
            except ValueError:
                logger.error(f"Invalid response for {command}: {response}")
        return None

    def get_diode_temp_c(self) -> Optional[float]:
        """Diode temperature in °C (LTA, 5000–50000 milli-°C)."""
        return self._query_millicelsius("LTA")

    def get_power_stage_temp_c(self) -> Optional[float]:
        """Power stage temperature in °C (LTT, 5000–65355 milli-°C)."""
        return self._query_millicelsius("LTT")

    def get_qswitch_temp_c(self) -> Optional[float]:
        """Q-switch temperature in °C (EOA, 5000–50000 milli-°C)."""
        return self._query_millicelsius("EOA")

    def get_controller_serial(self) -> Optional[str]:
        """
        Get controller serial number.

        Returns:
            Serial number string or None if error
        """
        return self._query("CSR")

    def get_head_serial(self) -> Optional[str]:
        """
        Get laser head serial number.

        Returns:
            Serial number string or None if error
        """
        return self._query("HSR")

    def get_status_registers(self) -> tuple:
        """
        Query LER, LCE, and CCE status registers.

        Each register is a bitmask (sum of flags). Non-zero values indicate
        active faults. Reset with reset_faults().

        Returns:
            Tuple of (ler, lce, cce) as ints, or None for each on error.
        """
        def _read_reg(cmd):
            resp = self._query(cmd)
            if resp is not None:
                try:
                    return int(resp)
                except ValueError:
                    logger.error(f"Invalid response for {cmd}: {resp}")
            return None

        ler = _read_reg("LER")
        lce = _read_reg("LCE")
        cce = _read_reg("CCE")
        return (ler, lce, cce)

    def reset_faults(self) -> bool:
        """
        Execute the controller reset sequence to clear status registers.

        Protocol-specified sequence: CCE 0 -> LCE 0 -> LER 0

        Returns:
            True if all three commands sent successfully
        """
        ok = True
        ok = self._send_command("CCE 0") and ok
        time.sleep(0.1)
        ok = self._send_command("LCE 0") and ok
        time.sleep(0.1)
        ok = self._send_command("LER 0") and ok
        if ok:
            logger.info("Fault reset sequence sent")
        return ok

    def get_remote_enable(self) -> Optional[bool]:
        """
        Query the remote enable state (LRE - activates utility connector pin 8).

        Returns:
            True if remote enable is active, False if not, None on error
        """
        response = self._query("LRE")
        if response is not None:
            try:
                return int(response) == 1
            except ValueError:
                logger.error(f"Invalid response for LRE: {response}")
        return None

    def send_raw_command(self, command: str) -> Optional[str]:
        """
        Send a raw command string and return the raw response.

        Useful for diagnostics. Sends *command* + CR, waits briefly,
        then reads whatever the device returns (up to the first CR or timeout).

        Returns:
            Raw response string (decoded, stripped) or None on error.
        """
        if not self.is_connected or not self.serial:
            logger.error("Not connected to laser")
            return None
        try:
            self.serial.reset_input_buffer()
            time.sleep(0.05)
            self.serial.write((command + '\r').encode('ascii'))
            time.sleep(0.3)
            raw = self.serial.read_until(b'\r')
            if not raw:
                raw = self.serial.read(self.serial.in_waiting)
            return raw.decode('ascii', errors='replace').strip()
        except Exception as e:
            logger.error(f"send_raw_command error: {e}")
            return None

    def set_remote_enable(self, enable: bool) -> bool:
        """
        Set the remote enable state (LRE - utility connector pin 8).

        Args:
            enable: True to activate remote enable, False to deactivate

        Returns:
            True if successful
        """
        command = f"LRE {1 if enable else 0}"
        return self._send_command(command)

    def __del__(self):
        """Destructor - ensure cleanup"""
        self.disconnect()