#
#    ICRAR - International Centre for Radio Astronomy Research
#    (c) UWA - The University of Western Australia, 2017
#    Copyright by UWA (in the framework of the ICRAR)
#    All rights reserved
#
#    This library is free software; you can redistribute it and/or
#    modify it under the terms of the GNU Lesser General Public
#    License as published by the Free Software Foundation; either
#    version 2.1 of the License, or (at your option) any later version.
#
#    This library is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
#    Lesser General Public License for more details.
#
#    You should have received a copy of the GNU Lesser General Public
#    License along with this library; if not, write to the Free Software
#    Foundation, Inc., 59 Temple Place, Suite 330, Boston,
#    MA 02111-1307  USA
#
"""Utilities to emulate the `dask.delayed` function"""

import base64
import contextlib
import logging
import pickle
import socket
import struct
import time

from . import utils, droputils
from .apps import pyfunc
from .common import dropdict, Categories
from .ddap_protocol import DROPStates
from .drop import BarrierAppDROP
from .exceptions import InvalidDropException

logger = logging.getLogger(__name__)


class ResultTransmitter(BarrierAppDROP):
    """Collects data from all inputs and transmits it to whomever connects to
    the given host/port"""

    def initialize(self, **kwargs):
        BarrierAppDROP.initialize(self, input_error_threshold=100, **kwargs)
        self.host = self._getArg(kwargs, "host", "127.0.0.1")
        self.port = self._getArg(kwargs, "port", None)
        if self.port is None:
            raise InvalidDropException(self, "Missing port parameter")

    def run(self):
        def read_result(x):
            if x.status == DROPStates.ERROR:
                return "Error"
            return pickle.loads(droputils.allDropContents(x))

        results = map(read_result, self.inputs)  # @UndefinedVariable
        results = list(results)
        if len(self.inputs) == 1:
            results = results[0]
        results = pickle.dumps(results)

        s = socket.socket()
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, True)
        s.bind((self.host, self.port))
        s.listen(1)
        client, _ = s.accept()
        with contextlib.closing(client):
            client = client.makefile("wb")
            client.write(struct.pack(">i", len(results)))
            client.write(results)


def _get_client(**kwargs):

    if "client" in kwargs:
        return kwargs["client"]

    from .manager.client import NodeManagerClient
    from .manager import constants

    host = kwargs.get("host", "127.0.0.1")
    port = kwargs.get("port", constants.NODE_DEFAULT_REST_PORT)
    timeout = kwargs.get("timeout", None)
    return NodeManagerClient(host, port, timeout)


def _is_list_of_delayeds(x):
    return isinstance(x, (list, tuple)) and len(x) > 0 and isinstance(x[0], _DataDrop)


def compute(value, **kwargs):
    """Returns the result of the (possibly) delayed computation by sending
    the graph to a Drop Manager and waiting for the result to arrive back"""

    # Support calling compute with a list of DelayedDrops
    if _is_list_of_delayeds(value):
        value = _DelayedDrops(*value)

    graph = value.get_graph()
    port = 10000
    # Add one final application that will wait for all results
    # and transmit them back to us
    transmitter_oid = "-1"
    transmitter = dropdict(
        {
            "type": "app",
            "app": "dlg.dask_emulation.ResultTransmitter",
            "oid": transmitter_oid,
            "port": port,
            "nm": "result transmitter",
        }
    )
    for leaf_oid in droputils.get_leaves(graph.values()):
        graph[leaf_oid].addConsumer(transmitter)
    graph[transmitter_oid] = transmitter

    graph = list(graph.values())

    # Submit and wait
    session_id = "session-%f" % time.time()
    client = _get_client(**kwargs)
    client.create_session(session_id)
    client.append_graph(session_id, graph)
    client.deploy_session(session_id, completed_uids=droputils.get_roots(graph))

    timeout = kwargs.get("timeout", None)
    s = utils.connect_to("localhost", port, timeout)
    s.settimeout(timeout)
    with contextlib.closing(s):
        s = s.makefile("rb")
        nbytes = struct.unpack(">i", s.read(4))[0]
        ret = pickle.loads(s.read(nbytes))
        logger.info("Received %r from graph computation", ret)
        return ret


class _DelayedDrop(object):
    _drop_count = 0

    def __init__(self, producer=None):
        self._dropdict = None
        self.producer = producer
        self.inputs = []

    @property
    def next_drop_oid(self):
        i = _DelayedDrop._drop_count
        _DelayedDrop._drop_count += 1
        return i

    @property
    def dropdict(self):
        if self._dropdict is None:
            self._dropdict = self.make_dropdict()
        return self._dropdict

    def reset(self):
        self._dropdict = None

    @property
    def oid(self):
        return self.dropdict["oid"]

    def compute(self, **kwargs):
        return compute(self, **kwargs)

    def get_graph(self):
        _DelayedDrop._drop_count = 0
        graph = {}
        visited = set()
        self._to_physical_graph(visited, graph)
        for d in visited:
            d.reset()
        return graph

    def _append_to_graph(self, visited, graph):
        if self in visited:
            return
        oid = str(self.next_drop_oid)
        dd = self.dropdict
        dd["oid"] = oid
        visited.add(self)
        graph[oid] = dd
        logger.debug("Appended %r/%s to the Physical Graph", self, oid)

    def _to_physical_graph(self, visited, graph):

        self._append_to_graph(visited, graph)

        dependencies = list(self.inputs)
        if self.producer:
            dependencies.append(self.producer)
        for d in dependencies:
            if isinstance(d, list):
                d = tuple(d)
            if d in visited:
                self._add_upstream(d)
                continue

            d = d._to_physical_graph(visited, graph)
            self._add_upstream(d)

        return self

    def _add_upstream(self, upstream):
        """Link the given drop as either a producer or input of this drop"""
        self_dd = self.dropdict
        up_dd = upstream.dropdict
        if isinstance(self, _DataDrop):
            self_dd.addProducer(up_dd)
            logger.debug(
                "Set %r/%s as producer of %r/%s", upstream, upstream.oid, self, self.oid
            )
        else:
            self_dd.addInput(up_dd)
            logger.debug(
                "Set %r/%s as input of %r/%s", upstream, upstream.oid, self, self.oid
            )


class _Listifier(BarrierAppDROP):
    """Returns a list with all objects as contents"""

    def run(self):
        self.outputs[0].write(
            pickle.dumps(
                [pickle.loads(droputils.allDropContents(x)) for x in self.inputs]
            )
        )


class _DelayedDrops(_DelayedDrop):
    """One or more _DelayedDrops treated as a single item"""

    def __init__(self, *drops):
        super(_DelayedDrops, self).__init__()
        self.drops = drops
        self.inputs.extend(drops)
        logger.debug("Created %r", self)

    def _to_physical_graph(self, visited, graph):
        output = _DataDrop(producer=self)
        output._append_to_graph(visited, graph)

        self._append_to_graph(visited, graph)
        output._add_upstream(self)

        for d in self.drops:
            d._to_physical_graph(visited, graph)
            self._add_upstream(d)

        return output

    def __iter__(self):
        return iter(self.drops)

    def __len__(self):
        return len(self.drops)

    def __getitem__(self, i):
        return self.drops[i]

    def make_dropdict(self):
        return dropdict(
            {"type": "app", "app": "dlg.dask_emulation._Listifier", "nm": "listifier"}
        )

    def __repr__(self):
        return "<_DelayedDrops n=%d>" % (len(self.drops),)


class _AppDrop(_DelayedDrop):
    """Defines a PyFuncApp drop for a given function `f`"""

    def __init__(self, f, nout):
        _DelayedDrop.__init__(self)
        self.f = f
        self.fname = None
        if hasattr(f, "__name__"):
            self.fname = f.__name__
        self.fcode, self.fdefaults = pyfunc.serialize_func(f)
        self.original_kwarg_names = []
        self.nout = nout
        logger.debug("Created %r", self)

    def make_dropdict(self):

        self.kwarg_names = list(self.original_kwarg_names)
        self.kwarg_names.reverse()
        my_dropdict = dropdict(
            {"type": "app", "app": "dlg.apps.pyfunc.PyFuncApp", "func_arg_mapping": {}}
        )
        if self.fname is not None:
            simple_fname = self.fname.split(".")[-1]
            my_dropdict["func_name"] = self.fname
            my_dropdict["nm"] = simple_fname
        if self.fcode is not None:
            my_dropdict["func_code"] = utils.b2s(base64.b64encode(self.fcode))
        if self.fdefaults:
            my_dropdict["func_defaults"] = self.fdefaults
        return my_dropdict

    def _add_upstream(self, dep):
        _DelayedDrop._add_upstream(self, dep)
        if self.kwarg_names:
            name = self.kwarg_names.pop()
            if name is not None:
                logger.debug(
                    "Adding %s/%s to function mapping for %s", name, dep.oid, self.fname
                )
                self.dropdict["func_arg_mapping"][name] = dep.oid

    def _to_delayed_arg(self, arg):

        logger.info("Turning into delayed arg for %r: %r", self, arg)
        if isinstance(arg, _DelayedDrop):
            return arg

        # Turn lists/tuples of _DataDrop objects into a _DelayedDrops
        if _is_list_of_delayeds(arg):
            return _DelayedDrops(*arg)

        # Plain data gets turned into a _DataDrop
        return _DataDrop(pydata=arg)

    def __call__(self, *args, **kwargs):

        logger.debug(
            "Delayed function %s called with %d args and %d kwargs",
            self.fname,
            len(args),
            len(kwargs),
        )
        for arg in args:
            self.inputs.append(self._to_delayed_arg(arg))
            self.original_kwarg_names.append(None)

        for name, arg in kwargs.items():
            self.inputs.append(self._to_delayed_arg(arg))
            self.original_kwarg_names.append(name)

        if self.nout is None:
            return _DataDrop(producer=self)

        return _DataDropSequence(nout=self.nout, producer=self)

    def __repr__(self):
        return "<_DelayedApp fname=%s, nout=%s>" % (self.fname, str(self.nout))


_no_data = object()


class _DataDrop(_DelayedDrop):
    """Defines an in-memory drop"""

    def __init__(self, producer=None, pydata=_no_data):
        _DelayedDrop.__init__(self, producer)

        if bool(producer is None) == bool(pydata is _no_data):
            raise ValueError("either producer or pydata must be not None")
        self.pydata = pydata
        logger.debug("Created %r", self)

    def make_dropdict(self):
        my_dropdict = dropdict({"type": "plain", "storage": Categories.MEMORY})
        if not self.producer:
            my_dropdict["pydata"] = pyfunc.serialize_data(self.pydata)
        return my_dropdict

    def __repr__(self):
        if not self.producer:
            return "<_DataDrop, pydata=%r>" % (self.pydata,)
        return "<_DataDrop, producer=%r>" % self.producer


class _DataDropSequence(_DataDrop):
    """One or more _DataDrops that can be subscribed"""

    def __init__(self, nout, producer):
        super(_DataDrop, self).__init__(producer=producer)
        self.nout = nout
        logger.debug("Created %r", self)

    def __iter__(self):
        for i in range(self.nout):
            yield self[i]

    def __len__(self):
        return self.nout

    def __getitem__(self, i):
        return delayed(lambda x, i: x[i])(self, i)

    def __repr__(self):
        return "<_DataDropSequence nout=%d, producer=%r>" % (self.nout, self.producer)


def delayed(x, *args, **kwargs):
    """Like dask.delayed, but quietly swallowing anything other than `nout`"""
    if "nout" in kwargs:
        nout = kwargs["nout"]
    elif args:
        nout = args[0]
    else:
        nout = None
    if callable(x):
        return _AppDrop(x, nout=nout)
    return _DataDrop(pydata=x)
