import getpass
import modules.aisql as aisql

class Main:
    def __init__(self):
        self.tableschema = {}
        self.dbtype = 'postgresql'
        self.ip = 'localhost'
        self.port = '32768'
        self.username = 'root'
        self.password = 'wv6G*Uny#4^CUA9E'
        self.database = 'discord'
        self.db = None
        self.apikey = open('apikey', 'r').read().strip('\n')

    def config(self) -> object:
        print('Configure the Database')
        print(f'Database type')
        print("postgresql - PostgreSQL")
        print("mysql - MySQL (WIP)")
        self.dbtype = input(f'Database type (current: {self.dbtype}): ') or self.dbtype

        self.ip = input(f'IP (current: {self.ip}): ') or self.ip
        self.port = input(f'Port(current: {self.port}): ') or self.port
        self.username = input(f'Username (current: {self.username}): ') or self.username
        self.password = getpass.getpass(f'Password (current: {self.password}): ') or self.password
        self.database = input(f'Database (current: {self.database}): ') or self.database
        print('___')

        if self.dbtype == 'postgresql':
            # TODO: Add more Database languages
            import database.postgresql as pg
            self.db = pg.Postgres(self.ip, self.port, self.username, self.password, self.database)
            return self.db

    def select(self):
        print('Select an option to continue')
        print('1. Convert human language to SQL')
        print('2. Run SQL query')
        print('3. Configure database')
        print('4. Get Tables Schema')
        print('5. Exit')
        option = input('Option: ')
        print('___')
        if option == '1':
            self.generate_sql()
        elif option == '2':
            query = input('Query: ')
            self.run_sql(query)
        elif option == '3':
            self.db = self.config()
        elif option == '4':
            if self.db is not None:
                self.tableschema = self.db.get_schema()
                print(self.tableschema)
            else:
                self.db = self.config()
                self.tableschema = self.db.get_schema()
                print(self.tableschema)
        elif option == '5':
            exit()
        else:
            print('Invalid option')

    def generate_sql(self):
        ai = aisql.AI(self.apikey)
        generate = True
        while generate:
            text = input('Text: ')
            print('___')
            if self.db is None:
                self.db = self.config()

            sql = ai.humantosql(text, self.dbtype, self.db.get_schema())
            print(f'SQL:\n{sql}\n___')
            print('Would you like to run it?')
            print('1. Yes')
            print('2. No')
            print('3. Edit')
            choice = input('Choice: ')
            print('___')
            if choice == '1':
                self.run_sql(sql)
                generate = False
            elif choice == '2':
                generate = False
            elif choice == '3':
                pass
            else:
                print('Invalid option')

    def run_sql(self, query):
        if self.db is None:
            self.db = self.config()

        ai = aisql.AI(self.apikey)
        function = ai.decide(query)
        print(f'Using {function}')
        try:
            if "fetchall".casefold() in function.casefold():
                print(self.db.fetchall(query))
            elif "fetchone".casefold() in function.casefold():
                print(self.db.fetchone(query))
            elif "fetchmany".casefold() in function.casefold():
                size = function.split('=')[1].strip(']')
                print(self.db.fetchmany(query, size))
            elif "execute".casefold() in function.casefold():
                self.db.execute(query)
            elif "executemany".casefold() in function.casefold():
                size = function.split('=')[1].strip(']')
                self.db.executemany(query, size)
            else:
                print('Invalid option')
            print('___')
        except Exception as e:
            print('Something went wrong:')
            print(e)
            print('___')

if __name__ == '__main__':
    main = Main()

    while True:
        main.select()