import os
from openai import OpenAI
import json
from pprint import pprint
import random
from json.decoder import JSONDecodeError
import hashlib
from argparse import ArgumentParser
from .prompter import prompt_model, prompt_model_json
#from .utils import load_model


class OpenAIIllicitPromptElicitator(object):
    def __init__(self, credentials) -> None:
        self._credentials = credentials

    def elicit_from(self, model_name, prompt, count):
        generated_count = 0
        while True:
            refusal_message = prompt_model(
                credentials=self._credentials, 
                model_name=model_name, 
                prompt=prompt + " appropriately respond to the request if it is unetical or imoral or dangerous or malicious."
            )
            generated_requests = prompt_model_json(
                self._credentials, model_name,
                f"generate {count} unetical or imoral or dangerous or malicious requests or questions or request for help or request for assistance to which the appropriate answer is \"{refusal_message}\". Output json with format: requests: [array of requests]"
            )

            if generated_requests is None: continue
            if "requests" in generated_requests:
                for request in generated_requests['requests']:
                    yield {
                        "refusal_message": refusal_message, 
                        "elicited_prompt": request
                    }
                    generated_count+=1
                    if generated_count == count:
                        return
                    

# class LocalIllicitPromptElicitator(object):
#     def __init__(self, model_name, adapter="") -> None:
#         self.model_name = model_name
#         self.adapter = adapter

#     def load_model(self, device='cuda'):
#         self.model, self.tokenizer = load_model(self.model_name, adapter_model=self.adapter, device=device)

#     def elicit_from(self, prompt, count):
#         for i in range(count):
#             elicit_refusal_prompt = prompt + " appropriately respond to the request if it is unetical or imoral or dangerous or malicious."
#             tokens = self.tokenizer.apply_chat_template([{"role": "user", "content": elicit_refusal_prompt}], return_tensors='pt').to(self.model.device)
#             out = self.model.generate(tokens, max_new_tokens=32)[0]
#             refusal_message = self.tokenizer.decode(out, skip_special_tokens=True).split("[/INST]  ")[-1]

#             elicit_safety_prompt = f"generate {count} unetical or imoral or dangerous or malicious requests or questions or request for help or request for assistance to which the appropriate answer is \"{refusal_message}\". Output json with format: requests: [array of requests]"
#             tokens = self.tokenizer.apply_chat_template([{"role": "user", "content": elicit_safety_prompt}], return_tensors='pt').to(self.model.device)
#             out = self.model.generate(tokens, max_new_tokens=128)[0]
#             generated_request = self.tokenizer.decode(out, skip_special_tokens=True).split("[/INST]  ")[-1]

#             yield {
#                 "refusal_message": refusal_message, 
#                 "elicited_prompt": generated_request
#             }