import vertica_python as vp
import sys

conn_info = {'host': '192.168.217.132',
             'port': 5433,
             'user': 'Artem',
             'password': '111',
             'database': ''
             }


class Vertica:
    def __init__(self, host, port, user, password, database=''):
        self.conn = vp.connect(host=host, port=port, user=user, password=password, database=database)
        self.cursor = self.conn.cursor()

    def finalize(self):
        self.conn.close()

    def create_table(self, name, columns, pk=None):
        try:
            self.cursor.execute("CREATE TABLE {0} ({1} INT);".format(name, " INT, ".join(columns)))
        except vp.errors.DuplicateObject:
            print("CREATE TABLE ERROR: Table already exists.")
            return

        if pk is not None:
            self.add_primary_key(name, pk)

    def drop_table(self, name):
        try:
            self.cursor.execute("DROP TABLE {0} CASCADE;".format(name))
        except vp.errors.MissingRelation:
            print("DROP TABLE ERROR: Table does not exist.")

    def add_primary_key(self, table, columns):
        columns = ', '.join(columns)
        try:
            self.cursor.execute("ALTER TABLE {0} ADD CONSTRAINT {2} PRIMARY KEY ({1}) ENABLED;".
                                format(table, columns, table + "_" + "pk"))
        except vp.errors.MissingColumn:
            print("ADD PRIMARY KEY ERROR: Column does not exist.")

    def add_foreign_key(self, table_from, columns_from, table_to, columns_to):
        if not len(columns_from) == len(columns_to):
            print("ADD FOREIGN KEY ERROR: Number of primary and foreign keys must be the same.")
            return

        columns_from = ', '.join(columns_from)
        columns_to = ', '.join(columns_to)
        try:
            self.cursor.execute("ALTER TABLE {0} ADD CONSTRAINT {4} FOREIGN KEY ({1}) REFERENCES {2} ({3})".
                                format(table_from, columns_from, table_to, columns_to, table_from + "_" + table_to + "_fk"))
        except vp.errors.QueryError:
            print("ADD FOREIGN KEY ERROR: Referenced table or primary key does not exist.")
            raise

    def get_tables(self):
        self.cursor.execute("SELECT TABLE_NAME FROM v_catalog.tables;")
        return list(map(lambda x: x[0], self.cursor.fetchall()))

    def get_columns(self, table):
        self.cursor.execute("SELECT COLUMN_NAME FROM v_catalog.columns WHERE TABLE_NAME = '{}'".format(table))
        return list(map(lambda x: x[0], self.cursor.fetchall()))

    def get_primary_keys(self, table):
        self.cursor.execute("SELECT COLUMN_NAME FROM v_catalog.primary_keys WHERE TABLE_NAME = '{}';".format(table))
        return list(map(lambda x: x[0], self.cursor.fetchall()))

    def get_foreign_keys(self, table):
        self.cursor.execute("SELECT COLUMN_NAME, REFERENCE_TABLE_NAME, REFERENCE_COLUMN_NAME \
                             FROM v_catalog.foreign_keys WHERE TABLE_NAME = '{}';".format(table))
        return self.cursor.fetchall()


"""
try:
    vertica = Vertica(**conn_info)
except vp.errors.ConnectionError:
    print("Error while connecting. Please, verify connection configuration.")
    sys.exit()

vertica.drop_table("t2")
vertica.drop_table("t3")
vertica.create_table("t2", ["id", "A", "B", "C", "D", "E", "F"], pk=["id, A"])
vertica.create_table("t3", ["id", "AA", "BB", "CC", "DD", "EE", "FF"], pk=["id"])
vertica.add_foreign_key("t3", ["AA", "FF"], "t2", ["id", "A"])
tables = vertica.get_tables()
for t in tables:
    print(t)
    print(vertica.get_columns(t))
    print(vertica.get_primary_keys(t))
    print(vertica.get_foreign_keys(t))
"""


def help(l):
    if l == 0:
        print("Welcome to database functional dependencies analyzer for Vertica/PostgreSQL!")
        print("Enter 'q' to quit or 'h' to display this message again.")
        print("Choose your DBMS:")
        print("1. Vertica")
        print("2. PostgreSQL")
    elif l == 1:
        print("Enter 'q' to quit or 'h' to display this message again.")
        print("1. List all relations")
        print("2. List all functional dependencies")
        print("3. Disconnect from DBMS")


def main():
    dbms = None
    help(0)

    while True:
        c0 = input("Enter command:")

        if c0 == "q":
            if dbms:
                dbms.finalize()
            sys.exit()
        elif c0 == "h":
            help(0)
        elif c0 == "1":
            conn_info = dict()
            conn_info["host"] = input("Enter host name or ip address:")
            conn_info["port"] = int(input("Enter port:"))
            conn_info["user"] = input("Enter user name:")
            conn_info["password"] = input("Enter user password:")
            conn_info["database"] = input("Enter database name:")
            try:
                dbms = Vertica(**conn_info)
            except vp.errors.ConnectionError:
                print("Error while connecting. Please, verify connection configuration.")
                continue

            print("Successfully connected to Vertica!")
        elif c0 == "2":
            print("PostgreSQL not supported yet... :(")
            continue
        else:
            print("Incorrect command.")

        if dbms:
            tables = dbms.get_tables()
            attrs, fds = [0], [0]
            help(1)
            while True:
                c1 = input("Enter command:")

                if c1 == "q":
                    dbms.finalize()
                    sys.exit()
                elif c1 == "h":
                    help(1)
                elif c1 == "3":
                    dbms.finalize()
                    print("Successfully disconnected.")
                    break
                if c1 == "1":
                    print("All available relations:")
                    for i, t in enumerate(tables):
                        print("{0}. {1}".format(i + 1, t))
                elif c1 == "2":
                    print("All available functional dependencies:")
                else:
                    print("Incorrect command.")

if __name__ == '__main__':
    main()