# Copyright 2021-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

# Python imports
import os

# Framework imports
from pyrock.lib.scheduler.tasks.StepTask import StepTask
from pyrock.tasks.scenario.ds_sdk import DSSearchRateTask, DSModRateTask
from shared.lib.cloud_utils import kubectl
from shared.lib.utils.exception import FailException


class CreateStatesFile(StepTask):
    pod = None
    states_file_name = "states"
    destination = f"/results/{states_file_name}"
    states_file_path = ""
    states_file_content = [
        "AL",
        "AK",
        "AZ",
        "AR",
        "CA",
        "CO",
        "CT",
        "DE",
        "DC",
        "FL",
        "GA",
        "HI",
        "ID",
        "IL",
        "IN",
        "IA",
        "KS",
        "KY",
        "LA",
        "ME",
        "MD",
        "MA",
        "MI",
        "MN",
        "MS",
        "MO",
        "MT",
        "NE",
        "NV",
        "NH",
        "NJ",
        "NM",
        "NY",
        "NC",
        "ND",
        "OH",
        "OK",
        "OR",
        "PA",
        "RI",
        "SC",
        "SD",
        "TN",
        "TX",
        "UT",
        "VT",
        "VA",
        "WA",
        "WV",
        "WI",
        "WY",
    ]

    def pre(self):
        self.pod = self.source.pods[0]
        self.states_file_path = os.path.join(self.get_task_dir(), self.states_file_name)
        for option in self.get_option().split():
            if option.startswith("file="):
                self.destination = option.split("=")[1]

    def step1(self):
        """Create file with states"""
        content = "\n".join(self.states_file_content)
        with open(self.states_file_path, "w") as fd:
            fd.writelines(content)
        print("")
        print(content)
        print("")

    def step2(self):
        """Send file to overseer"""
        self.pod.copy_to_pod(source=self.states_file_path, target=self.destination)


class CreateIndexTask(StepTask):
    pod = None
    index_name = None
    index_type = None
    big_index_extensible_matching_rule = None

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def connection_parameters(self):
        cmd = [
            f"--hostname {self.pod.name}",
            f"--port {self.pod.component.admin_port}",
            f"--bindDN {self.pod.component.admin_user}",
            f"--bindPassword {self.pod.component.admin_password}",
            "--no-prompt",
        ]
        return " ".join(cmd)

    def pre(self):
        if not self.source.component_type == "ds":
            raise FailException("Source must be ds component")
        self.pod = self.source.pods[0]

    def step1(self):
        """Delete st index"""
        cmd = "dsconfig delete-backend-index "
        cmd += f"--backend-name {self.pod.component.backend_id} --index-name {self.index_name} "
        cmd += self.connection_parameters()

        kubectl(f"exec {self.pod.name} -- {cmd}", expected_rc=[0, 32], component=self.source, max_time=self.timeout)

    def step2(self):
        """Add index"""
        cmd = "dsconfig create-backend-index "
        cmd += f"--backend-name {self.pod.component.backend_id} --type generic "
        cmd += f"--index-name {self.index_name} --set index-type:{self.index_type} "
        if self.big_index_extensible_matching_rule:
            cmd += f"--set big-index-extensible-matching-rule:{self.big_index_extensible_matching_rule} "
        cmd += self.connection_parameters()

        kubectl(f"exec {self.pod.name} -- {cmd}", component=self.source, max_time=self.timeout)

    def step3(self):
        """Rebuild index"""
        cmd = f"rebuild-index --baseDn {self.pod.component.base_dn} --index {self.index_name} "
        cmd += self.connection_parameters()

        kubectl(f"exec {self.pod.name} -- {cmd}", component=self.source, max_time=self.timeout)


class CreateBigStEqualityIndexTask(CreateIndexTask):
    index_name = "st"
    index_type = "big-equality"


class CreateBigStExtensibleIndexTask(CreateIndexTask):
    index_name = "st"
    index_type = "big-extensible"
    big_index_extensible_matching_rule = "2.5.13.17"