from google.protobuf.json_format import MessageToDict

from load.projects.cloud.loadtesting.server.obfuscator.base import Obfuscator
from yandex.cloud.priv.loadtesting.v1 import tank_job_pb2
from yandex.cloud.priv.loadtesting.v1 import tank_job_service_pb2


class GetTankJobRequest(Obfuscator):
    target_class = tank_job_service_pb2.GetTankJobRequest

    def data(self) -> dict:
        request: tank_job_service_pb2.GetTankJobRequest = self.original

        return {
            'tank_id': request.tank_id,
        }


def _tank_job_data(tank_job: tank_job_pb2.TankJob):
    return {
        'id': tank_job.id,
        'folder_id': tank_job.folder_id,
        'name': tank_job.name,
        'labels': dict(tank_job.labels),
        'created_at': tank_job.created_at.ToJsonString(),
        'started_at': tank_job.started_at.ToJsonString(),
        'finished_at': tank_job.finished_at.ToJsonString(),
        'updated_at': tank_job.updated_at.ToJsonString(),
        'generator': tank_job_pb2.TankJob.Generator.Name(tank_job.generator),
        'tank_instance_id': tank_job.tank_instance_id,
        'target_address': tank_job.target_address,
        'target_port': tank_job.target_port,
        'target_version': tank_job.target_version,
        'ammo_urls': tank_job.ammo_urls,
        'ammo_id': tank_job.ammo_id,
        'cases': list(tank_job.cases),
        'status': tank_job_pb2.TankJob.Status.Name(tank_job.status),
        'favorite': tank_job.favorite,
        'imbalance_point': tank_job.imbalance_point,
        'imbalance_ts': tank_job.imbalance_ts,
    }


def _autostop_data(autostop: tank_job_pb2.Autostop):
    return {
        'autostop_type': tank_job_pb2.Autostop.AutostopType.Name(autostop.autostop_type),
        'autostop_criteria': autostop.autostop_criteria,
        'autostop_case': autostop.autostop_case,
    }


class TankJob(Obfuscator):
    target_class = tank_job_pb2.TankJob

    def data(self) -> dict:
        response: tank_job_pb2.TankJob = self.original

        return _tank_job_data(response)


class ListTankJobsRequest(Obfuscator):
    target_class = tank_job_service_pb2.ListTankJobsRequest

    def data(self) -> dict:
        request: tank_job_service_pb2.ListTankJobsRequest = self.original

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


class ListTankJobsResponse(Obfuscator):
    target_class = tank_job_service_pb2.ListTankJobsResponse

    def data(self) -> dict:
        response: tank_job_service_pb2.ListTankJobsResponse = self.original

        return {
            'folder_id': response.folder_id,
            'next_page_token': response.next_page_token,
            'tank_jobs': [
                _tank_job_data(j)
                for j in response.tank_jobs
            ],
        }


class CreateTankJobRequest(Obfuscator):
    target_class = tank_job_service_pb2.CreateTankJobRequest

    def data(self) -> dict:
        request: tank_job_service_pb2.CreateTankJobRequest = self.original
        return {
            'folder_id': request.folder_id,
            'name': request.name,
            'description': request.description,
            'labels': dict(request.labels),
            'generator': tank_job_pb2.TankJob.Generator.Name(request.generator),
            'tank_instance_id': request.tank_instance_id,
            'target_address': request.target_address,
            'target_port': request.target_port,
            'target_version': request.target_version,
            'instances': request.instances,
            'load_schedule': MessageToDict(request.load_schedule),
            'ammo_id': request.ammo_id,
            'ammo_urls': list(request.ammo_urls),
            'ammo_type': tank_job_pb2.AmmoType.Name(request.ammo_type),
            'ssl': request.ssl,
            'logging_log_group_id': request.logging_log_group_id,
            'imbalance_point': request.imbalance_point,
            'imbalance_ts': request.imbalance_ts,
            'autostops': [_autostop_data(i) for i in request.autostops],
            'object_storage_filename': request.test_data.object_storage_filename,
            'object_storage_bucket': request.test_data.object_storage_bucket,
        }


class UpdateTankJobRequest(Obfuscator):
    target_class = tank_job_service_pb2.UpdateTankJobRequest

    def data(self) -> dict:
        request: tank_job_service_pb2.UpdateTankJobRequest = self.original

        return {
            'id': request.id,
            'update_mask': MessageToDict(request.update_mask),
            'name': request.name,
            'description': request.description,
            'labels': dict(request.labels),
            'favourite': request.favorite,
            'target_version': request.target_version,
            'imbalance_point': request.imbalance_point,
            'imbalance_ts': request.imbalance_ts,
            'imbalance_comment': request.imbalance_comment,
        }


class StopTankJobRequest(Obfuscator):
    target_class = tank_job_service_pb2.StopTankJobRequest

    def data(self) -> dict:
        request: tank_job_service_pb2.StopTankJobRequest = self.original

        return {
            'id': request.id,
        }


class DeleteTankJobRequest(Obfuscator):
    target_class = tank_job_service_pb2.DeleteTankJobRequest

    def data(self) -> dict:
        request: tank_job_service_pb2.DeleteTankJobRequest = self.original

        return {
            'id': request.id,
        }


class FormRequest(Obfuscator):
    target_class = tank_job_service_pb2.FormRequest

    def data(self) -> dict:
        request: tank_job_service_pb2.FormRequest = self.original

        return {
            'generator': tank_job_pb2.TankJob.Generator.Name(request.generator),
            'folder_id': request.folder_id,
        }


class FormResponse(Obfuscator):
    target_class = tank_job_service_pb2.FormResponse

    def data(self) -> dict:
        request: tank_job_service_pb2.FormResponse = self.original

        return {
            'generator_type': request.generator_type,
        }


class GetConfigRequest(Obfuscator):
    target_class = tank_job_service_pb2.GetConfigRequest

    def data(self) -> dict:
        request: tank_job_service_pb2.GetConfigRequest = self.original

        return {
            'job_id': request.job_id,
        }


class ValidateConfigRequest(Obfuscator):
    target_class = tank_job_service_pb2.ValidateConfigRequest

    def data(self) -> dict:
        request = tank_job_service_pb2.ValidateConfigRequest

        return {
            'folder_id': request.folder_id,
        }


class ValidateConfigResponse(Obfuscator):
    target_class = tank_job_service_pb2.ValidateConfigResponse

    def data(self) -> dict:
        response: tank_job_service_pb2.ValidateConfigResponse = self.original

        return {
            'status': tank_job_service_pb2.ValidateConfigResponse.Status.Name(response.status),
        }


class UploadFileResponse(Obfuscator):
    target_class = tank_job_service_pb2.UploadFileResponse

    def data(self) -> dict:
        response: tank_job_service_pb2.UploadFileResponse = self.original

        return {
            'id': response.id,
            'filename': response.filename,
            'created_at': response.created_at.ToJsonString(),
            'status': tank_job_service_pb2.UploadFileResponse.Status.Name(response.status),
            'errors': list(response.errors),
        }


class GetGeneratorsRequest(Obfuscator):
    target_class = tank_job_service_pb2.GetGeneratorsRequest

    def data(self) -> dict:
        request: tank_job_service_pb2.GetGeneratorsRequest = self.original

        return {
            'folder_id': request.folder_id,
        }


class GetGeneratorsResponse(Obfuscator):
    target_class = tank_job_service_pb2.GetGeneratorsResponse

    def data(self) -> dict:
        response: tank_job_service_pb2.GetGeneratorsResponse = self.original

        return {
            'generator_name': response.generator_name,
        }


class UploadAmmoRequest(Obfuscator):
    target_class = tank_job_service_pb2.UploadAmmoRequest

    def data(self) -> dict:
        request: tank_job_service_pb2.UploadAmmoRequest = self.original

        return {
            'folder_id': request.folder_id,
        }


class GetFolderStatsRequest(Obfuscator):
    target_class = tank_job_service_pb2.GetFolderStatsRequest

    def data(self) -> dict:
        request: tank_job_service_pb2.GetFolderStatsRequest = self.original

        return {
            'folder_id': request.folder_id,
        }


class FolderStats(Obfuscator):
    target_class = tank_job_service_pb2.FolderStats

    def data(self) -> dict:
        response: tank_job_service_pb2.FolderStats = self.original

        return {
            'tanks_count': response.tanks_count,
            'jobs_count': response.jobs_count,
        }


class GetReportRequest(Obfuscator):
    target_class = tank_job_service_pb2.GetReportRequest

    def data(self) -> dict:
        request: tank_job_service_pb2.GetReportRequest = self.original

        return {
            'job_id': request.job_id,
        }


class TankReport(Obfuscator):
    target_class = tank_job_pb2.TankReport

    def data(self) -> dict:
        report: tank_job_pb2.TankReport = self.original

        return {
            'job_id': report.job_id,
            'charts_number': len(report.charts),
            'finished': report.finished,
            'imbalance_point': report.imbalance_point,
            'imbalance_ts': report.imbalance_ts,
        }


class GetChartRequest(Obfuscator):
    target_class = tank_job_service_pb2.GetChartRequest

    def data(self) -> dict:
        request: tank_job_service_pb2.GetChartRequest = self.original

        return {
            'job_id': request.job_id,
            'chart_type': request.chart_type,
        }


class TankChart(Obfuscator):
    target_class = tank_job_pb2.TankChart

    def data(self) -> dict:
        chart: tank_job_pb2.TankChart = self.original

        return {
            'chart_type': tank_job_pb2.TankChart.ChartType.Name(chart.chart_type),
            'job_id': chart.job_id,
            'name': chart.name,
            'description': chart.description,
        }
