"""
Dispatchers
===========
"""

from __future__ import annotations

from asyncio import Future, Queue
from asyncio.tasks import FIRST_COMPLETED, ensure_future, gather, wait
from contextlib import asynccontextmanager
from logging import getLogger
from typing import (
    Any,
    AsyncIterator,
    Awaitable,
    Callable,
    Dict,
    List,
    NamedTuple,
    Sequence,
    Tuple,
    cast,
)
from weakref import WeakSet

from anyio import create_task_group
from jsonpatch import apply_patch, make_patch

from idom.core.vdom import VdomJson
from idom.utils import Ref

from .layout import LayoutEvent, LayoutUpdate
from .proto import LayoutType


logger = getLogger(__name__)


SendCoroutine = Callable[["VdomJsonPatch"], Awaitable[None]]
RecvCoroutine = Callable[[], Awaitable[LayoutEvent]]


async def dispatch_single_view(
    layout: LayoutType[LayoutUpdate, LayoutEvent],
    send: SendCoroutine,
    recv: RecvCoroutine,
) -> None:
    """Run a dispatch loop for a single view instance"""
    with layout:
        async with create_task_group() as task_group:
            task_group.start_soon(_single_outgoing_loop, layout, send)
            task_group.start_soon(_single_incoming_loop, layout, recv)


SharedViewDispatcher = Callable[[SendCoroutine, RecvCoroutine], Awaitable[None]]
_SharedViewDispatcherFuture = Callable[[SendCoroutine, RecvCoroutine], "Future[None]"]


@asynccontextmanager
async def create_shared_view_dispatcher(
    layout: LayoutType[LayoutUpdate, LayoutEvent],
) -> AsyncIterator[_SharedViewDispatcherFuture]:
    """Enter a dispatch context where all subsequent view instances share the same state"""
    with layout:
        (
            dispatch_shared_view,
            model_state,
            all_patch_queues,
        ) = await _make_shared_view_dispatcher(layout)

        dispatch_tasks: List[Future[None]] = []

        def dispatch_shared_view_soon(
            send: SendCoroutine, recv: RecvCoroutine
        ) -> Future[None]:
            future = ensure_future(dispatch_shared_view(send, recv))
            dispatch_tasks.append(future)
            return future

        yield dispatch_shared_view_soon

        gathered_dispatch_tasks = gather(*dispatch_tasks, return_exceptions=True)

        while True:
            (
                update_future,
                dispatchers_completed_future,
            ) = await _wait_until_first_complete(
                layout.render(),
                gathered_dispatch_tasks,
            )

            if dispatchers_completed_future.done():
                update_future.cancel()
                break
            else:
                patch = VdomJsonPatch.create_from(update_future.result())

            model_state.current = patch.apply_to(model_state.current)
            # push updates to all dispatcher callbacks
            for queue in all_patch_queues:
                queue.put_nowait(patch)


def ensure_shared_view_dispatcher_future(
    layout: LayoutType[LayoutUpdate, LayoutEvent],
) -> Tuple[Future[None], SharedViewDispatcher]:
    """Ensure the future of a dispatcher created by :func:`create_shared_view_dispatcher`"""
    dispatcher_future: Future[SharedViewDispatcher] = Future()

    async def dispatch_shared_view_forever() -> None:
        with layout:
            (
                dispatch_shared_view,
                model_state,
                all_patch_queues,
            ) = await _make_shared_view_dispatcher(layout)

            dispatcher_future.set_result(dispatch_shared_view)

            while True:
                patch = await render_json_patch(layout)
                model_state.current = patch.apply_to(model_state.current)
                # push updates to all dispatcher callbacks
                for queue in all_patch_queues:
                    queue.put_nowait(patch)

    async def dispatch(send: SendCoroutine, recv: RecvCoroutine) -> None:
        await (await dispatcher_future)(send, recv)

    return ensure_future(dispatch_shared_view_forever()), dispatch


async def render_json_patch(layout: LayoutType[LayoutUpdate, Any]) -> VdomJsonPatch:
    """Render a class:`VdomJsonPatch` from a layout"""
    return VdomJsonPatch.create_from(await layout.render())


class VdomJsonPatch(NamedTuple):
    """An object describing an update to a :class:`Layout` in the form of a JSON patch"""

    path: str
    """The path where changes should be applied"""

    changes: List[Dict[str, Any]]
    """A list of JSON patches to apply at the given path"""

    def apply_to(self, model: VdomJson) -> VdomJson:
        """Return the model resulting from the changes in this update"""
        return cast(
            VdomJson,
            apply_patch(
                model, [{**c, "path": self.path + c["path"]} for c in self.changes]
            ),
        )

    @classmethod
    def create_from(cls, update: LayoutUpdate) -> VdomJsonPatch:
        """Return a patch given an layout update"""
        return cls(update.path, make_patch(update.old or {}, update.new).patch)


async def _make_shared_view_dispatcher(
    layout: LayoutType[LayoutUpdate, LayoutEvent],
) -> Tuple[SharedViewDispatcher, Ref[Any], WeakSet[Queue[VdomJsonPatch]]]:
    update = await layout.render()
    model_state = Ref(update.new)

    # We push updates to queues instead of pushing directly to send() callbacks in
    # order to isolate the render loop from any errors dispatch callbacks might
    # raise.
    all_patch_queues: WeakSet[Queue[VdomJsonPatch]] = WeakSet()

    async def dispatch_shared_view(send: SendCoroutine, recv: RecvCoroutine) -> None:
        patch_queue: Queue[VdomJsonPatch] = Queue()
        async with create_task_group() as inner_task_group:
            all_patch_queues.add(patch_queue)
            effective_update = LayoutUpdate("", None, model_state.current)
            await send(VdomJsonPatch.create_from(effective_update))
            inner_task_group.start_soon(_single_incoming_loop, layout, recv)
            inner_task_group.start_soon(_shared_outgoing_loop, send, patch_queue)
        return None

    return dispatch_shared_view, model_state, all_patch_queues


async def _single_outgoing_loop(
    layout: LayoutType[LayoutUpdate, LayoutEvent], send: SendCoroutine
) -> None:
    while True:
        await send(await render_json_patch(layout))


async def _single_incoming_loop(
    layout: LayoutType[LayoutUpdate, LayoutEvent], recv: RecvCoroutine
) -> None:
    while True:
        await layout.deliver(await recv())


async def _shared_outgoing_loop(
    send: SendCoroutine, queue: Queue[VdomJsonPatch]
) -> None:
    while True:
        await send(await queue.get())


async def _wait_until_first_complete(
    *tasks: Awaitable[Any],
) -> Sequence[Future[Any]]:
    futures = [ensure_future(t) for t in tasks]
    await wait(futures, return_when=FIRST_COMPLETED)
    return futures
