"""
Implements partial support for Joy-IT JDS6600 function generator.
"""

import serial

from lab_control.function_generator import FunctionGenerator

def _checkChannel(channel: int):
    assert channel in JDS6600.AVAILABLE_CHANNELS, f"JDS6600: Invalid channel {channel}"

def _checkArg(arg: float, lowerBound: float, upperBound: float):
    valid = arg is not None
    if valid:
        valid &= arg >= lowerBound
        valid &= arg <= upperBound
    assert valid, f"JDS6600: Invalid argument {arg}"

class JDS6600(FunctionGenerator):
    """
    Instances of this class bind to the JDS6600 serial port
    and offer an API to control the device.
    """
    AVAILABLE_CHANNELS = [1, 2]

    def __init__(self, portName):
        super().__init__()
        self._port = serial.Serial(portName)
        self._port.baudrate = 115200
        self._port.bytesize = serial.EIGHTBITS
        self._port.stopbits = serial.STOPBITS_ONE
        self._port.parity = serial.PARITY_NONE

    def closePort(self) -> None:
        """
        Close the serial port. Instances of this class
        are no longer usable after this is called.
        """
        self._port.close()

    def _query(self, opcode: str) -> list[str]:
        # response format: ":{opcode}={v1},{v2}."
        response = self._sendRequest(opcode)
        return response[5:-1].split(",")

    def _sendRequest(self, opcode: str, args: str="") -> str:
        request = f":{opcode}={args}.\r\n"
        self._port.write(request.encode())
        responseRaw = self._port.readline()
        return responseRaw.decode().strip()

    def setOn(self, channel: int) -> None:
        _checkChannel(channel)

        state = self._query("r20")
        state[channel - 1] = "1"
        self._sendRequest("w20", f"{state[0]},{state[1]}")

    def setOff(self, channel: int) -> None:
        _checkChannel(channel)

        state = self._query("r20")
        state[channel - 1] = "0"
        self._sendRequest("w20", f"{state[0]},{state[1]}")

    def setFrequency(self, channel: int, frequency: float) -> None:
        _checkChannel(channel)
        _checkArg(frequency, 0.0, 60e6)

        opcode = 23 + channel - 1
        arg = int(frequency * 100.0)
        self._sendRequest(f"w{opcode}", f"{arg},0")

        # JDS6600 has been observed to sometimes incorrectly
        # ignore the frequency and set to 0 instead
        actual = self._query(f"r{opcode}")[0]
        if actual != arg:
            self._sendRequest(f"w{opcode}", f"{arg},0")

    def setAmplitude(self, channel: int, amplitude: float) -> None:
        _checkChannel(channel)
        _checkArg(amplitude, 0.0, 20.0)

        opcode = f"w{25 + channel - 1}"
        arg = int(amplitude * 1000.0)
        self._sendRequest(opcode, str(arg))

    def setFunction(self, channel: int, function: int) -> None:
        _checkChannel(channel)
        _checkArg(function, 0, 16)

        opcode = f"w{21 + channel - 1}"
        self._sendRequest(opcode, str(function))