"""
Mocks for the gRPC servers.

@sa @ref scripts/docs/en/userver/tutorial/grpc_service.md
"""

# pylint: disable=no-member
import asyncio
import contextlib
import functools
from typing import Callable
from typing import List
from typing import Optional

import grpc
import pytest

from testsuite.utils import callinfo

DEFAULT_PORT = 8091


class GrpcServiceMock:
    def __init__(self, servicer, methods):
        self.servicer = servicer
        self._known_methods = methods
        self._methods = {}

    def get(self, method, default):
        return self._methods.get(method, default)

    def reset_handlers(self):
        self._methods = {}

    @contextlib.contextmanager
    def mock(self):
        try:
            yield self.install_handler
        finally:
            self._methods = {}

    def install_handler(self, method: str):
        def decorator(func):
            if method not in self._known_methods:
                raise RuntimeError(
                    f'Trying to mock unknown grpc method {method}',
                )

            wrapped = callinfo.acallqueue(func)
            self._methods[method] = wrapped
            return wrapped

        return decorator


@pytest.fixture(scope='session')
def grpc_mockserver_endpoint(pytestconfig, get_free_port) -> str:
    """
    Returns the gRPC endpoint to start the mocking server that is set by
    command line `--grpc-mockserver-host` and `--grpc-mockserver-port` options.

    For port 0, picks some free port.

    Override this fixture to customize the endpoint used by gRPC mockserver.

    @snippet samples/grpc_service/tests/conftest.py  Prepare configs
    @ingroup userver_testsuite_fixtures
    """
    port = pytestconfig.option.grpc_mockserver_port
    if pytestconfig.option.service_wait or pytestconfig.option.service_disable:
        port = port or DEFAULT_PORT
    if port == 0:
        port = get_free_port()
    return f'{pytestconfig.option.grpc_mockserver_host}:{port}'


@pytest.fixture(scope='session')
async def grpc_mockserver(grpc_mockserver_endpoint) -> grpc.aio.Server:
    """
    Returns the gRPC mocking server.

    @snippet samples/grpc_service/tests/conftest.py  Prepare server mock
    @ingroup userver_testsuite_fixtures
    """
    server = grpc.aio.server()
    server.add_insecure_port(grpc_mockserver_endpoint)
    await server.start()

    try:
        yield server
    finally:

        async def stop_server():
            await server.stop(grace=None)
            await server.wait_for_termination()

        stop_server_task = asyncio.shield(asyncio.create_task(stop_server()))
        # asyncio.shield does not protect our await from cancellations, and we
        # really need to wait for the server stopping before continuing.
        try:
            await stop_server_task
        except asyncio.CancelledError:
            await stop_server_task
            # Propagate cancellation when we are done
            raise


@pytest.fixture(scope='session')
def create_grpc_mock():
    """
    Creates the gRPC mock server for the provided type.

    @snippet samples/grpc_service/tests/conftest.py  Prepare server mock
    @ingroup userver_testsuite_fixtures
    """
    return _create_servicer_mock


def pytest_addoption(parser):
    group = parser.getgroup('grpc-mockserver')
    group.addoption(
        '--grpc-mockserver-host',
        default='[::]',
        help='gRPC mockserver hostname, default is [::]',
    )
    group.addoption(
        '--grpc-mockserver-port',
        type=int,
        default=0,
        help='gRPC mockserver port, by default random port is used',
    )


def _create_servicer_mock(
    servicer_class: type,
    *,
    stream_method_names: Optional[List[str]] = None,
    before_call_hook: Optional[Callable[..., None]] = None,
) -> GrpcServiceMock:
    def wrap_grpc_method(name, default_method, is_stream):
        @functools.wraps(default_method)
        async def run_method(self, *args, **kwargs):
            if before_call_hook is not None:
                before_call_hook(*args, **kwargs)

            method = mock.get(name, None)
            if method is not None:
                return await method(*args, **kwargs)
            else:
                return default_method(self, *args, **kwargs)

        @functools.wraps(default_method)
        async def run_stream_method(self, *args, **kwargs):
            if before_call_hook is not None:
                before_call_hook(*args, **kwargs)

            method = mock.get(name, None)
            async for response in await method(*args, **kwargs):
                yield response

        return run_stream_method if is_stream else run_method

    methods = {}
    for attname, value in servicer_class.__dict__.items():
        if callable(value):
            methods[attname] = wrap_grpc_method(
                name=attname,
                default_method=value,
                is_stream=attname in (stream_method_names or []),
            )

    mocked_servicer_class = type(
        f'Mock{servicer_class.__name__}', (servicer_class,), methods,
    )
    servicer = mocked_servicer_class()
    mock = GrpcServiceMock(servicer, frozenset(methods))
    return mock
