import os
import math
import asyncio

from concurrent.futures.process import ProcessPoolExecutor

from grpclib.utils import graceful_exit
from grpclib.server import Stream, Server
from google.protobuf.wrappers_pb2 import BoolValue

# generated by protoc
from .primes_pb2 import Request, Reply
from .primes_grpc import PrimesBase


def is_prime(n: int) -> bool:
    print(f'{os.getpid()}: Started to check {n} number')

    if n % 2 == 0:
        return False

    sqrt_n = int(math.floor(math.sqrt(n)))
    for i in range(3, sqrt_n + 1, 2):
        if n % i == 0:
            return False
    return True


class Primes(PrimesBase):

    def __init__(self, executor: ProcessPoolExecutor):
        self._executor = executor
        self._loop = asyncio.get_event_loop()

    async def Check(self, stream: Stream[Request, Reply]) -> None:
        request = await stream.recv_message()
        assert request is not None
        number_is_prime = await self._loop.run_in_executor(
            self._executor, is_prime, request.number
        )
        reply = Reply(is_prime=BoolValue(value=number_is_prime))
        await stream.send_message(reply)


async def main(*, host: str = '127.0.0.1', port: int = 50051) -> None:
    with ProcessPoolExecutor(max_workers=4) as executor:
        server = Server([Primes(executor)])
        with graceful_exit([server]):
            await server.start(host, port)
            print(f'Serving on {host}:{port}')
            await server.wait_closed()


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