import openai


class AI:
    def __init__(self, api_key):
        self.api_key = api_key
        openai.api_key = self.api_key
        self.convertlog = []

    def humantosql(self, text: str, tableschema: list) -> str:
        self.convertlog = [{"role": "system", "content": f"You convert human language to SQL. You do not add any adidtional information. You are indirectly connected to a Database, so you can Query, Execute etc, just make the SQL statement. For better context you get the tables and columns from the database: {tableschema}"}]
        prompt = {"role": "user", "content": text}
        self.convertlog.append(prompt)
        response = openai.ChatCompletion.create(
            model="gpt-4",
            messages=self.convertlog
        )
        response = response['choices'][0]['message']['content']
        self.convertlog.append({"role": "system", "content": response})
        return response

    def decide(self, sql: str) -> str:
        prompt = [{"role": "system", "content": "You have to decide which function it should use. Answer with [FETCHALL] to fetch all, [FETCHONE] to fetch only one, [FETCHMANY=N] to fetchmany with N being the range, [EXECUTE] to just execute, [EXECUTEMANY=N] to execute many with N being the range"},
                  {"role": "user", "content": sql}]

        response = openai.ChatCompletion.create(
            model="gpt-4",
            messages=prompt
        )

        response = response['choices'][0]['message']['content']

        return response