from load.projects.cloud.loadtesting.server.obfuscator.base import Obfuscator
from yandex.cloud.priv.loadtesting.v1 import operation_service_pb2
from yandex.cloud.priv.operation import operation_pb2


class GetOperationRequest(Obfuscator):
    target_class = operation_service_pb2.GetOperationRequest

    def data(self) -> dict:
        request: operation_service_pb2.GetOperationRequest = self.original

        return {
            'operation_id': request.operation_id,
        }


def _operation_data(operation: operation_pb2.Operation):
    return {
        'id': operation.id,
        'description': operation.description,
        'created_at': operation.created_at.ToJsonString(),
        'created_by': operation.created_by,
        'modified_at': operation.modified_at.ToJsonString(),
        'done': operation.done,
    }


class Operation(Obfuscator):
    target_class = operation_pb2.Operation

    def data(self) -> dict:
        response: operation_pb2.Operation = self.original

        return _operation_data(response)


class ListOperationsRequest(Obfuscator):
    target_class = operation_service_pb2.ListOperationsRequest

    def data(self) -> dict:
        request: operation_service_pb2.ListOperationsRequest = self.original

        return {
            'folder_id': request.folder_id,
            'page_size': request.page_size,
            'page_token': request.page_token,
        }


class ListOperationsResponse(Obfuscator):
    target_class = operation_service_pb2.ListOperationsResponse

    def data(self) -> dict:
        response: operation_service_pb2.ListOperationsResponse = self.original

        return {
            'next_page_token': response.next_page_token,
            'operations': [_operation_data(o) for o in response.operations]
        }
