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

from lab_control.function_generator import FunctionGenerator
from lab_control.connection.serial_connection import SerialConnection

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

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

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]

    # TODO type hints
    def __init__(self, portName: str, overrideConnection=None):
        super().__init__()

        config = {
            "baudrate" : 115200,
            "bytesize" : 8,
            "stopbits" : 1,
            "parity" : "N"
        }

        if overrideConnection is not None:
            self._connection = overrideConnection
        else:
            self._connection = SerialConnection(portName)

        self._connection.configure(config)

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

    def _sendRequest(self, opcode: str, args: str="") -> str:
        request = f":{opcode}={args}.\r\n"
        response = self._connection.send(request)
        return response.strip()

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

    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)
        _checkBounds(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)
        _checkBounds(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)
        _checkBounds(function, 0, 16)

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