""" Autogenerated module for the {{ name }} QCodes driver. """
from zhinst.qcodes.base import ZISnapshotHelper, ZIBaseInstrument
from typing import Any, Dict, Union, List
from qcodes.instrument.base import Instrument
from zhinst.qcodes.qcodes_adaptions import ZIParameter
{% if base_module == "Instrument"%}
{% else %}
from zhinst.qcodes.modules.base_module import BaseInstrument
{% endif %}
from zhinst.toolkit.driver.modules.base_module import BaseModule as TKBaseModule
from zhinst.toolkit.driver.modules.{{ module_name }} import {{name}} as TK{{name}}

class {{ name }}({{ base_module }}):
    """QCodes driver for the Zurich Instrument {{ original_name }} module."""

{% if base_module == "Instrument"%}
    def __init__(self, tk_object: TK{{ name }}, name):
        super().__init__(f"zi_{{ name.lower() }}_{len(self.instances())}")
        self._tk_object = tk_object
        self._snapshot_cache = ZISnapshotHelper(self._tk_object, is_module = True)
        ZIBaseInstrument.init_nodetree(self, self._tk_object, self._snapshot_cache)

    def snapshot(self, update: bool = True) -> dict:
        """Decorate a snapshot dictionary with metadata.

        Override base method to make update default True and use the
        ZISnapshotHelper.

        Args:
            update: Passed to snapshot_base.

        Returns:
            dict: Base snapshot.
        """
        if not self._snapshot_cache.is_running:
            self._snapshot_cache.start_snapshot()
            val = super().snapshot(update)
            self._snapshot_cache.stop_snapshot()
        else:
            val = super().snapshot(update)
        return val

    def print_readable_snapshot(self, update: bool = True, max_chars: int = 80) -> None:
        """
        Prints a readable version of the snapshot.
        The readable snapshot includes the name, value and unit of each
        parameter.
        A convenience function to quickly get an overview of the
        status of an instrument.

        Args:
            update: If ``True``, update the state by querying the
                instrument. If ``False``, just use the latest values in memory.
                This argument gets passed to the snapshot function.
            max_chars: the maximum number of characters per line. The
                readable snapshot will be cropped if this value is exceeded.
                Defaults to 80 to be consistent with default terminal width.
        """
        self._snapshot_cache.start_snapshot()
        val = ZISnapshotHelper.print_readable_snapshot(
            self, update=update, max_chars=max_chars
        )
        self._snapshot_cache.stop_snapshot()
        return val

    @property
    def raw_module(self) -> Any:
        return self._tk_object.raw_module

    def subscribe(self, signal: ZIParameter):
        """Subscribe to a node.

        The node can either be a node of this module or of a connected device.

        Args:
            signal (Node): node that should be subcribed.
        """
        self._raw_module.subscribe(signal.zi_node)

    def unsubscribe(self, signal: ZIParameter):
        """Unsubscribe from a node.

        The node can either be a node of this module or of a connected device.

        Args:
            signal (Node): node that should be unsubscribe.
        """
        self._raw_module.unsubscribe(signal.zi_node)
{% else %}
    def __init__(self, tk_object: TK{{ name }}):
        super().__init__(tk_object, "{{ name.lower() }}")
{% endif %}

{% for function in functions %}
    def {{ function.name }}{{ function.signature }}:
        """{{ function.docstring }}"""
        {% if function.node_function %}
        data_tk = self._tk_object.{{ function.name }}({{ function.call_signature }})
        data = {}
        for tk_node, data in data_tk.items():
            device = self._session.devices[tk_node.root.prefix_hide]
            parameter = tk_node_to_parameter(device, tk_node)
            data[parameter] = data
        return data
        {%else%}
        return self._tk_object.{{ function.name }}({{ function.call_signature }})
        {% endif%}
{% endfor %}
