"""Autogenerated module for the {{ name }} QCoDeS driver."""
import typing as t
from typing import Union, Optional,List,Dict, Any, Tuple
from pathlib import Path
from zhinst.qcodes.driver.devices.base import ZIBaseInstrument
from zhinst.toolkit.driver.modules import ModuleType as TKModuleType
from zhinst.toolkit.driver.modules.{{ module_name }} import {{name}} as TK{{name}}
from zhinst.toolkit.driver.modules.base_module import ZIModule
from qcodes.instrument.base import Instrument
{% if base_module != "ZIInstrument"%}
from zhinst.qcodes.driver.modules.base_module import ZIBaseModule
{% endif -%}
from zhinst.toolkit.nodetree import Node as TKNode

from zhinst.qcodes.qcodes_adaptions import ZIParameter, NodeDict, ZIInstrument, init_nodetree, tk_node_to_parameter

if t.TYPE_CHECKING:
    from zhinst.qcodes.driver.devices import DeviceType
    from zhinst.qcodes.session import Session

class ZI{{ name }}({{ base_module }}):
    """{{ module_docstring }}
    Args:
        tk_object: Underlying zhinst-toolkit object.
        session: Session to the Data Server.
        name: Name of the module in QCoDeS.
    """
{% if base_module == "ZIInstrument"%}
    def __init__(self, tk_object: TKModuleType, session: "Session", name :str):
        self._tk_object = tk_object
        self._session = session
        super().__init__(
            f"zi_{name}_{len(self.instances())}", tk_object.root, is_module=True
        )
        init_nodetree(self, self._tk_object, self._snapshot_cache)

        self._tk_object.root.update_nodes(
            {
                "/device": {
                    "GetParser": lambda value: self._get_device(value),
                }
            },
            raise_for_invalid_node=False, 
        )


    def _get_device(self, serial: str) -> t.Union["DeviceType", str]:
        """Convert a device serial into a QCoDeS device object.

        Args:
            serial: Serial of the device

        Returns:
            QCoDeS device object. If the serial does not
                match to a connected device the serial is returned instead.
        """
        try:
            return self._session.devices[serial]
        except (RuntimeError, KeyError):
            return serial

    def _get_node(self, node: str) -> t.Union[ZIParameter, str]:
        """Convert a raw node string into a qcodes node.

        Args:
            node (str): raw node string

        Returns:
            Node: qcodes node. (if the node can not be converted the raw node
                string is returned)
        """
        tk_node = self._tk_object._get_node(node)
        if isinstance(tk_node, str):
            return tk_node
        device = self._session.devices[tk_node.root.prefix_hide]
        return tk_node_to_parameter(device, tk_node)

    @staticmethod
    def _set_node(signal: t.Union[ZIParameter, TKNode, str]) -> str:
        """Convert a toolkit node into a raw node string.

        Args:
            signal (Union[Node,str]): node

        Returns:
            str: raw string node
        """
        try:
            node = signal.zi_node  # type: ignore[union-attr]
        except AttributeError:
            node = TKBaseModule._set_node(signal)
        return node

    def subscribe(self, signal: t.Union[ZIParameter, str]):
        """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 subscribed.
        """
        try:
            self._tk_object.subscribe(signal.zi_node)  # type: ignore[union-attr]
        except AttributeError:
            self._tk_object.subscribe(signal)

    def unsubscribe(self, signal: t.Union[ZIParameter, str]):
        """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.
        """
        try:
            self._tk_object.unsubscribe(signal.zi_node)  # type: ignore[union-attr]
        except AttributeError:
            self._tk_object.unsubscribe(signal)

    @property
    def raw_module(self) -> ZIModule:  # type: ignore [type-var]
        """Underlying zhinst.core module."""
        return self._tk_object.raw_module
{% else %}
    def __init__(self, tk_object: TK{{ name }}, session: "Session", name :str = "{{ module_name.lower() }}"):
        super().__init__(tk_object, session, name)
        {% for node in node_param %}
        self._tk_object.root.update_nodes(
            {
                "/{{ node }}": {
                    "GetParser": lambda value: self._get_node(value),
                    "SetParser": lambda value: self._set_node(value),
                }
            },
            raise_for_invalid_node=False,
        )
        {% endfor %}
{% endif %}

{% for function in functions %}
    def {{ function.name }}{{ function.signature }}:
        """{{ function.docstring }}"""
        {% if function.is_node_dict -%}
        return NodeDict(self._tk_object.{{ function.name }}({{ function.call_signature }}))
        {%else-%}
        return self._tk_object.{{ function.name }}({{ function.call_signature }})
        {% endif-%}
{% endfor %}
