# coding: utf-8
# /*##########################################################################
#
# Copyright (c) 2017 European Synchrotron Radiation Facility
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#
# ###########################################################################*/

__authors__ = ["H.Payno"]
__license__ = "MIT"
__date__ = "17/12/2018"


from xml.etree.ElementTree import TreeBuilder, Element, ElementTree
from collections import defaultdict
from itertools import count, chain
import json
import pprint
import base64
import pickle
import logging
from .node import Node
from ..workflow import ActorFactory, StartActor, StopActor, JoinUntilStopSignal

_logger = logging.getLogger(__name__)


class Scheme(object):
    """
    class to define a workflow scheme from nodes and links

    :param list nodes:
    :param list links:
    """
    def __init__(self, nodes=None, links=None):
        self.title = ''
        self.description = ''
        self.links = {}
        """keys are link ID, values are Link"""
        if links is not None:
            for link in links:
                self.links[link.id] = link
        self.nodes = nodes or []
        """list of nodes"""
        self.nodes_dict = {}
        """dict with node id as key and node as value"""
        for node in self.nodes:
            self.nodes_dict[node.id] = node

        if links is not None:
            self._update_nodes_from_links()

        self._actor_factory = {}
        for node in self.nodes:
            self._actor_factory[node] = ActorFactory(node)

        # deal with connect
        for node in self.nodes:
            actor_factory = self._actor_factory[node]
            for downstream_node in node.downstream_nodes:
                downstream_actor_factory = self._actor_factory[downstream_node]
                actor_factory.connect(downstream_actor_factory)

        # add start actor
        self._start_actor = StartActor()
        for node in self.start_nodes():
            actor_factory = self._actor_factory[node]
            self._start_actor.connect(actor_factory)

        def connect_finals_nodes(actor):
            # add end actor
            for node in self.finalsNodes():
                actor_factory = self._actor_factory[node]
                actor_factory.connect(actor)

        self._end_actor = StopActor()

        if self.has_final_join():
            self._join_actor = JoinUntilStopSignal('stop join')
            connect_finals_nodes(self._join_actor)
            self._join_actor.connect(self._end_actor)
        else:
            connect_finals_nodes(self._end_actor)

    def finalsNodes(self):
        """

        :return: list of final nodes (with no output) and which hasn't any
                 control node upstream
        """
        res = []
        for node in self.nodes:
            assert isinstance(node, Node)
            if node.isfinal():
                res.append(node)
        return res

    def start_nodes(self):
        res = []
        for node in self.nodes:
            assert isinstance(node, Node)
            if node.isstart():
                res.append(node)
        return res

    def endlessNodes(self):
        res = []
        for node in self.nodes:
            assert isinstance(node, Node)
            if node.endless is True:
                res.append(node)
        return res

    def save_to(self, output_file):
        """
        Save the scheme as an xml formated file to `stream`
        """
        tree = self.scheme_to_etree(data_format="literal")
        indent(tree.getroot(), 0)

        tree.write(output_file)

    def scheme_to_etree(self, data_format="literal", pickle_fallback=False):
        """
        Return an `xml.etree.ElementTree` representation of the `scheme.
        """
        builder = TreeBuilder(element_factory=Element)
        builder.start("scheme", {"version": "2.0",
                                 "title": self.title or "",
                                 "description": self.description or ""})

        # Nodes
        node_ids = defaultdict(count().__next__)
        builder.start("nodes", {})
        for node in self.nodes:  # type: SchemeNode
            attrs = {"id": node.id,
                     # "name": node.name,
                     "qualified_name": node._qualified_name,
                     # "project_name": node.project_name or "",
                     # "version": node.version or "",
                     # "title": node.title,
                     }

            if type(node) is not Node:
                attrs["scheme_node_type"] = "%s.%s" % (type(node).__name__,
                                                       type(node).__module__)
            builder.start("node", attrs)
            builder.end("node")

        builder.end("nodes")

        # Links
        link_ids = defaultdict(count().__next__)
        builder.start("links", {})
        for link in self.links:  # type: SchemeLink
            source = link.source_node_id
            sink = link.sink_node_id
            source_id = node_ids[source]
            sink_id = node_ids[sink]
            attrs = {"id": str(link_ids[link]),
                     "source_node_id": str(source_id),
                     "sink_node_id": str(sink_id),
                     "source_channel": link.source_channel,
                     "sink_channel": link.sink_channel,
                     "enabled": "true" if link.enabled else "false",
                     }
            builder.start("link", attrs)
            builder.end("link")

        builder.end("links")

        # Annotations
        annotation_ids = defaultdict(count().__next__)
        builder.start("thumbnail", {})
        builder.end("thumbnail")

        # Node properties/settings
        builder.start("node_properties", {})
        for node in self.nodes:
            data = None
            if node.properties:
                try:
                    data, format = dumps(node.properties, format=data_format,
                                         pickle_fallback=pickle_fallback)
                except Exception:
                    _logger.error("Error serializing properties for node %r",
                              node.title, exc_info=True)
                if data is not None:
                    builder.start("properties",
                                  {"node_id": str(node_ids[node]),
                                   "format": format})
                    builder.data(data)
                    builder.end("properties")

        builder.end("node_properties")

        builder.end("scheme")
        root = builder.close()
        tree = ElementTree(root)
        return tree

    def _update_nodes_from_links(self):
        """
        Update upstream and downstream nodes from links definition
        """
        self._clear_nodes_connections()
        for link_id, link in self.links.items():
            source_node = self.nodes_dict[link.source_node_id]
            sink_node = self.nodes_dict[link.sink_node_id]
            source_node.downstream_nodes.add(self.nodes_dict[link.sink_node_id])
            sink_node.upstream_nodes.add(self.nodes_dict[link.source_node_id])

    def _clear_nodes_connections(self):
        """
        clear for all nodes downstream and upstream nodes
        """
        for node in self.nodes:
            assert isinstance(node, Node)
            node.downstream_nodes = set()
            node.upstream_nodes = set()

    def has_final_join(self):
        """True if we need to send a 'end' signal before closing the workflow
        This is needed for DataList and DataWatcher
        """
        for node in self.nodes:
            if node.need_stop_join:
                return True
        return False


def contains_control_nodes(nodes_list):
    for _node in nodes_list:
        if _node.endless or contains_control_nodes(_node.upstream_nodes):
            return True
    return False



def indent(element, level=0, indent="\t"):
    """
    Indent an instance of a :class:`Element`. Based on
    (http://effbot.org/zone/element-lib.htm#prettyprint).

    """
    def empty(text):
        return not text or not text.strip()

    def indent_(element, level, last):
        child_count = len(element)

        if child_count:
            if empty(element.text):
                element.text = "\n" + indent * (level + 1)

            if empty(element.tail):
                element.tail = "\n" + indent * (level + (-1 if last else 0))

            for i, child in enumerate(element):
                indent_(child, level + 1, i == child_count - 1)

        else:
            if empty(element.tail):
                element.tail = "\n" + indent * (level + (-1 if last else 0))

    return indent_(element, level, True)



def dumps(obj, format="literal", prettyprint=False, pickle_fallback=False):
    """
    Serialize `obj` using `format` ('json' or 'literal') and return its
    string representation and the used serialization format ('literal',
    'json' or 'pickle').

    If `pickle_fallback` is True and the serialization with `format`
    fails object's pickle representation will be returned

    """
    if format == "literal":
        try:
            return (literal_dumps(obj, prettyprint=prettyprint, indent=1),
                    "literal")
        except (ValueError, TypeError) as ex:
            if not pickle_fallback:
                raise

            _logger.debug("Could not serialize to a literal string")

    elif format == "json":
        try:
            return (json.dumps(obj, indent=1 if prettyprint else None),
                    "json")
        except (ValueError, TypeError):
            if not pickle_fallback:
                raise

            _logger.debug("Could not serialize to a json string")

    elif format == "pickle":
        return base64.encodebytes(pickle.dumps(obj)).decode('ascii'), "pickle"

    else:
        raise ValueError("Unsupported format %r" % format)

    if pickle_fallback:
        _logger.warning("Using pickle fallback")
        return base64.encodebytes(pickle.dumps(obj)).decode('ascii'), "pickle"
    else:
        raise Exception("Something strange happened.")



# This is a subset of PyON serialization.
def literal_dumps(obj, prettyprint=False, indent=4):
    """
    Write obj into a string as a python literal.
    """
    memo = {}
    NoneType = type(None)

    def check(obj):
        if type(obj) in [int, float, bool, NoneType, str, bytes]:
            return True

        if id(obj) in memo:
            raise ValueError("{0} is a recursive structure".format(obj))

        memo[id(obj)] = obj

        if type(obj) in [list, tuple]:
            return all(map(check, obj))
        elif type(obj) is dict:
            return all(map(check, chain(iter(obj.keys()), iter(obj.values()))))
        else:
            raise TypeError("{0} can not be serialized as a python "
                             "literal".format(type(obj)))

    check(obj)

    if prettyprint:
        return pprint.pformat(obj, indent=indent)
    else:
        return repr(obj)
