# -*- coding: utf-8 -*-

__author__ = r'wsb310@gmail.com'

import os
import typing
import socket
import asyncio
import contextlib
import contextvars

import grpc

from .server import GRPCServer, Router
from .server import UNARY_UNARY_MODE, UNARY_STREAM_MODE, STREAM_UNARY_MODE, STREAM_STREAM_MODE
from .client import GRPCClient

from ...extend.asyncio.base import Utils
from ...extend.asyncio.command import MainProcess, SubProcess


def get_metadata(remote_context: grpc.aio.ServicerContext) -> grpc.aio.Metadata:

    metadata = grpc.aio.Metadata((r'remote-peer', remote_context.peer()))

    for _key, _value in remote_context.invocation_metadata():
        metadata.add(_key, _value)

    return metadata


@contextlib.contextmanager
def grpc_context(
        remote_context: grpc.aio.ServicerContext,
        server_stream_stub: typing.Union[grpc.aio.UnaryStreamCall, grpc.aio.StreamStreamCall] = None
):

    remote_peer = remote_context.peer()

    Utils.log.info(f'grpc proxy open stream: {remote_peer}')

    try:
        yield
    except grpc.aio.UsageError:
        pass
    except Exception as err:
        Utils.log.exception(str(err))
    finally:
        if server_stream_stub is not None:
            server_stream_stub.cancel()

    Utils.log.info(f'grpc proxy close stream: {remote_peer}')


class ProxyUnaryStream:

    def __init__(
            self, server_proxy: GRPCClient, server_method: str,
            remote_request: bytes, remote_context: grpc.aio.ServicerContext
    ):

        self._server_proxy: GRPCClient = server_proxy

        self._server_stream_stub: grpc.aio.UnaryStreamCall = self._server_proxy.unary_stream(
            server_method, remote_request,
            metadata=get_metadata(remote_context)
        )

        self._remote_context: grpc.aio.ServicerContext = remote_context

        self._stream_task: asyncio.Task = asyncio.create_task(self._do_stream_task())

    async def _do_stream_task(self):

        with grpc_context(self._remote_context, self._server_stream_stub):
            async for _message in self._server_stream_stub:
                await self._remote_context.write(_message)

    async def join(self) -> typing.Any:
        return await self._stream_task


class ProxyStreamUnary:

    def __init__(
            self, server_proxy: GRPCClient, server_method: str,
            remote_request: typing.AsyncIterator, remote_context: grpc.aio.ServicerContext
    ):

        self._server_proxy: GRPCClient = server_proxy

        self._server_stream_stub: grpc.aio.StreamUnaryCall = self._server_proxy.stream_unary(
            server_method,
            metadata=get_metadata(remote_context)
        )

        self._remote_request: typing.AsyncIterator = remote_request
        self._remote_context: grpc.aio.ServicerContext = remote_context

        self._stream_task: asyncio.Task = asyncio.create_task(self._do_stream_task())

    async def _do_stream_task(self):

        with grpc_context(self._remote_context):
            async for _message in self._remote_request:
                await self._server_stream_stub.write(_message)

    async def join(self) -> typing.Any:
        return await self._server_stream_stub


class ProxyStreamStream:

    def __init__(
            self, server_proxy: GRPCClient, server_method: str,
            remote_request: typing.AsyncIterator, remote_context: grpc.aio.ServicerContext
    ):

        self._server_proxy: GRPCClient = server_proxy
        self._server_stream_stub: grpc.aio.StreamStreamCall = self._server_proxy.stream_stream(
            server_method,
            metadata=get_metadata(remote_context)
        )

        self._remote_request: typing.AsyncIterator = remote_request
        self._remote_context: grpc.aio.ServicerContext = remote_context

        self._stream_task: asyncio.Task = asyncio.create_task(self._do_stream_task())

    async def _read_upstream(self) -> bytes:

        try:
            return await anext(self._remote_request)
        except StopAsyncIteration:
            pass
        except Exception as err:
            Utils.log.error(f'read upstream error: {str(err)}')
            raise asyncio.CancelledError()

    async def _read_downstream(self) -> bytes:

        try:
            return await self._server_stream_stub.read()
        except Exception as err:
            Utils.log.error(f'read downstream error: {str(err)}')
            raise asyncio.CancelledError()

    async def _do_stream_task(self):

        with grpc_context(self._remote_context, self._server_stream_stub):

            upstream_task = asyncio.create_task(self._read_upstream())
            downstream_task = asyncio.create_task(self._read_downstream())

            while True:

                _tasks_done, _ = await asyncio.wait(
                    [upstream_task, downstream_task],
                    return_when=asyncio.FIRST_COMPLETED
                )

                for _task in _tasks_done:

                    _message = await _task

                    if _message == grpc.aio.EOF:
                        break

                    if _task is upstream_task:
                        await self._server_stream_stub.write(_message)
                        upstream_task = asyncio.create_task(self._read_upstream())
                    elif _task is downstream_task:
                        await self._remote_context.write(_message)
                        downstream_task = asyncio.create_task(self._read_downstream())

    async def join(self) -> typing.Any:
        return await self._stream_task


class ProxyInterceptor(grpc.aio.ServerInterceptor):

    def __init__(self, client: GRPCClient, routers: typing.List[Router]):

        self._method_context: contextvars.ContextVar = contextvars.ContextVar[str](r'grpc_method_context', default=None)

        self._client: GRPCClient = client

        self._handlers: typing.Dict[str, grpc.RpcMethodHandler] = {}

        self._init_handlers(routers)

    def _init_handlers(self, routers: typing.List[Router]):

        for router in routers:

            _service = router.service

            for _func, _mode in router.handlers:

                _method = f'/{_service}/{_func.__name__}'

                if _mode is UNARY_UNARY_MODE:
                    self._handlers[_method] = grpc.unary_unary_rpc_method_handler(self._unary_unary_handler)
                elif _mode is UNARY_STREAM_MODE:
                    self._handlers[_method] = grpc.unary_stream_rpc_method_handler(self._unary_stream_handler)
                elif _mode is STREAM_UNARY_MODE:
                    self._handlers[_method] = grpc.stream_unary_rpc_method_handler(self._stream_unary_handler)
                elif _mode is STREAM_STREAM_MODE:
                    self._handlers[_method] = grpc.stream_stream_rpc_method_handler(self._stream_stream_handler)

    async def _unary_unary_handler(self, request: bytes, context: grpc.aio.ServicerContext) -> typing.Any:

        method = self._method_context.get()

        return await self._client.unary_unary(
            method, request,
            metadata=get_metadata(context)
        )

    async def _unary_stream_handler(self, request: bytes, context: grpc.aio.ServicerContext) -> typing.Any:

        method = self._method_context.get()

        return await ProxyUnaryStream(self._client, method, request, context).join()

    async def _stream_unary_handler(self, request: typing.AsyncIterator, context: grpc.aio.ServicerContext) -> typing.Any:

        method = self._method_context.get()

        return await ProxyStreamUnary(self._client, method, request, context).join()

    async def _stream_stream_handler(self, request: typing.AsyncIterator, context: grpc.aio.ServicerContext) -> typing.Any:

        method = self._method_context.get()

        return await ProxyStreamStream(self._client, method, request, context).join()

    async def intercept_service(
            self,
            continuation: typing.Callable[[grpc.HandlerCallDetails], typing.Awaitable[grpc.RpcMethodHandler]],
            handler_call_details: grpc.HandlerCallDetails
    ) -> grpc.RpcMethodHandler:

        method = getattr(handler_call_details, r'method')

        self._method_context.set(method)

        if method in self._handlers:
            return self._handlers.get(method)
        else:
            return await continuation(handler_call_details)


class GRPCProxy:

    def __init__(
            self, address: typing.Union[str, typing.Tuple[str, int]], routers: typing.List[Router],
            grpc_family: int = socket.AF_INET, grpc_server_credentials: typing.Optional[grpc.ServerCredentials] = None
    ):

        self._routers: typing.List[Router] = routers

        self._grpc_client: typing.Optional[GRPCClient] = None
        self._grpc_server: typing.Optional[GRPCServer] = None

        self._grpc_address: typing.Union[str, typing.Tuple[str, int]] = address
        self._grpc_family: int = grpc_family
        self._grpc_server_credentials: typing.Optional[grpc.ServerCredentials] = grpc_server_credentials

    async def _register(self, request: typing.AsyncIterator, context: grpc.aio.ServicerContext):

        target = dict(context.invocation_metadata()).get(r'worker_entry_point')

        if target:
            context.add_done_callback(lambda _: Utils.create_task(self._grpc_client.clear_targets([target])))
        else:
            await context.abort(grpc.StatusCode.CANCELLED)

        await self._grpc_client.append_target(target)

        Utils.log.success(f'grpc proxy connected worker: {target}')

        with Utils.suppress(grpc.RpcError):
            async for _message in request:
                Utils.log.warning(f'grpc proxy received message: {_message}')

    async def open(self):

        self._grpc_client: GRPCClient = GRPCClient(
            request_serializer=None,
            response_deserializer=None
        )

        self._grpc_server: GRPCServer = GRPCServer(
            interceptors=[ProxyInterceptor(self._grpc_client, self._routers)],
            request_deserializer=None,
            response_serializer=None
        )

        self._grpc_server.register(r'', [self._register], mode=STREAM_STREAM_MODE)

        await self._grpc_server.start(self._grpc_address, self._grpc_family, self._grpc_server_credentials)

    async def close(self):

        await self._grpc_server.stop()
        await self._grpc_client.close()

    async def join(self):
        await self._grpc_server.wait()


class GRPCEntry:

    def __init__(self, proxy_address: typing.Union[str, typing.Tuple[str, int]], routers: typing.List[Router]):

        self._proxy_address: str = self._decode_address(proxy_address)
        self._server_address: str = f'/tmp/grpc_worker_{Utils.getpid()}.sock'

        self._grpc_server: typing.Optional[GRPCServer] = None
        self._grpc_routers: typing.List[Router] = routers

        self._proxy_client: typing.Optional[GRPCClient] = None
        self._proxy_client_stub: typing.Optional[grpc.aio.StreamStreamCall] = None

    @staticmethod
    def _decode_address(proxy_address: typing.Union[str, typing.Tuple[str, int]]) -> str:

        if isinstance(proxy_address, str):
            proxy_address = Utils.split_str(proxy_address, r':')

        if proxy_address[0] == r'0.0.0.0':
            proxy_address =(r'localhost', proxy_address[1])

        return Utils.join_str(proxy_address, r':')

    async def start(self):

        if os.path.exists(self._server_address):
            os.remove(self._server_address)

        self._grpc_server = GRPCServer()

        for _router in self._grpc_routers:
            self._grpc_server.bind_router(_router)

        await self._grpc_server.start(self._server_address, family=socket.AF_UNIX)

        self._proxy_client = GRPCClient()
        await self._proxy_client.open([self._proxy_address])

        self._proxy_client_stub = self._proxy_client.stream_stream(
            r'//_register',
            metadata=grpc.aio.Metadata((r'worker_entry_point', r'unix:' + self._server_address))
        )

    async def stop(self):
        await self._proxy_client.close()
        await self._grpc_server.stop()

    async def join(self):

        with Utils.suppress(grpc.RpcError):
            async for _message in self._proxy_client_stub:
                Utils.log.warning(f'grpc worker received message: {_message}')

        await self._grpc_server.wait()


class GRPCMainProcess(MainProcess):

    def __init__(
            self, address: typing.Union[str, typing.Tuple[str, int]], routers: typing.List[Router],
            worker: typing.Type[r'GRPCWorker'], sub_process_num: int, *,
            grpc_family: int = socket.AF_INET, grpc_server_credentials: typing.Optional[grpc.ServerCredentials] = None,
            cpu_affinity: bool = False, join_timeout: int = 10
    ):

        super().__init__(
            worker.create, sub_process_num,
            cpu_affinity=cpu_affinity, join_timeout=join_timeout,
            proxy_address=address, routers=routers
        )

        self._grpc_proxy: GRPCProxy = GRPCProxy(address, routers, grpc_family, grpc_server_credentials)

    async def initialize(self):
        await self._grpc_proxy.open()

    async def release(self):
        await self._grpc_proxy.close()


class GRPCWorker(SubProcess):

    def __init__(self, process_num: int, proxy_address: typing.Union[str, typing.Tuple[str, int]], routers: typing.List[Router]):

        super().__init__(process_num)

        self._grpc_entry: GRPCEntry = GRPCEntry(proxy_address, routers)

    async def initialize(self):
        await self._grpc_entry.start()

    async def release(self):
        await self._grpc_entry.stop()

    async def _execute(self):
        await self._grpc_entry.join()
