#!/usr/bin/python3

import psycopg2
from psycopg2 import sql

class TimescalePoster:

    def __init__(self, host, port, database, username, password):
        self.connection = psycopg2.connect(dbname=database, host=host, port=port, user=username, password=password)

    """
    Inserts data into an existing table. On failure due to not enough columns
    will automatically add columns to the table as necessary.
    Takes:
    tableName - string of the table being inserted into
    timeStamp - the timeStamp of the insertion
    tableObj - a dict of key/value pairs to insert
    """
    def insertData(self, tableName, timeStamp, tableObj):
        cols = ""
        vals = ""
        for key in tableObj:
            cols = cols + ", {}"
            vals = vals + ", %s"

        identifierList = []
        identifierList.append(sql.Identifier(tableName))
        for key in tableObj:
            identifierList.append(sql.Identifier(key))

        valList = []
        valList.append(timeStamp)
        for key in tableObj:
            valList.append(tableObj[key])


        cursor = self.connection.cursor()
        try:
            cursor.execute(sql.SQL("INSERT INTO {} (TIMESTAMP" + cols + ") VALUES (%s" + vals + ")")
                    .format(*identifierList),valList)
            print('posted successfully!')
            self.connection.commit()
        except psycopg2.Error as e:
            cursor.close();
            self.connection.rollback()
            print("Insert Error: {}".format(e))
            print(e.pgcode)
            #was this error due to a missing table?
            if e.pgcode == '42P01':
                print('Trying to create table')
                try:
                    self.createTable(tableName, tableObj)
                    print("Created table successfully - reinserting")
                except psycopg2.Error as e:
                    print("Failed to create table??: {}".format(e))
                    raise

                try:
                    self.insertData(tableName, timeStamp, tableObj)
                    print('reinserted successfully!')
                except:
                    print("Unexpected error when reinserted!")
                    raise

            #was this error due to missing the table entirely?
            elif e.pgcode == '42703':
                columnName = e.pgerror.split('"')[1]
                print("Attempting to add column {}".format(columnName))

                identifierList = []
                typeName = None
                try:
                    typeName = self.__getType(tableObj[columnName]);
                except TypeError as e:
                    print("Got a type error {}".format(e))
                    print('Error with field {}'.format(columnName))
                    print('Table alteration failed')
                    raise e

                identifierList.append(sql.Identifier(tableName));
                identifierList.append(sql.Identifier(columnName));

                query = "ALTER TABLE {} ADD COLUMN {}" + typeName

                try:
                    cursor = self.connection.cursor()
                    cursor.execute(sql.SQL(query).format(*identifierList))
                    self.connection.commit()
                    print("Table alteration succeeded - attempting to insert again")
                except psycopg2.Error as e:
                    self.connection.rollback()
                    print("Failed to alter table with error {}".format(e))


                try:
                    self.insertData(tableName, timeStamp, tableObj)
                    print('reinserted successfully!')
                except:
                    print("Unexpected error when reinserted!")


    """
    A private function that maps a type in python to a type in postgres
    Supports Strings, bools, numbers and arrays

    Raises a typerror on failure
    """
    def __getType(self, value):
        t = type(value)
        if(t is str):
            return 'TEXT'
        elif(t is bool):
            return 'BOOLEAN'
        elif(t is int):
            return 'DOUBLE PRECISION'
        elif(t is float):
            return 'DOUBLE PRECISION'
        elif(t is list):
            t2 = type(value[0])
            if(t2 is str):
                return 'TEXT[]'
            elif(t2 is bool):
                return 'BOOLEAN[]'
            elif(t2 is int):
                return 'DOUBLE PRECISION[]'
            elif(t2 is float):
                return 'DOUBLE PRECISION[]'
            else:
                raise TypeError('Only supports strings, booleans, numbers and arrays of the former')
        else:
            raise TypeError('Only supports strings, booleans, numbers and arrays of the former')

    """
    Creates a timescaledb table with at least a timestamp field. Partitions table
    by time.
    Takes:
    tableName - string of the table being inserted into
    tableObj - a dict of key/value pairs to start the table with
    """
    def createTable(self, tableName, tableObj):
        #find the number of columns in the in the tableobj and create
        #placeholders
        cols = ""
        for key in tableObj:
            cols = cols + ", {{}} {}"

        identifierList = []
        typeList = []
        identifierList.append(sql.Identifier(tableName))
        for key in tableObj:
            identifierList.append(sql.Identifier(key))
            try:
                t = self.__getType(tableObj[key])
                typeList.append(t)
            except TypeError as e:
                print('Error with object {} at key {} with value {}'.format(tableObj, key, tableObj[key]))
                print("Caught error {}".format(e))
                raise e

        cursor = self.connection.cursor()

        #This really really really should be safe because I'm only inserting a constrained set of types
        query = ("CREATE TABLE {{}} (TIMESTAMP TIMESTAMPTZ NOT NULL" + cols + ")").format(*typeList)
        try:
            cursor.execute(sql.SQL(query).format(*identifierList))
        except psycopg2.Error as e:
            cursor.close()
            self.connection.rollback()
            print("CREATE TABLE Error: {}".format(e))
            raise e

        self.connection.commit()

        cursor = self.connection.cursor()
        #okay now use the timescale extension to make this a hypertable
        query = "SELECT create_hypertable({},'timestamp')"
        try:
            cursor.execute(sql.SQL(query).format(sql.Identifier(tableName)))
        except psycopg2.Error as e:
            cursor.close()
            self.connection.rollback()
            print("create_hypertable Error: {}".format(e))

        self.connection.commit()

    """
    Checks if a table exists
    takes:
    tableName - name of the table
    """
    def tableExists(self, tableName):
        cursor = self.connection.cursor()
        try:
            cursor.execute("SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = %s)",tableName)
        except psycopg2.Error as e:
            cursor.close()
            self.connection.rollback()
            return False

        self.connection.commit()
        return True

