# Copyright 2016 - 2018  Ternaris.
# SPDX-License-Identifier: AGPL-3.0-only

from __future__ import annotations

import asyncio
import hashlib
from base64 import b32encode
from collections import OrderedDict, namedtuple
from contextlib import suppress
from itertools import count, product
from logging import getLogger
from pathlib import Path
from typing import TYPE_CHECKING

from marv_api import dag
from marv_api.ioctrl import NODE_SCHEMA, Abort, ResourceNotFoundError
from marv_api.iomsgs import GetLogger, GetResourcePath
from marv_api.utils import find_obj

from . import io
from .mixins import Keyed

if TYPE_CHECKING:
    from typing import Any, Callable, Dict

NODE_CACHE: Dict[Callable, Any] = {}


class InputSpec(Keyed, namedtuple('InputSpec', ('name', 'value', 'foreach'))):

    @property
    def key(self):
        value = self.value.key if hasattr(self.value, 'key') else self.value
        return (self.name, value, self.foreach)

    def clone(self, value):
        cls = type(self)
        value = StreamSpec(value) if isinstance(value, Node) else value
        return cls(self.name, value, self.foreach)

    def __repr__(self):
        foreach = 'foreach ' if self.foreach else ''
        return f'<{type(self)} {foreach}{self.name}={self.value!r}>'


class StreamSpec:  # noqa: SIM119  pylint: disable=too-few-public-methods

    def __init__(self, node, name=None):  # pylint: disable=redefined-outer-name
        assert isinstance(node, Node)
        self.node = node
        self.name = name
        self.key = (node.key,) if name is None else (node.key, name)
        self.args = (node,) if name is None else (node, name)


class Node(Keyed):  # pylint: disable=too-many-instance-attributes
    _key = None
    group = None

    @property
    def key(self):
        key = self._key
        if key:
            return key
        return '-'.join([self.specs_hash, self.fullname])

    @property
    def abbrev(self):
        return '.'.join([self.name, self.specs_hash[:10]])

    @staticmethod
    def genhash(specs):
        spec_keys = repr(tuple(x.key for x in sorted(specs.values()))).encode()
        return b32encode(hashlib.sha256(spec_keys).digest()).decode().lower()[:-4]

    @classmethod
    def from_dag_node(cls, func):
        if hasattr(func, '__marv_node__'):
            dnode = func.__marv_node__
        else:
            dnode = func
            func = None

        node = NODE_CACHE.get(dnode)
        if node is not None:
            return node

        if func is None:
            func = find_obj(dnode.function)
        namespace, name = dnode.function.rsplit('.', 1)
        schema = find_obj(dnode.message_schema) if dnode.message_schema is not None else None
        inputs = dnode.inputs
        specs = OrderedDict(
            (
                (
                    name,
                    InputSpec(
                        name=name,
                        value=(
                            value if not isinstance(value, dag.Stream) else
                            StreamSpec(node=Node.from_dag_node(value.node), name=value.name)
                        ),
                        foreach=name == dnode.foreach,
                    ),
                ) for name, value in
                ((name, getattr(inputs, name)) for name in inputs.__fields__.keys())  # pylint: disable=superfluous-parens
            ),
        )
        NODE_CACHE[dnode] = cls(
            func,
            schema=schema,
            version=dnode.version,
            name=name,
            namespace=namespace,
            specs=specs,
            group=dnode.group,
            dag_node=dnode,
        )
        return NODE_CACHE[dnode]

    def __init__(
        self,
        func,
        schema=None,
        version=None,
        name=None,
        namespace=None,
        specs=None,
        group=None,
        dag_node=None,
    ):
        # pylint: disable=too-many-arguments
        # TODO: assert no default values on func, or consider default
        # values (instead of) marv.input() declarations
        # Definitely not for node inputs as they are not passed directly!
        self.func = func
        self.version = version
        self.name = name = func.__name__ if name is None else name
        self.namespace = namespace = func.__module__ if namespace is None else namespace
        self.fullname = (name if not namespace else ':'.join([namespace, name]))
        self.specs_hash = self.genhash(specs or {})
        self.schema = schema
        self.specs = specs or {}
        assert group in (None, False, True, 'ondemand'), group
        self.group = (group if group is not None else any(x.foreach for x in self.specs.values()))
        # TODO: StreamSpec, seriously?
        self.deps = {x.value.node for x in self.specs.values() if isinstance(x.value, StreamSpec)}
        self.alldeps = self.deps.copy()
        self.alldeps.update(x for dep in self.deps for x in dep.alldeps)

        self.consumers = set()
        for dep in self.deps:
            assert self not in dep.consumers, (dep, self)
            dep.consumers.add(self)

        self.dependent = set()
        for dep in self.alldeps:
            assert self not in dep.dependent, (dep, self)
            dep.dependent.add(self)

        if self.fullname == 'marv_nodes:dataset':
            self.load = func.load

        assert dag_node is not None
        self.dag_node = dag_node

    def __call__(self, **inputs):
        return self.func(**inputs)

    async def invoke(self, key_abbrev, inputs=None, site=None):  # noqa: C901
        # pylint: disable=too-many-locals,too-many-branches,too-many-statements
        # We must not write any instance variables, a node is running
        # multiple times in parallel.

        common = []
        foreach_plain = []
        foreach_stream = []
        if inputs is None:
            for spec in self.specs.values():
                assert not isinstance(spec.value, Node), (self, spec.value)
                if isinstance(spec.value, StreamSpec):
                    value = yield io.GetStream(
                        setid=None,
                        node=spec.value.node,
                        name=spec.value.name or 'default',
                    )
                    target = foreach_stream if spec.foreach else common
                else:
                    value = spec.value
                    target = foreach_plain if spec.foreach else common
                target.append((spec.name, value))

        if foreach_plain or foreach_stream:
            log = getLogger(f'marv.node.{key_abbrev}')
            cross = product(*[[(k, x) for x in v] for k, v in foreach_plain])
            if foreach_stream:
                assert len(foreach_stream) == 1, self  # FOR NOW
                cross = list(cross)
                name, stream = foreach_stream[0]
                idx = count()
                while True:
                    value = yield io.Pull(stream, False)
                    if value is None:
                        log.noisy('finished forking')
                        break
                    for inputs in cross:  # pylint: disable=redefined-argument-from-local
                        inputs = dict(inputs)
                        inputs.update(common)
                        inputs[name] = value
                        # TODO: give fork a name
                        _idx = next(idx)
                        log.noisy('FORK %d with: %r', _idx, inputs)
                        yield io.Fork(f'{_idx}', inputs, False)
            else:
                # pylint: disable=redefined-argument-from-local
                for _idx, inputs in enumerate(cross):
                    # TODO: consider deepcopy
                    inputs = dict(inputs)
                    inputs.update(common)
                    log.noisy('FORK %d with: %r', _idx, inputs)
                    yield io.Fork(f'{_idx}', inputs, False)
                log.noisy('finished forking')
        else:
            if inputs is None:
                inputs = dict(common)
            qout, qin = asyncio.Queue(1), asyncio.Queue(1)
            task = asyncio.create_task(
                self.execnode(key_abbrev, inputs, qin=qout, qout=qin, site=site),
                name=key_abbrev,
            )
            while True:
                request = await qin.get()
                if request is None:
                    await task
                    break
                try:
                    response = yield request
                except GeneratorExit:
                    task.cancel()
                    with suppress(asyncio.CancelledError):
                        await task
                    break

                qout.put_nowait(response)

    async def execnode(self, key_abbrev, inputs, qin, qout, site=None):
        NODE_SCHEMA.set(self.schema)
        gen = self.func(**inputs)
        assert hasattr(gen, 'send')
        response = None
        log = getLogger(f'marv.node.{key_abbrev}')
        while True:
            try:
                if isinstance(response, Exception):
                    request = gen.throw(response)
                else:
                    request = gen.send(response)
            # TODO: Abort exception needs to invalidate previous node output
            except (Abort, StopIteration) as exc:
                msg = str(exc)
                if msg:
                    log.warning(msg)
                qout.put_nowait(None)
                return
            except BaseException:  # pylint: disable=broad-except
                qout.put_nowait(None)
                raise

            if isinstance(request, GetLogger):
                response = log
                continue

            if isinstance(request, GetResourcePath):
                rel = Path(request.name)
                if rel.anchor:
                    response = ResourceNotFoundError(request.name)
                else:
                    basepath = site.config.marv.resourcedir
                    fullpath = basepath.joinpath(rel).resolve()
                    if basepath.resolve() not in fullpath.parents or not fullpath.is_file():
                        response = ResourceNotFoundError(request.name)
                    else:
                        response = fullpath
                continue

            qout.put_nowait(request)
            response = await qin.get()

    def __str__(self):
        return self.key

    def __repr__(self):
        return f'<Node {self.abbrev}>'
