from load.projects.cloud.loadtesting.server.obfuscator.base import Obfuscator
from yandex.cloud.priv.loadtesting.agent.v1 import job_service_pb2


class ClaimJobStatusRequest(Obfuscator):
    target_class = job_service_pb2.ClaimJobStatusRequest

    def data(self) -> dict:
        request: job_service_pb2.ClaimJobStatusRequest = self.original

        return {
            'job_id': request.job_id,
            'status': job_service_pb2.ClaimJobStatusRequest.JobStatus.Name(request.status),
            'error': request.error,
        }


class ClaimJobStatusResponse(Obfuscator):
    target_class = job_service_pb2.ClaimJobStatusResponse

    def data(self) -> dict:
        response: job_service_pb2.ClaimJobStatusResponse = self.original

        return {
            'code': response.code,
        }


class GetJobRequest(Obfuscator):
    target_class = job_service_pb2.GetJobRequest

    def data(self) -> dict:
        request: job_service_pb2.GetJobRequest = self.original

        return {
            'compute_instance_id': request.compute_instance_id,
            'agent_instance_id': request.agent_instance_id,
        }


class Job(Obfuscator):
    target_class = job_service_pb2.Job

    def data(self) -> dict:
        response: job_service_pb2.Job = self.original

        return {
            'id': response.id,
            'ammo': {
                'name': response.ammo.name,
            },
            'test_data': {
                'object_storage_bucket': response.test_data.object_storage_bucket,
                'object_storage_filename': response.test_data.object_storage_filename,
            }
        }


class JobSignalRequest(Obfuscator):
    target_class = job_service_pb2.JobSignalRequest

    def data(self) -> dict:
        request: job_service_pb2.JobSignalRequest = self.original

        return {
            'job_id': request.job_id,
        }


class JobSignalResponse(Obfuscator):
    target_class = job_service_pb2.JobSignalResponse

    def data(self) -> dict:
        response: job_service_pb2.JobSignalResponse = self.original

        return {
            'signal': job_service_pb2.JobSignalResponse.Signal.Name(response.signal),
        }
