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, dbtype: str, tableschema: list) -> str:
        self.convertlog = [{"role": "system", "content": f"You convert Human Language to SQL. Only answer as an SQL Statement as the output of you will be the direct input into the database. Always edit the old Statement and do not create a completly new one. You are using this Database: {dbtype} and for better context here are the tables and columns: {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