#!/usr/bin/env python3
"""
SCPI Serial Interface - Easy-to-use serial connection for NEnG instruments

Provides a simple, reusable interface for communicating with SCPI instruments
via USB CDC serial connection with automatic port detection, error handling,
and convenient command/query methods.

Usage:
    from scpi_serial import SCPISerial

    # Auto-detect port
    with SCPISerial() as instr:
        print(instr.query("*IDN?"))
        instr.write("WIFI:ENABLE 1")

    # Or specify port
    with SCPISerial(port="/dev/cu.usbmodem2101") as instr:
        status = instr.query("WIFI:STATUS?")
        print(status)

(c) 2024-25 Prof. Flavio ABREU ARAUJO. All rights reserved.
"""

import contextlib
import os
import platform
import sys
import tempfile
import time

# Ensure we load the real PySerial package (and not the similarly named 'serial' stub)
try:
    import serial as pyserial

    # PySerial exposes Serial and SerialException; if they're missing we likely imported
    # the wrong package (e.g., the 'serial' stub from PyPI).
    if not hasattr(pyserial, "Serial") or not hasattr(pyserial, "SerialException"):
        raise ImportError(
            f"Imported 'serial' from {getattr(pyserial, '__file__', 'unknown')} but it lacks Serial/SerialException"
        )

    serial = pyserial  # Alias to keep the rest of the module unchanged
    serial.tools.list_ports = __import__("serial.tools.list_ports", fromlist=["list_ports"])
except ImportError as exc:  # pragma: no cover - defensive guard
    raise ImportError(
        "PySerial is missing or shadowed by another 'serial' package. "
        "Uninstall any 'serial' package and install the correct one with 'pip install pyserial'."
    ) from exc


class DeviceInUseError(RuntimeError):
    """Raised when the requested serial port is already held by another process.

    Attributes:
        port:      Port name, e.g. '/dev/ttyACM0' or 'COM3'.
        lock_path: Path to the lock file on disk.
    """

    def __init__(self, port: str, lock_path: str, pid_info: str = "") -> None:
        self.port = port
        self.lock_path = lock_path
        system = platform.system()
        delete_cmd = f'del "{lock_path}"' if system == "Windows" else f'rm "{lock_path}"'
        msg = (
            f"Port {port} is already in use by another process{pid_info}.\n"
            "\nPossible causes:\n"
            f"  - Another script or application holds {port} open\n"
            "  - A previous session did not close the connection cleanly\n"
            "\nSolutions:\n"
            f"  1. Close every other application or script that uses {port}\n"
            "  2. If you are certain nothing else is using the port, the lock\n"
            "     file may be stale. Remove it with:\n"
            f"       {delete_cmd}\n"
        )
        super().__init__(msg)


class PortLock:
    """Cross-platform exclusive lock for a serial port.

    Stores a lock file in the system temp directory as
    ``neng_lock_<port_basename>`` (e.g. ``/tmp/neng_lock_ttyACM0`` or
    ``C:\\...\\Temp\\neng_lock_COM3``).

    The underlying byte-range / advisory lock is automatically released by the
    OS when the process exits, so a stale lock *file* left behind by a crashed
    process does **not** prevent the next process from acquiring the lock.

    Platforms:
        Linux / macOS: ``fcntl.flock(LOCK_EX | LOCK_NB)``
        Windows:       ``msvcrt.locking(LK_NBLCK, …)``
    """

    def __init__(self, port: str) -> None:
        self._port = port
        self._lock_fd = None
        safe_name = os.path.basename(port).replace("\\", "_").replace("/", "_")
        self._lock_path = os.path.join(tempfile.gettempdir(), f"neng_lock_{safe_name}")

    # ------------------------------------------------------------------
    # Public interface
    # ------------------------------------------------------------------

    def acquire(self) -> None:
        """Acquire exclusive access.

        Raises:
            DeviceInUseError: if another process already holds the lock.
        """
        if platform.system() == "Windows":
            self._acquire_windows()
        else:
            self._acquire_unix()

    def release(self) -> None:
        """Release the lock and remove the lock file."""
        if self._lock_fd is None:
            return
        try:
            if platform.system() == "Windows":
                import msvcrt

                self._lock_fd.seek(0)
                msvcrt.locking(self._lock_fd.fileno(), msvcrt.LK_UNLCK, 1)
            else:
                import fcntl

                fcntl.flock(self._lock_fd, fcntl.LOCK_UN)
            self._lock_fd.close()
        except Exception:
            pass
        finally:
            self._lock_fd = None
        with contextlib.suppress(Exception):
            os.remove(self._lock_path)

    @property
    def lock_path(self) -> str:
        return self._lock_path

    def __enter__(self) -> "PortLock":
        self.acquire()
        return self

    def __exit__(self, *_) -> None:
        self.release()

    # ------------------------------------------------------------------
    # Platform-specific internals
    # ------------------------------------------------------------------

    def _read_pid(self) -> str:
        """Return ' (PID <n>)' if a PID can be read from the lock file."""
        try:
            with open(self._lock_path) as fh:
                content = fh.read().split()[0]
            if content.isdigit():
                return f" (PID {content})"
        except Exception:
            pass
        return ""

    def _acquire_unix(self) -> None:
        import fcntl

        fd = None
        try:
            # 'a+' opens for read+append and does NOT truncate the existing
            # file, so we can still read the PID written by the owner if the
            # lock attempt fails.
            fd = open(self._lock_path, "a+")  # noqa: SIM115
            fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
            # We now own the lock — overwrite with our PID.
            fd.seek(0)
            fd.truncate()
            fd.write(str(os.getpid()))
            fd.flush()
            self._lock_fd = fd
        except OSError as exc:
            if fd is not None:
                with contextlib.suppress(Exception):
                    fd.close()
                # File opened but lock refused → port in use by another process.
                pid_info = self._read_pid()
                raise DeviceInUseError(self._port, self._lock_path, pid_info) from exc
            # File could not even be opened (e.g. permission denied on /tmp).
            raise RuntimeError(
                f"Cannot create lock file for port {self._port}: {exc}\n"
                f"Attempted path: {self._lock_path}"
            ) from exc

    def _acquire_windows(self) -> None:
        import msvcrt

        fd = None
        try:
            # O_RDWR | O_CREAT: open for read/write, create if absent.
            raw = os.open(self._lock_path, os.O_RDWR | os.O_CREAT, 0o644)
            fd = os.fdopen(raw, "r+")
            fd.seek(0)
            # LK_NBLCK raises OSError immediately if the range is already locked.
            msvcrt.locking(fd.fileno(), msvcrt.LK_NBLCK, 1)
            # We own the lock — write our PID (fixed width so seek(0) always
            # covers the same locked byte on subsequent acquisitions).
            fd.seek(0)
            fd.write(f"{os.getpid():<10}")
            fd.flush()
            self._lock_fd = fd
        except OSError as exc:
            if fd is not None:
                with contextlib.suppress(Exception):
                    fd.close()
                pid_info = self._read_pid()
                raise DeviceInUseError(self._port, self._lock_path, pid_info) from exc
            raise RuntimeError(
                f"Cannot create lock file for port {self._port}: {exc}\n"
                f"Attempted path: {self._lock_path}"
            ) from exc


class SCPISerial:
    """
    Easy-to-use serial interface for SCPI instruments.

    Features:
    - Automatic port detection
    - Context manager support (with statement)
    - Simple write/query methods
    - Built-in error handling
    - Connection verification
    - Cross-platform support (macOS, Linux, Windows)
    """

    def __init__(self, port=None, baudrate=115200, timeout=1, auto_connect=True, exclusive=True):
        """
        Initialize SCPI serial interface.

        Args:
            port: Serial port path (e.g., "/dev/cu.usbmodem2101", "COM3")
                  If None, will attempt auto-detection
            baudrate: Serial baudrate (default: 115200)
            timeout: Read timeout in seconds (default: 1)
            auto_connect: Automatically connect on initialization (default: True)
            exclusive: Enforce single-client access via a lock file (default: True).
                       Raises DeviceInUseError if another process already holds the port.
        """
        self.port = port
        self.baudrate = baudrate
        self.timeout = timeout
        self.serial = None
        self._connected = False
        self._rx_partial = ""  # buffer for incomplete lines (non-blocking reads)
        self._exclusive = exclusive
        self._port_lock: PortLock | None = None

        if auto_connect:
            self.connect()

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

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit - closes connection"""
        self.close()
        return False

    @staticmethod
    def list_ports():
        """
        List all available serial ports with descriptions.

        Returns:
            List of tuples: [(port, description, hwid), ...]
        """
        ports = []
        for port in serial.tools.list_ports.comports():
            ports.append((port.device, port.description, port.hwid))
        return ports

    @staticmethod
    def find_neng_port():
        """
        Attempt to auto-detect NEnG instrument port.

        Looks for common USB CDC patterns:
        - macOS: /dev/cu.usbmodem*
        - Linux: /dev/ttyACM*, /dev/ttyUSB*
        - Windows: COM*

        Returns:
            Port path string or None if not found
        """
        system = platform.system()

        for port in serial.tools.list_ports.comports():
            device = port.device

            # Check for common CircuitPython/USB CDC patterns
            if system == "Darwin":  # macOS
                if device.startswith("/dev/cu.usbmodem"):
                    return device
            elif system == "Linux":
                # Check for USB CDC and USB serial devices
                if device.startswith("/dev/ttyACM") or device.startswith("/dev/ttyUSB"):
                    return device
            elif system == "Windows" and device.startswith("COM"):
                return device

        return None

    @staticmethod
    def get_useful_ports():
        """
        Get list of potentially useful serial ports (USB, not system ports).

        Filters out system tty ports (/dev/ttyS*) which are usually not USB devices.

        Returns:
            List of tuples: [(device, description), ...]
        """
        system = platform.system()
        useful_ports = []

        for port in serial.tools.list_ports.comports():
            device = port.device
            desc = port.description

            # Filter based on platform
            if system == "Darwin":  # macOS
                # Include USB modem and serial ports
                if device.startswith("/dev/cu.usbmodem") or device.startswith("/dev/cu.usbserial"):
                    useful_ports.append((device, desc))
            elif system == "Linux":
                # Include USB CDC and USB serial, exclude system ttyS ports
                if device.startswith("/dev/ttyACM") or device.startswith("/dev/ttyUSB"):
                    useful_ports.append((device, desc))
            elif system == "Windows" and device.startswith("COM"):
                # Include COM ports (usually USB on Windows)
                useful_ports.append((device, desc))

        return useful_ports

    def connect(self):
        """
        Connect to serial port.

        If port was not specified during initialization, attempts auto-detection.

        Raises:
            serial.SerialException: If connection fails
            ValueError: If no port specified and auto-detection fails
        """
        if self._connected:
            return

        # Auto-detect port if not specified
        if self.port is None:
            self.port = self.find_neng_port()
            if self.port is None:
                available = self.get_useful_ports()
                error_msg = (
                    "❌ No compatible device found!\n"
                    "Possible causes:\n"
                    "  • Device is disconnected\n"
                    "  • Device is powered off\n"
                    "  • USB cable is loose or damaged\n"
                    "  • Device is in use by another application\n\n"
                    "Solution:\n"
                    "  1. Check the USB connection\n"
                    "  2. Check the device is powered on\n"
                    "  3. Try specifying the port explicitly\n"
                )
                if available:
                    error_msg += "\nAvailable USB ports:\n"
                    for port_name, port_desc in available:
                        error_msg += f"  • {port_name}: {port_desc}\n"
                else:
                    error_msg += "\nNo USB serial ports found on this system.\n"

                raise ValueError(error_msg)

        # Acquire exclusive port lock before opening the serial port.
        if self._exclusive:
            self._port_lock = PortLock(self.port)
            self._port_lock.acquire()  # raises DeviceInUseError if already locked

        # Connect
        try:
            self.serial = serial.Serial(self.port, baudrate=self.baudrate, timeout=self.timeout)
            time.sleep(0.1)  # Allow connection to stabilize
            self._connected = True

            # Verify connection with *IDN? query
            try:
                idn = self.query("*IDN?")
                if idn:
                    pass  # Connection verified
            except Exception:
                pass  # Continue even if *IDN? fails

        except serial.SerialException as e:
            if self._port_lock:
                self._port_lock.release()
                self._port_lock = None
            raise serial.SerialException(f"Failed to connect to {self.port}: {e}") from e

    def close(self):
        """Close serial connection and release the exclusive port lock."""
        if self.serial and self._connected:
            try:
                self.serial.close()
            except Exception:
                pass
            finally:
                self._connected = False
                self.serial = None
        if self._port_lock:
            self._port_lock.release()
            self._port_lock = None

    def write(self, command):
        """
        Send a command to the instrument.

        Args:
            command: SCPI command string (newline added automatically)

        Raises:
            RuntimeError: If not connected
        """
        if not self._connected or not self.serial:
            raise RuntimeError("Not connected to instrument")

        # Ensure command ends with newline
        if not command.endswith("\n"):
            command += "\n"

        self.serial.write(command.encode("ascii"))
        self.serial.flush()

    def read(self):
        """
        Read a line from the instrument.

        Returns:
            Response string (stripped of whitespace)

        Raises:
            RuntimeError: If not connected
        """
        if not self._connected or not self.serial:
            raise RuntimeError("Not connected to instrument")

        response = self.serial.readline().decode("ascii", errors="ignore").strip()
        return response

    def bytes_available(self):
        """Return the number of bytes available in the RX buffer (non-blocking).

        Notes:
            - Based on `pyserial.Serial.in_waiting`.
            - 0 means no data is currently waiting.
        """
        if not self._connected or not self.serial:
            raise RuntimeError("Not connected to instrument")
        try:
            return int(getattr(self.serial, "in_waiting", 0) or 0)
        except Exception:
            return 0

    def has_data(self):
        """Return True if data is available to read (non-blocking)."""
        return self.bytes_available() > 0

    def read_available(self, max_bytes=None, decode=True):
        """Read everything that's already available (without waiting).

        Args:
            max_bytes: read limit. If None, read everything currently available.
            decode: if True, returns an ASCII str (errors='ignore'), otherwise bytes.

        Returns:
            str (or bytes); empty if nothing is available.
        """
        if not self._connected or not self.serial:
            raise RuntimeError("Not connected to instrument")

        available = self.bytes_available()
        if available <= 0:
            return "" if decode else b""

        if max_bytes is not None:
            with contextlib.suppress(Exception):
                available = min(int(max_bytes), available)

        # Force a non-blocking read even if the instance has timeout > 0
        old_timeout = getattr(self.serial, "timeout", None)
        try:
            self.serial.timeout = 0
            raw = self.serial.read(available) or b""
        finally:
            with contextlib.suppress(Exception):
                self.serial.timeout = old_timeout

        if not decode:
            return raw
        return raw.decode("ascii", errors="ignore")

    def read_available_lines(self, max_bytes=None):
        """Read all available *complete lines* without blocking.

        Useful when the device sends multiple lines (or logs) and you want
        to drain output without risking a block.

        Splitting is done on '\n' (supports '\r\n'). Incomplete line fragments
        are kept internally and reassembled on the next call.

        Args:
            max_bytes: optional read limit, in bytes.

        Returns:
            List[str]: lines (without '\r'/'\n').
        """
        chunk = self.read_available(max_bytes=max_bytes, decode=True)
        if not chunk:
            return []

        data = self._rx_partial + chunk

        # Normalize CRLF -> LF then split on LF
        data = data.replace("\r\n", "\n").replace("\r", "\n")

        if "\n" not in data:
            # Still an incomplete line
            self._rx_partial = data
            return []

        parts = data.split("\n")
        self._rx_partial = parts[-1]  # last fragment may be incomplete
        lines = [p for p in parts[:-1] if p != ""]
        return lines

    def read_binary_until(self, terminator=b"\n"):
        """
        Read binary data until a terminator is found.

        Args:
            terminator: Binary terminator to stop reading (default: b"\n")

        Returns:
            Binary data including terminator

        Raises:
            RuntimeError: If not connected
        """
        if not self._connected or not self.serial:
            raise RuntimeError("Not connected to instrument")

        return self.serial.read_until(terminator)

    def read_ieee_block(self, timeout=1.0):
        """
        Read IEEE 488.2 definite-length block data.

        Format: #<digit><length><data><terminator>
        Example: #120<20 bytes of binary data>\n

        Args:
            timeout: Read timeout in seconds

        Returns:
            Binary data (without the header and terminator)

        Raises:
            RuntimeError: If not connected or timeout occurs
            ValueError: If block format is invalid
        """
        if not self._connected or not self.serial:
            raise RuntimeError("Not connected to instrument")

        # Save old timeout
        old_timeout = self.serial.timeout
        self.serial.timeout = timeout

        try:
            # Read until we get the header marker
            header = self.serial.read_until(b"#")
            if not header or header[-1:] != b"#":
                raise ValueError("Invalid IEEE block: missing '#' marker")

            # Read the digit count
            digit_count_byte = self.serial.read(1)
            if not digit_count_byte:
                raise ValueError("Invalid IEEE block: missing digit count")

            digit_count = int(chr(digit_count_byte[0]))
            if digit_count <= 0 or digit_count > 9:
                raise ValueError(f"Invalid IEEE block: digit count {digit_count}")

            # Read the length value
            length_str = self.serial.read(digit_count)
            if not length_str:
                raise ValueError("Invalid IEEE block: missing length")

            data_length = int(length_str)

            # Read the binary data
            data = self.serial.read(data_length)
            if len(data) != data_length:
                raise ValueError(f"Expected {data_length} bytes, got {len(data)}")

            # Read the terminator (usually \n)
            self.serial.read(1)

            return data

        finally:
            self.serial.timeout = old_timeout

    def query(self, command, delay=0.1):
        """
        Send a query and read the response.

        Drains any stale data in the serial input buffer before sending
        the command, ensuring the response read corresponds to this query.

        Args:
            command: SCPI query string (e.g., "*IDN?")
            delay: Delay between write and read in seconds (default: 0.1)

        Returns:
            Response string

        Raises:
            RuntimeError: If not connected
        """
        if not self._connected or not self.serial:
            raise RuntimeError("Not connected to instrument")

        # Drain stale data from the input buffer before sending
        self.serial.reset_input_buffer()
        self._rx_partial = ""

        self.write(command)
        time.sleep(delay)
        return self.read()

    def query_int(self, command, delay=0.1):
        """
        Query and return integer value.

        Args:
            command: SCPI query string
            delay: Delay between write and read in seconds

        Returns:
            Integer value or None if conversion fails
        """
        response = self.query(command, delay)
        return self.parse_int(response)

    @staticmethod
    def parse_int(value, base=None):
        """Parse a value (str/int) into an int.

            Supports:
        - decimal: "123"
        - hex with prefix: "0xFFFF"
        - hex without prefix (typical for masks): "FFFF"
        - binary with prefix: "0b1010"

            Args:
                value: raw value (often the SCPI response)
                base: explicit base (2/10/16). If None, auto-detect.

            Returns:
                int, or None if conversion fails.
        """
        if value is None:
            return None

        if isinstance(value, int):
            return value

        s = str(value).strip()
        if s == "":
            return None

        try:
            if base is not None:
                return int(s, base)

            # auto-detect
            low = s.lower()
            if low.startswith("0x"):
                return int(low, 16)
            if low.startswith("0b"):
                return int(low, 2)

            # If the string looks like hex without a prefix (A-F)
            hex_chars = set("0123456789abcdef")
            if all(c in hex_chars for c in low) and any(c in "abcdef" for c in low):
                return int(low, 16)

            return int(s, 10)
        except (ValueError, TypeError):
            return None

    def query_hex_int(self, command, delay=0.1):
        """Query and explicitly parse a hexadecimal value (with or without 0x)."""
        response = self.query(command, delay)
        return self.parse_int(response, base=16)

    def query_mask_bits(self, command, width=16, delay=0.1):
        """Query a bitmask (often returned in hex) and return its binary form.

        Example: "FFFF" with width=16 -> "1111111111111111"
        """
        mask = self.query_int(command, delay=delay)
        if mask is None:
            return None
        return format(mask, f"0{width}b")

    def query_float(self, command, delay=0.1):
        """
        Query and return float value.

        Args:
            command: SCPI query string
            delay: Delay between write and read in seconds

        Returns:
            Float value or None if conversion fails
        """
        response = self.query(command, delay)
        try:
            return float(response)
        except (ValueError, TypeError):
            return None

    def query_bool(self, command, delay=0.1):
        """
        Query and return boolean value.

        Args:
            command: SCPI query string
            delay: Delay between write and read in seconds

        Returns:
            True if response is "1", False otherwise
        """
        response = self.query(command, delay)
        return response == "1"

    def wait_opc(self, timeout=10):
        """
        Wait for operation complete (OPC).

        Polls *OPC? until it returns "1" or timeout occurs.

        Args:
            timeout: Maximum wait time in seconds

        Returns:
            True if OPC received, False if timeout
        """
        start_time = time.time()
        while time.time() - start_time < timeout:
            try:
                if self.query_bool("*OPC?", delay=0.05):
                    return True
            except Exception:
                pass
            time.sleep(0.1)
        return False

    @property
    def connected(self):
        """Check if connected to instrument"""
        return self._connected

    @property
    def port_name(self):
        """Get connected port name"""
        return self.port if self.port else "Not connected"

    def get_identification(self):
        """
        Get instrument identification.

        Returns:
            Dictionary with manufacturer, model, serial, firmware
            or None if query fails
        """
        try:
            idn = self.query("*IDN?")
            parts = idn.split(",")
            if len(parts) >= 4:
                return {
                    "manufacturer": parts[0].strip(),
                    "model": parts[1].strip(),
                    "serial": parts[2].strip(),
                    "firmware": parts[3].strip(),
                    "full": idn,
                }
        except Exception:
            pass
        return None

    def __repr__(self):
        """String representation"""
        status = "connected" if self._connected else "disconnected"
        return f"SCPISerial(port='{self.port}', {status})"


# Convenience function for quick connections
def connect(port=None, **kwargs):
    """
    Convenience function to create and connect to instrument.

    Args:
        port: Serial port (optional, will auto-detect if None)
        **kwargs: Additional arguments passed to SCPISerial()

    Returns:
        Connected SCPISerial instance

    Example:
        instr = connect()
        print(instr.query("*IDN?"))
        instr.close()
    """
    return SCPISerial(port=port, **kwargs)


# Example usage and testing
if __name__ == "__main__":
    print("SCPI Serial Interface - Test")
    print("=" * 50)

    # List available ports
    print("\nAvailable serial ports:")
    ports = SCPISerial.list_ports()
    if ports:
        for port, desc, hwid in ports:
            print(f"  {port}")
            print(f"    Description: {desc}")
            print(f"    Hardware ID: {hwid}")
    else:
        print("  No serial ports found")

    # Try auto-detection
    print("\nAttempting auto-detection...")
    detected = SCPISerial.find_neng_port()
    if detected:
        print(f"  Detected port: {detected}")
    else:
        print("  No instrument detected")
        sys.exit(1)

    # Connect and test
    print("\nConnecting to instrument...")
    try:
        with SCPISerial() as instr:
            print(f"  Connected to: {instr.port_name}")

            # Get identification
            idn = instr.get_identification()
            if idn:
                print("\nInstrument Information:")
                print(f"  Manufacturer: {idn['manufacturer']}")
                print(f"  Model: {idn['model']}")
                print(f"  Serial: {idn['serial']}")
                print(f"  Firmware: {idn['firmware']}")

            # Test some queries
            print("\nTesting queries...")
            print(f"  *OPC? = {instr.query('*OPC?')}")
            print(f"  *ESR? = {instr.query('*ESR?')}")

            print("\n✓ Connection test successful!")

    except Exception as e:
        print(f"\n✗ Connection test failed: {e}")
        sys.exit(1)
