"""Module defining Real Node types, which can be run using our DeltaPySimulator."""
from __future__ import annotations
from inspect import _empty, signature
import logging
from queue import Full, Queue
import sys
import textwrap
from threading import Event
from time import sleep

from typing import (TYPE_CHECKING,
                    Any,
                    Callable,
                    Dict,
                    List,
                    Optional,
                    OrderedDict,
                    Tuple,
                    Type,
                    Union)

import dill

from deltalanguage.data_types import (BaseDeltaType,
                                      DeltaTypeError,
                                      DOptional,
                                      ForkedReturn,
                                      Void,
                                      DeltaIOError,
                                      as_delta_type,
                                      delta_type)
from deltalanguage.logging import MessageLog, make_logger
from deltalanguage._utils import (NamespacedName,
                                  QueueMessage)

from .._body_templates import BodyTemplate, MethodBodyTemplate
from .abstract_node import AbstractNode, ForkedNode
from .node_bodies import (PyConstBody,
                          PyFuncBody,
                          PyMethodBody,
                          PyInteractiveBody,
                          PyMigenBody,
                          Body)
from .port_classes import InPort, OutPort

if TYPE_CHECKING:
    from deltalanguage.runtime import DeltaQueue, DeltaPySimulator
    from .. import DeltaGraph


class RealNode(AbstractNode):
    """Class to represent a non-abstract node that can form part of
    :py:class:`DeltaGraph`.

    Parameters
    ----------
    graph : DeltaGraph
        Graph this node is a member of.
    bodies : List[Body]
        List of bodies that can represent the implementation for this node.
    out_type : BaseDeltaType
        The type of what we expect the body to output.
    name : Optional[str]
        The name of the node (an index is appended to the end).
    lvl : int
        Logging level for the node.
        By default only error logs are displayed.
    is_autogenerated : bool
        True if this node is created automatically to provide an input to
        another node. For instance:

        .. code-block:: python

            with dl.DeltaGraph() as graph:
                printer(42)

        has an autogenerated node that provides 42. For strict typing this
        node should send data in the same format as printer's input.
    """

    def __init__(self,
                 graph,
                 bodies: List[Body],
                 in_params: Dict[str, BaseDeltaType] = {},
                 out_type: BaseDeltaType = None,
                 name: Optional[str] = None,
                 lvl: int = logging.ERROR,
                 is_autogenerated: bool = False):
        self.graph = graph
        self.graph.add_node(self)  # Registering self with parent graph

        self.is_autogenerated = is_autogenerated

        self._body = None
        self.bodies = bodies

        if len(self.bodies) == 1:
            self._body = bodies[0]

        self.in_params = in_params
        self.out_type = out_type

        idx = RealNode.get_next_index()
        if name is None:
            # set my name to the next unique available name
            self._name = f"node_{idx}"
        else:
            self._name = f"{name}_{idx}"

        # Ports in/out to this node
        self.in_ports: Dict[NamespacedName, InPort] = {}
        self.out_ports: List[OutPort] = []

        self.log = make_logger(lvl,
                               f"{self.__class__.__name__} {self._name}")

        # See MessageLog for detail
        self._clock = 0

        self.fork_names = None
        if isinstance(self.out_type, ForkedReturn):
            self.fork_names = self.out_type.keys

            for fork_name in self.fork_names:
                if fork_name in dir(self):
                    raise NameError("Invalid fork name: " +
                                    fork_name + " for node " + self.name)

    def select_body(self,
                    exclusions: List[object] = [],
                    preferred: List[object] = [],
                    override=True):
        """Select a body for this node, out of the existing list of bodies
        attached to it.
        Exclusions take priority over preferences. If multiple choices
        remain the earliest added node will be selected.

        Parameters
        ----------
        exclusions : List[object]
            List of keys to exclude from selection
        preferred : List[object]
            List of keys to be preferred for selection if available
        override : Bool
            If true, then this selection will override existing selection
            Otherwise this selection will not occur if there has already
            been a selection

        Returns
        -------
        bool
            ``True`` if a ``body`` is defined at the end of selection.

        """

        if override or not self._body:

            not_excluded_list = []
            for body in self.bodies:
                # If body has no excluded tag
                if not [tag for tag in exclusions if tag in body.access_tags]:
                    not_excluded_list.append(body)

            if not_excluded_list:
                prefered_list = []
                for body in not_excluded_list:
                    # If body has a preferred tag
                    if [tag for tag in preferred if tag in body.access_tags]:
                        prefered_list.append(body)

                if prefered_list:
                    self._body = prefered_list[0]
                else:
                    self._body = not_excluded_list[0]
            else:
                self._body = None
        return bool(self._body)

    def __str__(self) -> str:

        indent_prefix = "    "

        ret = f'ports:\n'
        if self.in_ports:
            ret += textwrap.indent(
                f'in:\n' + textwrap.indent(
                    ''.join(
                        [f"{in_port}\n" for in_port in self.in_ports.values()]
                    ),
                    prefix=indent_prefix),
                prefix=indent_prefix)
        if self.out_ports:
            ret += textwrap.indent(
                f'out:\n' + textwrap.indent(
                    ''.join(
                        [f"{out_port}\n" for out_port in self.out_ports]
                    ),
                    prefix=indent_prefix),
                prefix=indent_prefix)

        ret += f'bodies:\n'
        if not self.bodies:
            ret += textwrap.indent('None\n\n', prefix=indent_prefix)
        else:
            ret += textwrap.indent(
                ''.join(
                    [
                        f'{"*" if body is self._body else ""}' +
                        f'{body.__class__.__name__}\n' +
                        textwrap.indent(
                            f'tags: {next(i for i in body.access_tags if i is not body.__class__)}\n',
                            prefix=indent_prefix
                        ) for body in self.bodies
                    ]
                ),
                prefix=indent_prefix
            ) + '\n'

        ret = f"node[{self.name}]:\n" + textwrap.indent(
            ret, prefix=indent_prefix
        )

        return ret

    def __repr__(self):
        return self.name

    def __getattr__(self, item):
        if self.fork_names and item in self.fork_names:
            return ForkedNode(self, item)
        elif self.fork_names is None:
            raise AttributeError(
                f"Cannot fetch {item} from {self.name}, as we don't "
                f"have multiple outputs. Suggest using the node on it's own."
            )
        elif item == '__del__':
            raise AttributeError('Migen tries to get this attribute')
        else:
            raise DeltaIOError(
                f"Node {self.name} has out ports {list(self.fork_names)}. "
                f"The requested out port \'{item}\' is not found. "
                "Please check the node's definition."
            )

    def __getstate__(self):
        return self.__dict__

    def __setstate__(self, state):
        self.__dict__ = state

    def add_out_port(self, port_destination: InPort, index=None):
        """Creates an out-port and adds it to my out-port store.

        Parameters
        ----------
        index : Optional[str]
            If the out-port is one of several for this node, index specifies
            what part of the output is sent via this port. If the node has
            only one output, then this is `None`.
        port_destination : InPort
            The in-port that this out-port exports to.
        """
        try:
            if issubclass(self.out_type, Void):
                raise DeltaIOError(
                    f"Cannot make an out-port on node {self.name} "
                    "with return type \'Void\'.\n"
                    "Please either add the proper return type to the "
                    "node definition or do not try to create an "
                    "output data channel from this node."
                )
        except TypeError:
            pass

        if index is None:  # out-port type is whole node return type
            # due to the strict typing the out type should match the in type
            if self.is_autogenerated:
                type_out = port_destination.port_type
            else:
                type_out = as_delta_type(self.out_type)

            self.out_ports.append(
                OutPort(NamespacedName(self.name, None),
                        type_out,
                        port_destination, self)
            )
        else:  # out-port type is indexed node return type
            self.out_ports.append(
                OutPort(NamespacedName(self.name, index),
                        as_delta_type(self.out_type.elem_dict[index]),
                        port_destination, self)
            )

        # If this port is going into a port on a different graph,
        # flatten this graph into said graph
        into_graph = port_destination.node.graph
        if into_graph is not self.graph:
            into_graph.flatten(self.graph)

    def add_in_port(self, arg_name: str, in_type: Type, in_port_size: int = 0):
        """Creates an in-port and adds it to my in-port store.

        Parameters
        ----------
        arg_name : str
            Name of the argument this port supplies.
        in_type : Type
            The type that is expected to be supplied for this port.
        in_port_size: int
            Maximum size of the in ports.
            If 0 then size is unlimited.

        Returns
        -------
        InPort
            The created port.
        """
        my_port = InPort(NamespacedName(self.name, arg_name),
                         as_delta_type(in_type),
                         self,
                         in_port_size)
        self.in_ports[my_port.port_name] = my_port
        return my_port

    def _create_upstream_ports(self,
                               required_in_ports: Dict[str, Type],
                               given_nodes: Dict[str, AbstractNode],
                               in_port_size: int = 0):
        """Create the ports going into this node and their
        corresponding out-ports.

        Parameters
        ----------
        required_in_ports
            Dictionary that describes what in-ports we want coming in to
            this node.
        given_nodes
            The nodes that are expected to send results to out in-ports.
        in_port_size
            The maximum size of the node's in ports. If 0 then unlimited size.
        """
        for arg_name, type_wanted in required_in_ports.items():
            if not arg_name in given_nodes:
                raise DeltaIOError(
                    f"Node {self.name} has a mismatch between defined and "
                    f"used input ports.\nArgument \'{arg_name}\' cannot be "
                    f"found in given ports {list(given_nodes.keys())}."
                )

            in_port = self.add_in_port(arg_name, type_wanted, in_port_size)
            given_nodes[arg_name].add_out_port(in_port, None)

    def _ports_from_arguments(self,
                              required_in_params,
                              pos_in_nodes,
                              kw_in_nodes,
                              in_port_size=0):
        """Manages the creation of upstream ports by getting the given nodes
        into the correct data structure.

        Parameters
        ----------
        required_in_params
            Dictionary that describes what in-ports we want coming in to
            this node.
        pos_in_nodes
            Nodes expected to send results to this node, specified
            positionally.
        kw_in_nodes
            Nodes expected to send results to this node, specified by keyword.
        in_port_size
            The maximum size of the node's in ports. If 0 then unlimited size.
        """
        if not isinstance(required_in_params, OrderedDict) and pos_in_nodes:
            raise TypeError("Please specify input parameter types to blocks "
                            "via an OrderedDict if you want to specify "
                            "parameters by position")

        # wrap the given positional in_nodes up with the name of the param they
        # satisfy
        pos_in_nodes_dict = {
            param_name: given_node
            for (given_node, (param_name, _)) in zip(
                pos_in_nodes,
                required_in_params.items()
            )
        }

        self._create_upstream_ports(required_in_params,
                                    {**pos_in_nodes_dict, **kw_in_nodes},
                                    in_port_size=in_port_size)

    @property
    def name(self) -> str:
        return self._name

    @property
    def body(self):
        if self._body:
            return self._body
        else:
            if self.bodies:
                raise ValueError("Please call select_body on node "
                                 f"{self.name} before body access")
            else:
                raise ValueError("No bodies available for selection on node "
                                 f"{self.name}.")


class PythonNode(RealNode):
    """Parent Node type for all Python constructs.

    Attributes
    ----------
    in_queues : Dict[str, DeltaQueue]
        Queues providing input(s).
    out_queues : Dict[str, DeltaQueue]
        Queues consumins output(s).
    sig_stop : Event
        Communication channel through which a runtime simulator or a runtime
        signals `thread_worker` to stop.
    node_key : Optional[str]
        Keyword argument used for providing the node to the block, included for
        debugging purposes.
    lvl : int
        Logging level for the node.
        By default only error logs are displayed.
    """

    def __init__(self,
                 graph,
                 bodies,
                 in_params,
                 pos_in_nodes,
                 kw_in_nodes,
                 in_port_size: int = 0,
                 node_key: Optional[str] = None,
                 out_type=None,
                 name: str = None,
                 lvl: int = logging.ERROR,
                 is_autogenerated: bool = False):
        super().__init__(graph=graph,
                         bodies=bodies,
                         in_params=in_params,
                         out_type=out_type,
                         name=name,
                         lvl=lvl,
                         is_autogenerated=is_autogenerated)

        self._is_const = False

        const_body_no = sum(isinstance(body, PyConstBody) for body in bodies)
        if const_body_no > 0:
            if const_body_no == len(bodies):
                self._is_const = True
                for in_type in self.in_params.values():
                    if isinstance(in_type, DOptional):
                        raise TypeError('Optional input is not allowed '
                                        'for constant nodes')
                    elif not isinstance(in_type, BaseDeltaType):
                        raise TypeError('Unsupported data type')
            else:
                raise ValueError("For a multi-body node, if one body is const "
                                 "all bodies must be const")

        self.in_queues: Dict[str, DeltaQueue] = None
        self.out_queues: Dict[str, DeltaQueue] = None
        self.sig_stop: Event = None
        self.node_key = node_key
        self.in_port_size = in_port_size

        self._ports_from_arguments(self.in_params, pos_in_nodes,
                                   kw_in_nodes, self.in_port_size)

    def set_communications(self, runtime: DeltaPySimulator):
        """Get the in and out queues relating to this node, as well as the
        utility events such as ``sig_stop`` from a runtime simulator or
        a runtime and save them in the instance.

        Parameters
        ----------
        runtime : DeltaPySimulator
            API of a runtime simulator or a runtime.
        """
        self.in_queues = runtime.in_queues[self.name]
        self.out_queues = runtime.out_queues[self.name]
        self.sig_stop = runtime.sig_stop

    def check_stop(self):
        """Check the stop signal, which can be set by a runtime simulator or
        a runtime. If set, it stops the current thread.
        """
        if self.sig_stop.is_set():
            self.log.info(f"Stopped {repr(self)}.")
            sys.exit()

    def receive(self, *args: str) -> Union[Dict[str, Any], Any]:
        """Receives the node's input(s) via in ports.

        Compulsory inputs block the further execution until the data is
        provided, whereas optional inputs do not block.

        Check if the node should stop.

        Parameters
        ----------
        args : str
            Optionally filter inputs. Only the specified ones will be received.

        Returns
        -------
        Union[Dict[str, Any], Any]
            If there is one input is specified via ``args`` it is received as
            an object, overwise the input values are returned as a dictionary.
        """
        # receive all or only selected inputs
        if args:
            in_queues = {name: in_q
                         for name, in_q in self.in_queues.items()
                         if name in args}
        else:
            in_queues = self.in_queues

        values = {}
        for name, in_q in in_queues.items():
            values[name] = in_q.get()
            if not isinstance(values[name], QueueMessage):
                raise TypeError(
                    f"Queue {in_q} from port {repr(in_q._src)} contained "
                    f"an item which was not a QueueMessage: {values[name]}"
                )
            self.msg_log.add_message(self.name, name, values[name])

        # logical clock update
        self._clock = max([self._clock] + [v.clk for v in values.values()])

        # unpack the inner msg
        val = {k: v.msg for k, v in values.items()}

        self.check_stop()

        if all(v == None for v in val.values()):
            # let the Python GIL take a look at the other threads
            sleep(1e-9)

        # if there is just one value to return, unpack it from the dict
        if len(val) == 1 and args:
            val = list(val.values())[0]

        if val:
            self.log.info(f"<- {val}")

        return val

    def send(self, ret: Union[object, Tuple]):
        """Sends the node's output(s) via out ports.

        If sending is blocked this method will lock the execution of the node
        until unblocked or a stop signal is raised, which shall stop
        the execution of the node.

        Parameters
        ----------
        ret : Union[object, Tuple]
            The return value. It is implied by construction that if it is a
            single object then the node has only one out port, otherwise
            a named tuple is used, with the names of the fields matching
            the names of the out ports.
        """
        # Log only non-trivial output(s)
        if not ((isinstance(ret, Tuple) and all(map(lambda x: x is None, ret)))
                or ret is None):
            self.log.info(f"-> {ret}")

        self._clock += 1
        message = QueueMessage(ret, clk=self._clock)

        for out_q in self.out_queues.values():
            while True:
                try:
                    out_q.put(message)
                    self.check_stop()
                    break
                except Full:
                    self.check_stop()
                except:
                    raise

        # let the Python GIL take a look at the other threads
        sleep(1e-9)

    def thread_worker(self, runtime: DeltaPySimulator):
        """Run a regular Python node.

        Waits for input on all the mandatory inputs.
        Then, de-queues the optional inputs and executes its body.
        The output is written to the appropriate output queues.

        Parameters
        ----------
        runtime : DeltaPySimulator
            API of a runtime simulator or a runtime.
        """
        if isinstance(self.body, PyConstBody):
            # Thread worker should never run on a const body
            raise NotImplementedError

        if isinstance(self.body, PyInteractiveBody):
            self.log.debug(f"Running... {self}")
            self.body.eval(self)

        else:
            while True:
                values = self.receive()

                # If a node keyword has been specified for debugging then add
                # the node to the arguments.
                if self.node_key:
                    # the self arg is effectively a const message,
                    # so from time 0
                    values[self.node_key] = self

                self.log.debug("Running...")
                ret = self.body.eval(**values)
                self.send(ret)

    def run_once(self, runtime: DeltaPySimulator):
        """Compute the value of the node and pass it to the output queues.

        The output queues are
        :py:class:`ConstQueue<deltalanguage.runtime.ConstQueue>` that store
        this value and retrive a deepcopy at each request of the receiving
        node.
        This is done purely for optimisation purposes.

        Parameters
        ----------
        runtime : DeltaPySimulator
            API of a runtime simulator or a runtime.
        """
        ret = self.body.eval()
        self.send(ret)

    def capnp(self, capnp_node, capnp_bodies):
        """Generate ``capnp`` form of this node.

        Parameters
        ----------
        capnp_node
            The capnp object of this node.
        capnp_bodies
            List of bodies so we can check if a body is already serialised.
        """
        capnp_node.name = self.name
        capnp_node.init("bodies", len(self.bodies))

        for i_bod, bod in enumerate(self.bodies):

            body_impl = bod.as_serialized

            if isinstance(bod, PyMigenBody):
                body_id = 'migen'
                impl_type = 'verilog'
            elif isinstance(bod, PyInteractiveBody):
                body_id = 'interactive'
                impl_type = 'dillImpl'
            else:
                body_id = 'python'
                impl_type = 'dillImpl'

            def get_body_impl(body):
                return body.__getattr__(body_id).__getattr__(impl_type)

            def set_body_impl(body, impl):
                body.__getattr__(body_id).__setattr__(impl_type, impl)

            for i, body in enumerate(capnp_bodies):
                if body.which() == body_id:
                    if get_body_impl(body) == body_impl:
                        capnp_node.bodies[i_bod] = i
                        break
            else:
                body = capnp_bodies.add()
                body.init(body_id)
                set_body_impl(body, body_impl)
                capnp_node.bodies[i_bod] = len(capnp_bodies) - 1

        # 2. save I/O ports
        self.capnp_ports(capnp_node)

    def capnp_ports(self, capnp_node):
        """Helper method, generates capnp for in/out ports of the node.

        Parameters
        ----------
        capnp_node
            The node of the interest.
        """
        in_ports = capnp_node.init("inPorts", len(self.in_ports))
        for capnp_in_port, in_port in zip(in_ports, self.in_ports.values()):
            in_port.capnp(capnp_in_port)

        out_ports = capnp_node.init("outPorts", len(self.out_ports))
        for capnp_out_port, out_port in zip(out_ports, self.out_ports):
            out_port.capnp(capnp_out_port)

    def capnp_wiring(self, capnp_nodes, capnp_wiring):
        """Generate capnp form of this node's wires.

        Parameters
        ----------
        capnp_nodes
            List of nodes so indexes can be found.
        capnp_wiring
            List of wires so we can add our relevant wires.
        """
        for i, capnp_node in enumerate(capnp_nodes):
            if capnp_node.name == self.name:
                capnp_node_index = i
                break

        for i, out_port in enumerate(self.out_ports):
            capnp_wire = capnp_wiring.add()
            capnp_wire.srcNode = capnp_node_index
            capnp_wire.srcOutPort = i
            out_port.capnp_wiring(capnp_nodes, capnp_wire)

    def set_msg_log(self, msg_log: MessageLog):
        """Sets the log for messages received.

        Parameters
        ----------
        msg_log : MessageLog
            Instance of the message log.
        """
        self.msg_log = msg_log

    def add_body(self, body: Union[BodyTemplate, Callable]):
        """Add some body to this node, using a ``BodyTemplate`` as a recipe
        for how to construct the body to add. This will also enable checks
        about body compatibility to be made.

        Parameters
        ----------
        body : Union[BodyTemplate, Callable]
            ``BodyTemplate`` to add or some object for which
            ``template`` is defined, where ``template`` is a ``BodyTemplate``
        """
        if self.is_const():
            raise NotImplementedError("Adding a body to a const node "
                                      "post-creation is currently unsupported")

        template = body
        if not isinstance(template, BodyTemplate):
            template = body.template

        if template.compatible(self.node_key, self.in_port_size,
                               self.in_params, self.out_type):
            if isinstance(template, MethodBodyTemplate):
                self.bodies.append(template.construct_body(body.__self__))
            else:
                self.bodies.append(template.construct_body())
            self.select_body(override=False)

    def is_const(self):
        """Method to return state of self._is_const.

        Returns
        -------
        bool
            If the current body is constant or not.
        """
        return self._is_const


def as_node(potential_node: Union[AbstractNode, object],
            graph: DeltaGraph) -> PythonNode:
    """Ensures argument is a node and if not makes it into a constant node.

    Parameters
    ----------
    potential_node : Union[AbstractNode, object]
        Node that could be a node or not.
    graph : DeltaGraph
        Graph the node would be in

    Returns
    -------
    PythonNode
        Made for potential_node or potential_node as it was already a node.
    """
    if isinstance(potential_node, AbstractNode):
        return potential_node
    else:
        return PythonNode(graph,
                          [PyConstBody(lambda: potential_node)],
                          {},
                          [],
                          {},
                          out_type=delta_type(potential_node),
                          is_autogenerated=True)


def get_func_in_params_out_type(
    a_func: Callable,
    is_method: bool,
    node_key: Optional[str] = None
) -> Tuple[OrderedDict[str, Union[BaseDeltaType, DOptional]],
           Union[BaseDeltaType, ForkedReturn]]:
    """Helper function to extract input and output types of a node function.

    Parameters
    ----------
    a_func : Callable
        The function to analyse.
    is_method : bool
        Flag to specify if function is a class method.
    node_key : Optional[str]
        Keyword argument used for providing the node to the block, included for
        some logic purposes.

    Returns
    -------
    OrderedDict[str, Union[BaseDeltaType, DOptional]]
        Types of the in parameters.
    Union[BaseDeltaType, ForkedReturn]
        Type of the output the node to be made.

    Raises
    ------
    TypeError
        Raised if either the input or output types aren't specified in the
        function signature.
    """
    func_args = signature(a_func).parameters
    out_type = signature(a_func).return_annotation

    in_params = OrderedDict()
    for i, (arg_name, arg_param) in enumerate(func_args.items()):

        # first argument should always be 'self' for a method
        if i == 0 and is_method:
            continue

        if arg_param.annotation == _empty:
            raise TypeError(
                "Must specify the type of argument " +
                f"'{arg_name}' as annotation in " +
                f"function '{a_func.__name__}'"
            )
        in_params[arg_name] = arg_param.annotation

    in_params = in_params_as_delta_types(in_params, node_key)

    if out_type == _empty:
        out_type = Void

    delta_type_out = as_delta_type(out_type)
    if delta_type_out is not Void:
        if not isinstance(delta_type_out, (BaseDeltaType, ForkedReturn)):
            raise DeltaTypeError(f"Unsupported type={out_type}")

    return in_params, delta_type_out


def in_params_as_delta_types(in_params: Dict[str, Type],
                             node_key: Optional[str] = None
                             ) -> Dict[str, Union[BaseDeltaType, DOptional]]:
    """Take ``in_params`` and convert to DeltaTypes, raising appropriate
    errors, skips and removes ``node_key``.

    Parameters
    ----------
    in_params : Dict[str, Type]
        Dictionary of in_params that need to become DeltaTypes
    node_key : Optional[str], optional
        Node key for node param if available, by default None

    Returns
    -------
    Dict[str, Union[BaseDeltaType, DOptional]]
        Dictionary of in_params as DeltaTypes
    """
    for param_name, param_type in in_params.items():

        if node_key == param_name and param_type == PythonNode:
            in_params.pop(node_key)
            continue

        delta_type_in = as_delta_type(param_type)
        if not isinstance(delta_type_in, (BaseDeltaType, DOptional)):
            raise DeltaTypeError(f"Unsupported type={param_type}")

        in_params[param_name] = delta_type_in
    return in_params
