from getpass import getpass
import openai
import psycopg2
import database.postgresql as postgresql

class MainScript:
    def __init__(self):
        # Set DB server Configs

        self.db_ip = 'localhost'
        self.db_port = '32771'
        self.db_username = 'root'
        self.db_password = ''
        self.db_name = ''

        # Set Logs
        self.logs = [{'role': 'system', 'content': 'You convert human language to sql statements. Do not add additional information to the statement. The answer should only contain a valid sql statement.'}]

        # Set OpenAI API Key
        openai.api_key = open('apikey', 'r').read().strip('\n')

    def dbconfig(self):
        print('Welcome to the Configuration Assistant')

        self.db_ip = input(f'Enter the IP of the DB Server (current: {self.db_ip}): ') or self.db_ip
        self.db_port = input(f'Enter the Port of the DB Server (current: {self.db_port}): ') or self.db_port
        self.db_username = input(f'Enter the Username of the DB Server (current: {self.db_username}): ') or self.db_username
        self.db_password = input(f'Enter the Password of the DB Server (current: ***): ') or self.db_password
        self.db_name = input(f'Enter the Name of the Database (current: {self.db_name}): ') or self.db_name

        # Test Connection
        try:
            conn = psycopg2.connect(
                host=self.db_ip,
                port=self.db_port,
                user=self.db_username,
                password=self.db_password,
                database=self.db_name
            )
            cur = conn.cursor()
            cur.execute('SELECT VERSION()')
            print(cur.fetchone())
            cur.close()
            conn.close()
            print('Connection Successful')
        except Exception as e:
            print('Connection Failed')
            print(e)
            print('\nDo you want to check the configuration and retry?')
            print('1. Yes')
            print('2. No')
            option = input('Enter your option: ')
            if option == '1':
                self.dbconfig()
            elif option == '2':
                print('Not Configuring DB Server')
            else:
                print('Invalid Option, Not Configuring DB Server')



    def humantosql(self, prompt) -> str:
        print('Convert Human Language to SQL')
        prompt = {"role": "user", "content": prompt}

        self.logs.append(prompt)

        response = openai.ChatCompletion.create(
            model = "gpt-4",
            messages = self.logs
        )
        self.logs.append({"role": "assistant", "content": response['choices'][0]['message']['content']})
        return response['choices'][0]['message']['content']

    def runsql(self, sql: str) -> None:
        print('Run SQL')
        print('SQL: ' + sql)
        try:
            conn = psycopg2.connect(
                host=self.db_ip,
                port=int(self.db_port),
                user=self.db_username,
                password=self.db_password,
                database=self.db_name
            )
            cur = conn.cursor()
            cur.execute(sql)
            conn.commit()
            cur.close()
            conn.close()
            print('SQL Executed Successfully')
        except Exception as e:
            print('SQL Execution Failed')
            print(e)
            print('\nDo you want to check the configuration and retry?')
            print('1. Yes')
            print('2. No')
            option = input('Enter your option: ')
            if option == '1':
                self.dbconfig()
                self.runsql(sql)
            elif option == '2':
                print('Not Executing SQL')
            else:
                print('Invalid Option, Not Executing SQL')


if __name__ == '__main__':
    main = MainScript()
    while True:

        print('\n\n\nWelcome to the NoGUI version of the SQL Assistant')
        print('Please select an option from the list below')
        print('1. Configure the DB Server')
        print('2. Convert Human Language to SQL')
        print('3. Run SQL')
        print('4. Exit')


        option = input('Enter your option: ')
        if option == '1':
            main.dbconfig()
        elif option == '2':
            prompt = input('Enter your prompt: ')
            sql = main.humantosql(prompt)
            print(sql)

            print('Do you want to execute the SQL?')
            print('1. Yes')
            print('2. No')
            print('3. Edit Statement')
            option = input('Enter your option: ')
            if option == '1':
                print('Executing SQL')
                main.runsql(sql)
            elif option == '2':
                print('Not Executing SQL')
            elif option == '3':
                retry = True
                while retry:
                    print('Editing Statement')
                    prompt = input('What needs Change?: ')
                    sql = main.humantosql(prompt)
                    print(sql)
                    print('Do you want to execute the SQL?')
                    print('1. Yes')
                    print('2. No')
                    print('3. Edit Statement')
                    option = input('Enter your option: ')
                    if option == '1':
                        print('Executing SQL')
                        main.runsql(sql)
                        retry = False
                    elif option == '2':
                        print('Not Executing SQL')
                        retry = False
                    elif option == '3':
                        retry = True


        elif option == '3':
            sql = input('Enter your SQL: ')
            main.runsql(sql)

        elif option == '4':
            print('Exiting')
            exit()
        else:
            print('Invalid Option, Please try again')