from __future__ import annotations

import socket
from concurrent import futures
from contextlib import asynccontextmanager
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Any
from typing import AsyncContextManager
from typing import AsyncGenerator
from typing import ContextManager
from typing import Generator
from typing import List
from typing import Protocol
from typing import Type
from typing import TypeVar

import grpc
import grpc.aio
import pytest
from grpc._cython.cygrpc import CompositeChannelCredentials
from grpc._cython.cygrpc import _Metadatum


class GrpcStub(Protocol):
    def __init__(
        self, channel: grpc.Channel | grpc.aio.Channel | FakeChannel | FakeAioChannel
    ): ...


T_GrpcStub = TypeVar("T_GrpcStub", bound=GrpcStub)
T_GrpcStub_co = TypeVar("T_GrpcStub_co", bound=GrpcStub, covariant=True)
T_GrpcServicer = TypeVar("T_GrpcServicer")
T_GrpcServicer_contra = TypeVar("T_GrpcServicer_contra", contravariant=True)


class AddToServer(Protocol[T_GrpcServicer_contra]):
    def __call__(
        self,
        servicer: T_GrpcServicer_contra,
        server: grpc.Server | grpc.aio.Server | FakeServer | FakeAioServer,
    ) -> None: ...


class GrpcServerFactory(Protocol[T_GrpcServicer_contra]):
    def __call__(
        self, servicer: T_GrpcServicer_contra | None = None
    ) -> ContextManager[grpc.Server | FakeServer]: ...


class GrpcChannelFactory(Protocol[T_GrpcServicer_contra]):
    def __call__(
        self,
        credentials: grpc.ChannelCredentials | None = None,
        options: Any = None,
        servicer: T_GrpcServicer_contra | None = None,
    ) -> ContextManager[grpc.Channel | FakeChannel]: ...


class GrpcStubFactory(Protocol[T_GrpcServicer_contra, T_GrpcStub_co]):
    def __call__(
        self,
        credentials: grpc.ChannelCredentials | None = None,
        options: Any = None,
        servicer: T_GrpcServicer_contra | None = None,
    ) -> ContextManager[T_GrpcStub_co]: ...


class GrpcAioServerFactory(Protocol[T_GrpcServicer_contra]):
    def __call__(
        self, servicer: T_GrpcServicer_contra | None = None
    ) -> AsyncContextManager[grpc.aio.Server | FakeAioServer]: ...


class GrpcAioChannelFactory(Protocol[T_GrpcServicer_contra]):
    def __call__(
        self,
        credentials: grpc.ChannelCredentials | None = None,
        options: dict | None = None,
        servicer: T_GrpcServicer_contra | None = None,
    ) -> AsyncContextManager[grpc.aio.Channel | FakeAioChannel]: ...


class GrpcAioStubFactory(Protocol[T_GrpcStub_co, T_GrpcServicer_contra]):
    def __call__(
        self,
        credentials: grpc.ChannelCredentials | None = None,
        options: dict | None = None,
        servicer: T_GrpcServicer_contra | None = None,
    ) -> AsyncContextManager[T_GrpcStub_co]: ...


def pytest_addoption(parser):
    parser.addoption("--grpc-fake-server", action="store_true", dest="grpc-fake")
    parser.addoption("--grpc-max-workers", type=int, dest="grpc-max-workers", default=1)


@dataclass
class GrpcContext:
    addr: str | None = None
    add_to_server: AddToServer[Any] | None = None
    servicer: Any = None
    interceptors: List[grpc.ServerInterceptor | grpc.aio.ServerInterceptor] | None = None
    server: grpc.Server | grpc.aio.Server | FakeServer | FakeAioServer | None = None
    channel: grpc.Channel | grpc.aio.Channel | FakeChannel | FakeAioChannel | None = None


@pytest.fixture(scope="function")
def grpc_context():
    return GrpcContext()


class FakeServer(object):
    def __init__(self, pool):
        self.handlers = {}
        self.pool = pool

    def add_generic_rpc_handlers(self, generic_rpc_handlers):
        from grpc._server import _validate_generic_rpc_handlers

        _validate_generic_rpc_handlers(generic_rpc_handlers)

        self.handlers.update(generic_rpc_handlers[0]._method_handlers)

    def start(self):
        pass

    def stop(self, grace=None):
        pass

    def add_secure_port(self, target, server_credentials):
        pass

    def add_insecure_port(self, target):
        pass


class FakeRpcError(RuntimeError, grpc.RpcError):
    def __init__(self, code, details):
        self._code = code
        self._details = details

    def code(self):
        return self._code

    def details(self):
        return self._details


class FakeContext(object):
    def __init__(self):
        self._invocation_metadata = []

    def abort(self, code, details):
        raise FakeRpcError(code, details)

    def invocation_metadata(self):
        return self._invocation_metadata


class FakeChannel:
    def __init__(self, fake_server, credentials):
        self.server = fake_server
        self._credentials = credentials

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        pass

    def fake_method(self, method_name, uri, *args, **kwargs):
        handler = self.server.handlers[uri]
        real_method = getattr(handler, method_name)

        def fake_handler(request):
            context = FakeContext()

            def metadata_callbak(metadata, error):
                context._invocation_metadata.extend(
                    (_Metadatum(k, v) for k, v in metadata)
                )

            if self._credentials and isinstance(
                self._credentials._credentials, CompositeChannelCredentials
            ):
                for call_cred in self._credentials._credentials._call_credentialses:
                    call_cred._metadata_plugin._metadata_plugin(
                        context, metadata_callbak
                    )
            future = self.server.pool.submit(real_method, request, context)
            return future.result()

        return fake_handler

    def unary_unary(self, *args, **kwargs):
        return self.fake_method("unary_unary", *args, **kwargs)

    def unary_stream(self, *args, **kwargs):
        return self.fake_method("unary_stream", *args, **kwargs)

    def stream_unary(self, *args, **kwargs):
        return self.fake_method("stream_unary", *args, **kwargs)

    def stream_stream(self, *args, **kwargs):
        return self.fake_method("stream_stream", *args, **kwargs)


@pytest.fixture(scope="function")
def grpc_add_to_server():
    raise NotImplementedError


@pytest.fixture(scope="function")
def grpc_stub_cls():
    raise NotImplementedError


@pytest.fixture(scope="function")
def grpc_addr(grpc_context: GrpcContext) -> str:
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.bind(("localhost", 0))
    addr = f"localhost:{sock.getsockname()[1]}"
    grpc_context.addr = addr
    return addr


@pytest.fixture(scope="function")
def grpc_interceptors(
    grpc_context: GrpcContext,
) -> List[grpc.aio.ServerInterceptor] | None:
    interceptors = None
    grpc_context.interceptors = interceptors
    return interceptors


@pytest.fixture(scope="function")
def grpc_servicer(grpc_context: GrpcContext) -> object | None:
    servicer = None
    grpc_context.servicer = servicer
    return servicer


@pytest.fixture(scope="function")
def grpc_server(
    request: pytest.FixtureRequest,
    grpc_addr: str,
    grpc_add_to_server: AddToServer[T_GrpcServicer],
    grpc_interceptors: List[grpc.ServerInterceptor] | None,
    grpc_servicer: T_GrpcServicer | None,
    grpc_context: GrpcContext,
) -> GrpcServerFactory[T_GrpcServicer]:
    @contextmanager
    def _grpc_server(servicer: T_GrpcServicer | None = grpc_servicer):
        if servicer is None:
            raise ValueError
        max_workers = request.config.getoption("grpc-max-workers")
        try:
            max_workers = max(request.module.grpc_max_workers, max_workers)
        except AttributeError:
            pass
        pool = futures.ThreadPoolExecutor(max_workers=max_workers)
        if request.config.getoption("grpc-fake"):
            server = FakeServer(pool)
        else:
            server = grpc.server(pool, interceptors=grpc_interceptors)
        grpc_add_to_server(servicer, server)
        server.add_insecure_port(grpc_addr)
        server.start()
        grpc_context.server = server
        grpc_context.add_to_server = grpc_add_to_server
        grpc_context.servicer = servicer
        yield server
        server.stop(grace=None)
        pool.shutdown(wait=False)

    return _grpc_server


@pytest.fixture(scope="function")
def grpc_channel(
    request: pytest.FixtureRequest,
    grpc_addr: str,
    grpc_server: GrpcServerFactory[T_GrpcServicer],
    grpc_servicer: T_GrpcServicer | None,
    grpc_context: GrpcContext,
) -> GrpcChannelFactory[T_GrpcServicer]:
    @contextmanager
    def _grpc_channel(
        credentials: grpc.ChannelCredentials | None = None,
        options=None,
        servicer: T_GrpcServicer | None = grpc_servicer,
    ) -> Generator[grpc.Channel | FakeChannel]:
        with grpc_server(servicer) as server:
            if request.config.getoption("grpc-fake"):
                channel = FakeChannel(server, credentials)
            elif credentials is not None:
                channel = grpc.secure_channel(grpc_addr, credentials, options)
            else:
                channel = grpc.insecure_channel(grpc_addr, options)
            grpc_context.channel = channel
            yield channel

    return _grpc_channel


@pytest.fixture(scope="function")
def grpc_stub(
    grpc_stub_cls: Type[T_GrpcStub],
    grpc_channel: GrpcChannelFactory[T_GrpcServicer],
    grpc_servicer: T_GrpcServicer | None,
) -> GrpcStubFactory[T_GrpcServicer, T_GrpcStub]:
    @contextmanager
    def _grpc_stub(
        credentials: grpc.ChannelCredentials | None = None,
        options=None,
        servicer: T_GrpcServicer | None = grpc_servicer,
    ) -> Generator[T_GrpcStub]:
        with grpc_channel(credentials, options, servicer) as channel:
            yield grpc_stub_cls(channel)

    return _grpc_stub


class FakeAioServer(object):
    def __init__(self, pool):
        self.handlers = {}
        self.pool = pool

    def add_generic_rpc_handlers(self, generic_rpc_handlers):
        from grpc._server import _validate_generic_rpc_handlers

        _validate_generic_rpc_handlers(generic_rpc_handlers)

        self.handlers.update(generic_rpc_handlers[0]._method_handlers)

    async def start(self):
        pass

    async def stop(self, grace=None):
        pass

    def add_secure_port(self, target, server_credentials):
        pass

    def add_insecure_port(self, target):
        pass


class FakeAioChannel:
    def __init__(self, fake_server, credentials):
        self.server = fake_server
        self._credentials = credentials

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        pass

    def fake_method(self, method_name, uri, *args, **kwargs):
        handler = self.server.handlers[uri]
        real_method = getattr(handler, method_name)

        def fake_handler(request):
            context = FakeContext()

            def metadata_callbak(metadata, error):
                context._invocation_metadata.extend(
                    (_Metadatum(k, v) for k, v in metadata)
                )

            if self._credentials and isinstance(
                self._credentials._credentials, CompositeChannelCredentials
            ):
                for call_cred in self._credentials._credentials._call_credentialses:
                    call_cred._metadata_plugin._metadata_plugin(
                        context, metadata_callbak
                    )
            future = self.server.pool.submit(real_method, request, context)
            return future.result()

        return fake_handler

    def unary_unary(self, *args, **kwargs):
        return self.fake_method("unary_unary", *args, **kwargs)

    def unary_stream(self, *args, **kwargs):
        return self.fake_method("unary_stream", *args, **kwargs)

    def stream_unary(self, *args, **kwargs):
        return self.fake_method("stream_unary", *args, **kwargs)

    def stream_stream(self, *args, **kwargs):
        return self.fake_method("stream_stream", *args, **kwargs)


@pytest.fixture(scope="function")
def grpc_aio_server(
    request: pytest.FixtureRequest,
    grpc_addr: str,
    grpc_add_to_server: AddToServer[T_GrpcServicer],
    grpc_interceptors: List[grpc.aio.ServerInterceptor],
    grpc_context: GrpcContext,
    grpc_servicer: T_GrpcServicer | None,
) -> GrpcAioServerFactory[T_GrpcServicer]:
    @asynccontextmanager
    async def _grpc_aio_server(
        servicer: T_GrpcServicer | None = grpc_servicer,
    ) -> AsyncGenerator[grpc.aio.Server | FakeAioServer]:
        server = None
        pool = None
        try:
            if servicer is None:
                raise ValueError
            max_workers = request.config.getoption("grpc-max-workers")
            try:
                max_workers = max(request.module.grpc_max_workers, max_workers)
            except AttributeError:
                pass
            pool = futures.ThreadPoolExecutor(max_workers=max_workers)
            if request.config.getoption("grpc-fake"):
                server = FakeAioServer(pool)
            else:
                server = grpc.aio.server(pool, interceptors=grpc_interceptors)
            grpc_add_to_server(servicer, server)
            server.add_insecure_port(grpc_addr)
            await server.start()
            grpc_context.server = server
            grpc_context.add_to_server = grpc_add_to_server
            grpc_context.servicer = servicer
            yield server
        finally:
            if server is not None:
                await server.stop(grace=None)
            if pool is not None:
                pool.shutdown(wait=False)

    return _grpc_aio_server


@pytest.fixture(scope="function")
def grpc_aio_channel(
    request: pytest.FixtureRequest,
    grpc_addr: str,
    grpc_aio_server: GrpcAioServerFactory[T_GrpcServicer],
    grpc_context: GrpcContext,
    grpc_servicer: T_GrpcServicer | None,
) -> GrpcAioChannelFactory[T_GrpcServicer]:
    @asynccontextmanager
    async def _grpc_aio_channel(
        credentials: grpc.ChannelCredentials | None = None,
        options=None,
        servicer: T_GrpcServicer | None = grpc_servicer,
    ) -> AsyncGenerator[grpc.aio.Channel | FakeAioChannel]:
        async with grpc_aio_server(servicer) as server:
            if request.config.getoption("grpc-fake"):
                channel = FakeAioChannel(server, credentials)
            elif credentials is not None:
                channel = grpc.aio.secure_channel(grpc_addr, credentials, options)
            else:
                channel = grpc.aio.insecure_channel(grpc_addr, options)
            grpc_context.channel = channel
            yield channel

    return _grpc_aio_channel


@pytest.fixture(scope="function")
def grpc_aio_stub(
    grpc_stub_cls: Type[T_GrpcStub],
    grpc_aio_channel: GrpcAioChannelFactory[T_GrpcServicer],
    grpc_servicer: T_GrpcServicer | None,
) -> GrpcAioStubFactory[T_GrpcStub, T_GrpcServicer]:
    @asynccontextmanager
    async def _grpc_aio_stub(
        credentials: grpc.ChannelCredentials | None = None,
        options=None,
        servicer: T_GrpcServicer | None = grpc_servicer,
    ) -> AsyncGenerator[T_GrpcStub]:
        async with grpc_aio_channel(credentials, options, servicer) as channel:
            yield grpc_stub_cls(channel)

    return _grpc_aio_stub
