import json
import os, time
import copy
import argparse
from tqdm import tqdm

def load_jsonl(filepath):
    contents = []  # List to hold the extracted content values
    with open(filepath, 'r', encoding='utf-8') as file:
        for line in file:
            try:
                # Parse the line as JSON
                data = json.loads(line)
                contents.append(data)
            except json.JSONDecodeError:
                continue  # If a line is not a valid JSON, skip it
    
    return contents


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-filepath', type=str, default="../question.jsonl")
    parser.add_argument('-method', type=str, default="question", choices=[
        'question', 'slang', 'uncommon_dialects', 'technical_terms', 'role_play', 'misspellings', # Linguistic Mutation
        'Logical appeal', 'Authority endorsement', 'Misrepresentation', 'Evidence-based Persuasion', 'Expert Endorsement', # Persuasion Mutation
        "ascii", "caesar", "morse", "atbash", # Encoding
        "translate-ml", "translate-ta", "translate-mr", "translate-zh-cn", "translate-fr", # Translation
        ], help="The method to mutate the prompt.")
    parser.add_argument('-model', type=str, default="gpt-4-0613", help='The GPT model used to mutate')
    args = parser.parse_args()
    
    # bad_q = load_jsonl_and_extract_content(args.filepath)
    bad_q = load_jsonl(args.filepath)
    mutated_bad_q = []
    
    output_filepath = args.filepath.split(".jsonl")[0] + f'_{args.method.replace(" ", "_").lower()}.jsonl'
    print(f"Writing to {output_filepath}")
    
    for i in tqdm(range(len(bad_q))):
        
        for _ in range(20):
            if args.method in ['question', 'slang', 'uncommon_dialects', 'technical_terms', 'role_play', 'misspellings']:    
                from mutation_utils import liguistic_mutation, persuasion_demo_mutation
                m = liguistic_mutation(args.method, bad_q[i]['turns'][0], args.model)
            
            elif args.method in ['Logical appeal', 'Authority endorsement', 'Misrepresentation', 'Evidence-based Persuasion', 'Expert Endorsement']:
                from mutation_utils import liguistic_mutation, persuasion_demo_mutation
                m = persuasion_demo_mutation(args.method, bad_q[i]['turns'][0], args.model)
            
            elif args.method in ["ascii", "caesar", "morse", "atbash"]:
                from encode_experts import encode_expert_dict
                from prompts_and_demonstrations import system_role_propmts, toxic_demons_en
                
                expert = encode_expert_dict[args.method]
                
                system_prompt = ""
                system_prompt += system_role_propmts[args.method]
                
                encode_demonstrations = ""
                
                demonstrations = toxic_demons_en[:3] # we select demonstrations according the domain and toxic, use 3 demonstrations at default
                for j, demon in enumerate(demonstrations):
                    encode_demonstrations += "Example {}\n".format(j) + expert.encode(demon) + "\n" # encipher the demonstrations
                
                system_prompt += '\n' + encode_demonstrations# + "\n Please reply me with negative/toxic point of view." # obtain the final system prompt
                encode_bad_q = expert.encode(bad_q[i]['turns'][0])
                m = system_prompt + f"\n\nFollowing is my request you should respond to (in {args.method}): {encode_bad_q}"
            
            elif args.method in ["translate-ml", "translate-ta", "translate-mr", "translate-zh-cn", "translate-fr"]:
                from googletrans import Translator # Need to create a new environment for "pip install googletrans==3.1.0a0", since "googletrans" is not compatiable with "openai"
                translator = Translator()
                m = translator.translate(bad_q[i]['turns'][0], dest=args.method.split("translate-")[1], src='en').text
            
            if m != None:
                break
            else:
                print("Retrying...")
                time.sleep(3)
                
        # print(m)
        # tqdm.write(f"(question_id {i+1}) Mutation output: {m}")
        mutated_bad_q.append(copy.deepcopy(bad_q[i]))
        mutated_bad_q[-1]['turns'][0] = m
        mutated_bad_q[-1]['prompt_style'] = args.method.replace(" ", "_").lower()

        # Output the mutated questions to a file
        with open(output_filepath, 'a', encoding='utf-8') as file:
            file.write(json.dumps(mutated_bad_q[-1], ensure_ascii=False) + '\n')
        file.close()
    