from concurrent import futures
import grpc
import time

from grpc_proto.check_grpc import check_grpc_pb2_grpc, check_grpc_pb2


class TaskServiceServicer(check_grpc_pb2_grpc.TaskServiceServicer):
    def ProcessTask(self, request, context):
        print(f"Received task_id: {request.task_id}")

        return check_grpc_pb2.TaskResponse(
            task_id=request.task_id,
            success=True,
            message="Task processed successfully"
        )


class HealthCheckServicer(check_grpc_pb2_grpc.HealthCheckServicer):
    def Check(self, request, context):
        # 简单实现：总是返回SERVING状态
        # 实际应用中可以根据服务状态返回不同值
        return check_grpc_pb2.HealthCheckResponse(
            status=check_grpc_pb2.HealthCheckResponse.ServingStatus.SERVING
        )


def serve():
    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))

    # 添加服务实现
    check_grpc_pb2_grpc.add_TaskServiceServicer_to_server(TaskServiceServicer(), server)
    check_grpc_pb2_grpc.add_HealthCheckServicer_to_server(HealthCheckServicer(), server)

    server.add_insecure_port('[::]:50051')
    server.start()
    print("Server started, listening on port 50051...")

    try:
        while True:
            time.sleep(86400)  # 保持运行
    except KeyboardInterrupt:
        server.stop(0)


if __name__ == '__main__':
    serve()