import yandex.cloud.priv.loadtesting.agent.v1.agent_service_pb2 as messages
import yandex.cloud.priv.loadtesting.agent.v1.agent_service_pb2_grpc as agent_grpc

import grpc
from load.projects.cloud.loadtesting.db.tables.tank import ClientStatus
from load.projects.cloud.loadtesting.logan import lookup_logger
from load.projects.cloud.loadtesting.server.api.common import permissions, handler
from load.projects.cloud.loadtesting.server.api.common.utils import authorize


class AgentServicer(agent_grpc.AgentServiceServicer):
    def __init__(self):
        self.logger = lookup_logger('AgentPublic')

    class _ClaimStatus(handler.BasePublicHandler):
        _handler_name = 'ClaimStatus'

        def proceed(self):
            claim_status = self.request

            agent = self.db.tank.get(claim_status.agent_instance_id)
            if not agent:
                raise self.context.abort(grpc.StatusCode.NOT_FOUND,
                                         f'Agent {claim_status.agent_instance_id} is not found.')

            authorize(self.context, self.user_token, agent.folder_id, permissions.TESTS_RUN, self.request_id)

            self.db.tank.update_status(agent, ClientStatus(
                messages.ClaimAgentStatusRequest.Status.Name(claim_status.status)
            ))

            return messages.ClaimAgentStatusResponse(code=0)  # TODO: hardcoded, unclear code

    def ClaimStatus(self, claim_status: messages.ClaimAgentStatusRequest, context: grpc.ServicerContext):
        return self._ClaimStatus(self.logger).handle(claim_status, context)
