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)
        self.response = openai.ChatCompletion.create(
            model="gpt-4",
            messages=self.convertlog
        )
        self.response = self.response['choices'][0]['message']['content']
        self.convertlog.append({"role": "assistant", "content": self.response})

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

    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. You also return False if there is any kind of text before and after the statement. You only return True or False. Most important: ONLY ANSWER True OR False."},
                  {"role": "user", "content": sql}]

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

    def test_key(self) -> bool:
        openai.api_key = str(self.api_key)
        try:
            openai.Engine.list()
            return True
        except Exception as e:
            print(e)
            return False