from collections import defaultdict
from dataclasses import dataclass
from pyvisa_mock.base.base_mocker import BaseMocker, scpi
from pyvisa_mock.base.register import register_resource


class BK9132_mocker_channel(BaseMocker):
    """
    A mocker channel for a multi channel voltage source.
    Voltages are zero by default
    """

    def __init__(self, call_delay: float = 0.0):
        super().__init__(call_delay=call_delay)
        self._voltage = 0

    # Lets define handler functions.

    @scpi(":VOLTage <voltage>")
    def _set_voltage(self, voltage: float) -> None:
        self._voltage = voltage

    @scpi(":MEAS:VOLT:ALL?")
    def _get_voltage(self) -> float:
        return self._voltage

@dataclass
class Channel:
    voltage: float = 0
    current: float = 0
    output: float = 0
    
class BK9132_mocker(BaseMocker):
    channels = [Channel() for _ in range(3)] 
    
    def __init__(self, call_delay: float = 0.0):
        super().__init__(call_delay=call_delay)
        self._channels = defaultdict(BK9132_mocker_channel)

    @scpi("*IDN?")
    def idn(self) -> str:
        """
        'vendor', 'model', 'serial', 'firmware'
        """
        return "Mocker,testing,00000,0.01"

    @scpi(":MEAS:VOLT:ALL?")
    def _get_voltages(self) -> str:
        voltages = [channel.voltage for channel in self.channels]
        return  ','.join(map(str, voltages)) 


register_resource("MOCK0::mock1::INSTR", BK9132_mocker())