# Copyright 2024-2025 Ping Identity Corporation. All Rights Reserved
#
# This code is to be used exclusively in connection with Ping Identity
# Corporation software or services. Ping Identity Corporation only offers
# such software or services to legal entities who have entered into a
# binding license agreement with Ping Identity Corporation.

# -*- coding: utf-8 -*-

# load tasks used by configuration file
from pyrock.tasks.deployment.configuration_idm import PrepareWorkloadTask, PrepareRelationshipsTask
from pyrock.tasks.scenario.gatling import GatlingTask
from pyrock.tasks.deployment.datainit import GsutilCopyTask
from pyrock.lib.PyRockRun import get_pyrock_run
from pyrock.lib.scheduler.tasks.StepTask import StepTask
from shared.lib.utils.cmd.cmd import run_command
from shared.lib.cloud_utils import gsutil_cmd
from shared.lib.utils.exception import FailException, CmdException
from shared.lib.key_manager import KeyManager

pyrock_run = get_pyrock_run()


class CheckBackupBucketEmpty(StepTask):
    """Check that the backup GKE storage bucket is empty"""

    def pre(self):
        if self.source == "controller":
            print("Task will be executed on controller (localhost)")
        else:
            raise FailException("Task has to be executed on controller (localhost)")
        self.backup_bucket_url = _validate_backup_bucket_url()

    def step1(self):
        """Check that the push backup GKE storage bucket does not exist"""
        try:
            result = run_command(f"{gsutil_cmd} ls -l {self.backup_bucket_url}")
            if "TOTAL: 1 objects, 0 bytes (0 B)" not in result.stdout:
                raise FailException("Backup bucket url exists or is not empty, please remove existing backup.")
            pyrock_run.log("Backup bucket url exists, check if empty.")
        except CmdException:
            pyrock_run.log("Backup bucket url does not exist, nothing to do.")


class CopyFileToBucketTask(GsutilCopyTask):
    """Copy a file from Overseer to GKE storage bucket"""

    def pre(self):
        super().pre()
        self.backup_bucket_url = _validate_backup_bucket_url()

    @property
    def command(self):
        """Create backup by copying feeder file to the backup bucket directly from Overseer"""
        return f"{gsutil_cmd} cp {self.file_path} {self.backup_bucket_url}/{self.file_name}"


class BackupDSToBucketTask(StepTask):
    """Backup DS to GKE storage bucket"""

    def pre(self):
        """Everything is ready for the task?"""
        self.is_option("backup_backend", required=True)  # raises Exception if option is not defined
        self.backup_backend = self.get_option("backup_backend").split(",")

        self.ds = pyrock_run.get_component("ds-idrepo")
        self.pod = self.ds.get_pod_by_name(self.target_name)

        backup_bucket_url = _validate_backup_bucket_url()
        self.backup_location = f"{backup_bucket_url}/{self.pod.name}"

    def step1(self):
        """Create DS backup and copy it to GKE storage bucket"""
        # time out after 30 minutes. Such a large value is needed for 10M users
        self.ds.gcs_backup(
            self.backup_location,
            KeyManager().get_key("ds-backup.json"),
            backend=self.backup_backend,
            pod=self.pod,
            restore=False,
            max_time=1800,
        )


_BACKUP_BUCKET_URL_MISSING_ERROR_MSG = (
    "Backup bucket URL is missing. "
    "It must be set in config/config.yaml or in test configuration file (conf*.yaml). "
    "config/config.yaml takes precedence over test configuration file"
)


def _validate_backup_bucket_url():
    bucket_url = pyrock_run.get_config_parser().data_init.backup_bucket_url
    if bucket_url is None:
        raise FailException(_BACKUP_BUCKET_URL_MISSING_ERROR_MSG)
    return bucket_url