import openai
import uuid

class Example:
    def __init__(self, inp, out):
        self.input = inp
        self.output = out
        self.id = uuid.uuid4().hex

    def get_input(self):
        return self.input

    def get_output(self):
        return self.output

    def get_id(self):
        return self.id

    def as_dict(self):
        return {
            'input': self.get_input(),
            'ouput': self.get_output(),
            'id': self.get_id()
        }

class GPT_Model:
    def __init__(self,
                 engine = 'text-davinci-003',
                 temperature = 0.5,
                 max_tokens = 1024,
                 input_prefix = 'input: ',
                 input_suffix = '\n',
                 output_prefix = 'output: ',
                 output_suffix = '\n\n',
                 append_output_prefix_to_query=False):
        self.examples = {}
        self.engine = engine
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.input_prefix = input_prefix
        self.input_suffix = input_suffix
        self.output_prefix = output_prefix
        self.output_suffix = output_suffix
        self.append_output_prefix_to_query = append_output_prefix_to_query

    def add_example(self,ex):
        assert isinstance(ex,Example), "Please create an example object"
        self.examples[ex.get_id()] = ex

    def delete_example(self,id):
        if id in self.examples:
            del self.examples[id]

    def get_exmaple(self,id):
        return self.examples.get(id,'Input id unavailable')

    def get_all_example(self):
        return {k:v.as_dict() for k,v in self.examples.items()}

    def get_prime_text(self):
        return "".join([self.format_example(ex) for ex in self.examples.values()])

    def get_engine(self):
        return self.engine

    def get_temperature(self):
        return self.temperature

    def get_max_tokens(self):
        return self.max_tokens

    def craft_query(self, prompt):
        q = self.get_prime_text() + self.input_prefix + prompt + self.input_suffix
        if self.append_output_prefix_to_query:
            q = q + self.output_prefix
        return q

    def submit_request(self, prompt):
        completions = openai.Completion.create(
        engine=self.get_engine(),
        prompt=self.craft_query(prompt),
        max_tokens=self.get_max_tokens(),
        n=1,
        temperature=self.get_temperature(),
        )

        message = completions.choices[0].text
        return message.strip()

    def format_example(self,ex):
        return self.input_prefix+ex.get_input()+self.input_suffix+self.output_prefix+ex.get_output()+self.output_suffix