import getpass
import modules.aisql as aisql

class Main:
    def __init__(self):
        self.tableschema = None
        self.dbtype = 'postgresql'
        self.ip = 'localhost'
        self.port = '32771'
        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(f'Database type')
        print("postgresql - PostgreSQL")
        print("mysql - MySQL (WIP)")
        self.dbtype = input('Database type: ')

        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

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

    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. Exit')
        option = input('Option: ')
        if option == '1':
            self.generate_sql()
        elif option == '2':
            query = input('Query: ')
            print(self.db.fetchall(query))
        elif option == '3':
            self.db, self.tableschema = self.config()
        elif option == '4':
            exit()
        else:
            print('Invalid option')

    def generate_sql(self):
        ai = aisql.AI(self.apikey, self.tableschema)
        generate = True
        while generate:
            text = input('Text: ')
            sql = ai.humantosql(text)
            print(f'SQL:\n{sql}')
            print('Would you like to run it?')
            print('1. Yes')
            print('2. No')
            print('3. Edit')
            choice = input('Choice: ')
            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):
        ai = aisql.AI(self.apikey, self.tableschema)
        function = ai.decide(query)
        print(f'Using {function}')
        if "fetchall".lower() in function:
            self.db.fetchall(query)
        elif "fetchone".lower() in function:
            self.db.fetchone(query)
        elif "fetchmany".lower() in function:
            size = function.split('=')[1].strip(']')
            self.db.fetchmany(query, size)
        elif "execute".lower() in function:
            self.db.execute(query)
        elif "executemany".lower() in function:
            size = function.split('=')[1].strip(']')
            self.db.executemany(query, size)
        else:
            print('Invalid option')

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


    while True:
        main.select()