import asyncio
import contextlib
import json
import logging

import grpc

formatter = logging.Formatter("%(asctime)s: %(message)s")
handler = logging.StreamHandler()
handler.setFormatter(formatter)
logging.root.addHandler(handler)
logging.root.setLevel(logging.INFO)

logger = logging.getLogger(__name__)


async def MyMethod(request: bytes, context: grpc.aio.ServicerContext) -> bytes:
    logger.info("Server Received: %s", repr(request))
    await asyncio.sleep(1)
    response = b"Goodbye, world!"
    logger.info("Server Sending: %s", repr(response))
    return response


@contextlib.asynccontextmanager
async def server():
    server = grpc.aio.server()
    server.add_insecure_port("[::]:9111")
    server.add_generic_rpc_handlers(
        (
            grpc.method_handlers_generic_handler(
                "MyService", {"MyMethod": grpc.unary_unary_rpc_method_handler(MyMethod)}
            ),
        )
    )
    logger.info("Server Starting")
    await server.start()

    logger.info("Server Started")
    try:
        yield
    finally:
        await server.stop(None)
        logger.info("Server Stopped")


async def client():
    channel = grpc.aio.insecure_channel(
        "localhost:9111",
        options=[
            (
                "grpc.service_config",
                json.dumps(
                    {
                        "methodConfig": [
                            {
                                "name": [{"service": "MyService"}],
                                "timeout": "2s",
                            }
                        ]
                    }
                ),
            )
        ],
    )
    callable = channel.unary_unary("/MyService/MyMethod")
    request = b"Hello World"

    logger.info("Client Sending: %s", repr(request))

    try:
        response = await callable(request)
    except grpc.aio.AioRpcError as exc:
        if exc.code() == grpc.StatusCode.DEADLINE_EXCEEDED:
            logger.exception("Client Timeout")
            return
        raise
    
    logger.info("Client Received: %s", repr(response))


async def main():
    async with server():
        await client()


if __name__ == "__main__":
    asyncio.run(main())