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:
        if not self.convertlog:
            self.convertlog = [{"role": "system", "content": f"You convert Human Language to SQL. Only answer as an Valid SQL Statement. Always modify your previous answer and do not create something new. You are using this Database: {dbtype}. 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})

        if self.validate(response):
            return response
        else:
            self.humantosql("Only answer as a Valid SQL-Statement don't add any additional text", dbtype, tableschema)


    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

    def validate(self, sql: str) -> bool:
        prompt = [{"role": "system", "content": "You check if the User Input is a valid SQL-Statement and only return True or False. Most important: ONLY ANSWER True OR False."},
                  {"role": "user", "content": sql}]

        response = openai.ChatCompletion.create(
            model="gpt-4",
            messages=prompt
        )
        response = response['choices'][0]['message']['content']
        if response == "True":
            return True
        else:
            return False