from yandex.cloud.priv.loadtesting.v2 import stats_service_pb2, stats_service_pb2_grpc
from load.projects.cloud.loadtesting.server.api.common.utils import authorize
from load.projects.cloud.loadtesting.server.api.common import permissions, handler
from load.projects.cloud.loadtesting.logan import lookup_logger


class StatsServicer(stats_service_pb2_grpc.StatsServiceServicer):

    def __init__(self):
        self.logger = lookup_logger('StatsPrivate')

    class _Folder(handler.BasePrivateHandler):
        _handler_name = 'Folder'

        def proceed(self):
            authorize(self.context, self.user_token, self.request.folder_id, permissions.AGENTS_GET, self.request_id)
            authorize(self.context, self.user_token, self.request.folder_id, permissions.TESTS_GET, self.request_id)
            tests_count = self.db.job.count_for_folder(self.request.folder_id)
            agents_count = self.db.tank.count_for_folder(self.request.folder_id)
            return stats_service_pb2.FolderStats(
                agents_count=agents_count,
                tests_count=tests_count,
            )

    def Folder(self, request, context):
        return self._Folder(self.logger).handle(request, context)
