import serial
import re

from lab_control.function_generator import FunctionGenerator

class JDS6600(FunctionGenerator):
    AVAILABLE_CHANNELS = [1, 2]

    def __init__(self, portName):
        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 _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 closePort(self) -> None:
        self._port.close()

    def _queryOnOff(self) -> list[str, str]:
        # response format: ":r20=0,0."
        response = self._sendRequest("r20")
        return [response[5], response[7]]

    def setOn(self, channel: int) -> None:
        assert channel in JDS6600.AVAILABLE_CHANNELS, f"JDS6600: Invalid channel {channel}"

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

    def setOff(self, channel: int) -> None:
        assert channel in JDS6600.AVAILABLE_CHANNELS, f"JDS6600: Invalid channel {channel}"

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

    def setFrequency(self, channel: int, frequency: float) -> None:
        assert channel in JDS6600.AVAILABLE_CHANNELS, f"JDS6600: Invalid channel {channel}"
        assert frequency is not None and frequency >= 0.0 and frequency <= 60e6, f"JDS6600: Invalid frequency {frequency}"

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

    def setAmplitude(self, channel: int, amplitude: float) -> None:
        assert channel in JDS6600.AVAILABLE_CHANNELS, f"JDS6600: Invalid channel {channel}"
        assert amplitude is not None and amplitude >= 0.0 and amplitude <= 20.0, f"JDS6600: Invalid amplitude {amplitude}"

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

    def setFunction(self, channel: int, function: int) -> None:
        assert channel in JDS6600.AVAILABLE_CHANNELS, f"JDS6600: Invalid channel {channel}"
        assert function is not None and function >= 0 and function <= 16, f"JDS6600: Invalid function code {function}"

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