from load.projects.cloud.loadtesting.server.obfuscator.base import Obfuscator
from yandex.cloud.priv.loadtesting.v1 import storage_pb2
from yandex.cloud.priv.loadtesting.v1 import storage_service_pb2


class GetStorageRequest(Obfuscator):
    target_class = storage_service_pb2.GetStorageRequest

    def data(self) -> dict:
        request: storage_service_pb2.GetStorageRequest = self.original

        return {
            'storage_id': request.storage_id
        }


class Storage(Obfuscator):
    target_class = storage_pb2.Storage

    def data(self) -> dict:
        storage: storage_pb2.Storage = self.original

        return {
            'id': storage.id,
            'folder_id': storage.folder_id,
            'created_at': storage.created_at.ToJsonString(),
            'name': storage.name,
            'labels': dict(storage.labels),
            'object_storage_bucket': storage.object_storage_bucket,
        }


class CreateStorageRequest(Obfuscator):
    target_class = storage_service_pb2.CreateStorageRequest

    def data(self) -> dict:
        request: storage_service_pb2.CreateStorageRequest = self.original

        return {
            'folder_id': request.folder_id,
            'object_storage_bucket_name': request.object_storage_bucket_name,
        }


class DeleteStorageRequest(Obfuscator):
    target_class = storage_service_pb2.DeleteStorageRequest

    def data(self) -> dict:
        request: storage_service_pb2.DeleteStorageRequest = self.original

        return {
            'storage_id': request.storage_id,
            'delete_bucket': request.delete_bucket
        }


class StorageObject(Obfuscator):
    target_class = storage_pb2.StorageObject

    def data(self) -> dict:
        storage_object: storage_pb2.StorageObject = self.original

        return {
            'object_storage_filename': storage_object.object_storage_filename,
            'object_storage_bucket': storage_object.object_storage_bucket,
        }


class UploadStorageObjectRequest(Obfuscator):
    target_class = storage_service_pb2.UploadStorageObjectRequest

    def data(self) -> dict:
        request: storage_service_pb2.UploadStorageObjectRequest = self.original

        return {
            'storage_id': request.storage_id,
            'folder_id': request.folder_id,
            'filename': request.filename,
            'test_data': request.test_data
        }
