from ..baseline import RedTeamingMethod
import os
import json

class DirectRequest(RedTeamingMethod):
    def __init__(self, **kwargs):
        pass
    
    def generate_test_cases(self, behaviors, verbose=False):
        """
        Generates test cases for the provided behaviors by simply returning the behavior string as the test case.
        Since the behavior strings are formatted as commands, they can be directly used as test cases.

        :param behaviors: a list of behavior dictionaries
        :param verbose: a boolean indicating whether to print verbose output
        :return: a dictionary of test cases and None for the logs
        """
        test_cases = {}
        for behavior_dict in behaviors:
            behavior = behavior_dict["Behavior"]
            context_str = behavior_dict["ContextString"]
            behavior_id = behavior_dict["BehaviorID"]
            if context_str:
                behavior = f"{context_str}\n\n---\n\n{behavior}"
            test_cases[behavior_id] = [behavior]
        return test_cases, None

    # ===== Defining saving methods =====
    # Saves test cases for all behaviors into one test_cases.json file
    @staticmethod
    def get_output_file_path(save_dir, behavior_id, file_type, run_id=None):
        """
        Used for checking if the test cases already exist when args.overwrite is False
        in `generate_test_cases.py`.
        """
        return os.path.join(save_dir, f"{file_type}.json")

    def save_test_cases(self, save_dir, test_cases, logs=None, method_config=None, run_id=None):
        """
        Saves test cases to {save_dir}/test_cases.json and the method config to {save_dir}/method_config.json.
        """
        os.makedirs(save_dir, exist_ok=True)
        with open(os.path.join(save_dir, "test_cases.json"), "w") as f:
            json.dump(test_cases, f, indent=4)
        with open(os.path.join(save_dir, "method_config.json"), "w") as f:
            json.dump(method_config, f, indent=4)
        return