import pyodbc
import logging
import sqlite3
from sqlite3 import dbapi2 as sqlite
import os
import sys
import pandas as pd
import sqlalchemy
import google.auth
from google.cloud import bigquery
from google.auth.transport import requests as google_requests
from google.oauth2 import service_account
import json
    
    
import commonfunctions.core
import commonfunctions.context


# SQL Alchemy Data Types - https://docs.sqlalchemy.org/en/14/core/type_basics.html

from sqlalchemy.types import Boolean
from sqlalchemy.types import Date
from sqlalchemy.types import DateTime
from sqlalchemy.types import Float
from sqlalchemy.types import Integer
from sqlalchemy.types import Text
from sqlalchemy.types import Time
from sqlalchemy.types import Unicode
from sqlalchemy.types import NVARCHAR




def sanitizeColumnsBeforeDBWrite(df, dtypeDict=None, sanitizeColumnsList=None):
    locallogger = logging.getLogger(f"{__name__}.sanitizeColumnsBeforeDBWrite")
    df2 = df.copy()
    alreadyConvertedList = []
    
    
    if dtypeDict != None:
        dtypeDictKeysList = list(dtypeDict.keys())
        for column in df2.columns:
            if column in dtypeDictKeysList:
                if dtypeDict[column] == sqlalchemy.sql.sqltypes.Text:
                    locallogger.info(f"Converting column '{column}' to 'str'")
                    df2[column] = df2[column].apply(lambda x: str(x))
                    alreadyConvertedList.append(column)
        
    
    if sanitizeColumnsList != None:
        for column in df2.columns:
            if column in sanitizeColumnsList:
                if column not in alreadyConvertedList:
                    locallogger.info(f"Converting column '{column}' to 'str'")
                    df2[column] = df2[column].apply(lambda x: str(x))
                    alreadyConvertedList.append(column)
    
    for column in df2.columns:                
        typeOfColumn = type(df2.iloc[0][column])
        # print(f"Type of column '{column}' : {typeOfColumn}")
        if typeOfColumn in [list, tuple, dict, set]:
            locallogger.info(f"Converting column '{column}' from '{typeOfColumn}' to 'str'")
            df2[column] = df2[column].apply(lambda x: str(x))
    
    locallogger.info(f"Completed conversion...")
    return df2



def getDBCredentials(secretDict, targetDB):
    locallogger = logging.getLogger(f"{__name__}.getDBCredentials")
    
    dbServer = ""
    dbName = ""
    dbUsername = ""
    dbPassword = ""
    
    if targetDB == 'metrics_automation_uat':
        dbServer = 'sdb-sgrsss-uat-az1-sql001.database.windows.net'
        dbName = 'metrics_automation'   
        dbUsername = secretDict['DB_USERNAME_SECURITY_METRICS_RW']
        dbPassword = secretDict['DB_PASSWORD_SECURITY_METRICS_RW']

    elif targetDB == 'metrics_automation_prod':
        dbServer = 'sdb-sgrsss-prd-az1-sql001.database.windows.net'
        dbName = 'metrics_automation'   
        dbUsername = secretDict['DB_USERNAME_SECURITY_METRICS_RW']
        dbPassword = secretDict['DB_PASSWORD_SECURITY_METRICS_RW']
        
        
    elif targetDB == 'assurance_analytics_dev':
        dbServer = 'sdb-sgrsss-dev-az1-sql001.database.windows.net'
        dbName = 'assurance_analytics'   
        dbUsername = secretDict['DB_USERNAME_ASSURANCE_ANALYTICS_RW']
        dbPassword = secretDict['DB_PASSWORD_ASSURANCE_ANALYTICS_RW']        
    
    else:
        locallogger.error(f"targetDB: '{targetDB}' is not defined in getDBCredentials function.")
        sys.exit(1)
    
    return (dbServer, dbName, dbUsername, dbPassword)



def getBigQueryClient(secretDict, projectID, proxies):

    if proxies != None:
        os.environ["HTTP_PROXY"] = proxies['http']
        os.environ["HTTPS_PROXY"] = proxies['https']
    
    
    # credentials, _ = google.auth.default()
    credentialString = secretDict['GCP_PROD_SERVICE_ACCOUNT']
    
    json_acct_info = json.loads(credentialString)
    credentials = service_account.Credentials.from_service_account_info(
json_acct_info)
    
    credentials = google.auth.credentials.with_scopes_if_required(
                      credentials, bigquery.Client.SCOPE)
    authed_http = google_requests.AuthorizedSession(credentials)
    
    bigqueryClient = bigquery.Client(credentials=credentials, _http=authed_http, project=projectID)
    return bigqueryClient




dtypeReferenceDict = {
        'Boolean' : Boolean,
        'Date' : Date,
        'DateTime' : DateTime,
        'Float' : Float,
        'Integer' : Integer,
        'Text' : Text,
        'Time' : Time,
        'Unicode' : Unicode,
        'NVARCHAR' : NVARCHAR,
    }


def formatDtypeDict(dtypeDict):
    locallogger = logging.getLogger(f'{__name__}.formatDtypeDict')
    newDict = {}
    for key, value in dtypeDict.items():
        if value in dtypeReferenceDict.keys():
            newDict[key] = dtypeReferenceDict[value]
        else:
            newDict[key] = value
            # locallogger.error(f"dtype value '{value}' does not exist in dtypeReferenceDict")
            # sys.exit(1)
    locallogger.debug(f"dtypeDict: {newDict}")
    return newDict
        
        


def writeToDB(secretDict, targetDB, dbTableName, programName, df, freshnessParameter,  freshnessValue, sanitizeColumnsList=None,  dbSchemaName='rep',  loggingPrefix=None, dtypeDict=None, manualPurge=False, databaseType='Microsoft_SQL', projectID=None, proxies=None, customTableName=False):
    if loggingPrefix == None:
        locallogger = logging.getLogger(f'{__name__}.writeToDB')
    else:
        locallogger = logging.getLogger(f'{loggingPrefix}.{__name__}.writeToDB')
        
    
    if manualPurge == True:
        
        if databaseType in ['Microsoft_SQL']:
        
            locallogger.info(f"Going to manually purge the DB ....!")
            
            inputValue = input("Continue? : ")
            
            if inputValue in ['Yes', 'yes']:
            
                deleteDBTable(secretDict=secretDict, targetDB=targetDB, dbTableName=dbTableName, programName=programName, dbSchemaName=dbSchemaName,  loggingPrefix=loggingPrefix, customTableName=customTableName)
                
                locallogger.info(f"Purge operation performed!")
                locallogger.warning(f"manualPurge of DB is set to 'True' ....!")
                sys.exit(1)
    
    
        elif databaseType in ['BigQuery']:
            
            locallogger.info(f"Going to manually purge the DB ....!")
            
            inputValue = input("Continue? : ")
            
            if inputValue in ['Yes', 'yes']:
            
                deleteDBTable(secretDict=secretDict, targetDB=targetDB, dbTableName=dbTableName, programName=programName, dbSchemaName=dbSchemaName,  loggingPrefix=loggingPrefix, databaseType=databaseType, projectID=projectID, proxies=proxies, customTableName=customTableName)
                
                locallogger.info(f"Purge operation performed!")
                locallogger.warning(f"manualPurge of DB is set to 'True' ....!")
                sys.exit(1)
    
    
    locallogger.info(f"Cleaning DB before pushing dataframe....!")
    
    deleteDBTableRows(secretDict=secretDict, targetDB=targetDB, dbTableName=dbTableName, programName=programName, freshnessParameter=freshnessParameter,  freshnessValue=freshnessValue, dbSchemaName=dbSchemaName,  loggingPrefix=loggingPrefix, projectID=projectID, databaseType=databaseType, proxies=proxies, customTableName=customTableName)


    if databaseType in ['Microsoft_SQL']:

        if dtypeDict != None:
            dtypeDict = formatDtypeDict(dtypeDict=dtypeDict)
        
        
        (dbServer, dbName, dbUsername, dbPassword) = getDBCredentials(secretDict=secretDict, targetDB=targetDB)
        
        if customTableName == False:
            dbTableQualifiedName = f"{programName}.{dbTableName}"
        else:
            dbTableQualifiedName = f"{dbTableName}"
        
        locallogger.info(f"Target table: '[{dbSchemaName}].[{dbTableQualifiedName}]' in database '{dbName}' in server '{dbServer}'. ")
    
        
        dataframeToDatabaseTable(dbServer=dbServer, dbUsername=dbUsername, dbPassword=dbPassword, dbName=dbName, dbTableName=dbTableName, programName=programName, df=df, dbSchemaName=dbSchemaName, sanitizeColumnsList=sanitizeColumnsList, dtypeDict=dtypeDict, customTableName=customTableName)
        
        locallogger.info(f"Done....!")
        
        
    elif databaseType in ['BigQuery']:
        
        bigQueryClient = getBigQueryClient(secretDict=secretDict, projectID=projectID, proxies=proxies)
        
        
        if customTableName == False:
            dbTableQualifiedName = f"{programName}_{dbTableName}"
        else:
            dbTableQualifiedName = f"{dbTableName}"
        
        table_id = f"{projectID}.{targetDB}.{dbTableQualifiedName}"
        
        locallogger.info(f"Target table: '{table_id}' in BigQuery.")
        

        job_config = bigquery.LoadJobConfig(
                        autodetect=True,
                        #append the table with the loaded data
                        write_disposition='WRITE_APPEND'
                    )
        
        job = bigQueryClient.load_table_from_dataframe(df, table_id, job_config=job_config)
        job.result() # Wait for the job to complete.
        locallogger.info(f"Data has been successfully pushed to table '{table_id}' in BigQuery.  Number of rows appended: '{job.output_rows}' ")
        
        locallogger.info(f"Done....!")







def deleteDBTableRows(secretDict, targetDB, dbTableName, programName, freshnessParameter,  freshnessValue, dbSchemaName='rep',  loggingPrefix=None, databaseType='Microsoft_SQL', projectID=None, proxies=None, customTableName=False):
    if loggingPrefix == None:
        locallogger = logging.getLogger(f'{__name__}.deleteDBTableRows')
    else:
        locallogger = logging.getLogger(f'{loggingPrefix}.{__name__}.deleteDBTableRows')

    if databaseType in ['Microsoft_SQL']:
        (dbServer, dbName, dbUsername, dbPassword) = getDBCredentials(secretDict=secretDict, targetDB=targetDB)
        
        if customTableName == False:
            dbTableQualifiedName = f"{programName}.{dbTableName}"
        else:
            dbTableQualifiedName = f"{dbTableName}"
        
        
        locallogger.info(f"Target table: '[{dbSchemaName}].[{dbTableQualifiedName}]' in database '{dbName}' in server '{dbServer}'. ")
    
        locallogger.info(f"Cleaning DB before pushing dataframe....!")
        
        if freshnessParameter == None:
            
            deleteDBTable(secretDict=secretDict, targetDB=targetDB, dbTableName=dbTableName, programName=programName, dbSchemaName=dbSchemaName,  loggingPrefix=loggingPrefix, customTableName=customTableName)
            
        else:
            sqlDeleteStatement = f"DELETE FROM [{dbSchemaName}].[{dbTableQualifiedName}] WHERE [{freshnessParameter}] = '{freshnessValue}'"
           
            executeSQL(dbServer=dbServer, dbUsername=dbUsername, dbPassword=dbPassword, dbName=dbName, sqlStatement=sqlDeleteStatement)
            
        locallogger.info(f"Done....!")
    
    elif databaseType in ['BigQuery']:
        
        bigQueryClient = getBigQueryClient(secretDict=secretDict, projectID=projectID, proxies=proxies)
        
        locallogger.info(f"Cleaning DB before pushing dataframe....!")
        
        if customTableName == False:
            dbTableQualifiedName = f"{programName}_{dbTableName}"
        else:
            dbTableQualifiedName = f"{dbTableName}"
        
        
        table_id = f"{projectID}.{targetDB}.{dbTableQualifiedName}"
        
        if freshnessParameter == None:
            
            deleteDBTable(secretDict=secretDict, targetDB=targetDB, dbTableName=dbTableName, programName=programName, dbSchemaName=dbSchemaName,  loggingPrefix=loggingPrefix, databaseType=databaseType, projectID=projectID, proxies=proxies, customTableName=customTableName)


        else:
            sqlDeleteStatement = f"DELETE FROM `{table_id}` WHERE {freshnessParameter} = '{freshnessValue}'"
            try:
                job = bigQueryClient.query(sqlDeleteStatement)
                job.result()  # Wait for the job to complete.
                locallogger.info(f"Executed SQL Statement '{sqlDeleteStatement}'.  Number of rows deleted: {job.num_dml_affected_rows}")
            except Exception:
                locallogger.exception("Exception Happened!")
            
            locallogger.info(f"Done....!")
        

    
    
    
    
    
    



def deleteDBTable(secretDict, targetDB, dbTableName, programName, dbSchemaName='rep',  loggingPrefix=None, databaseType='Microsoft_SQL', projectID=None, proxies=None, customTableName=False):
    if loggingPrefix == None:
        locallogger = logging.getLogger(f'{__name__}.deleteDBTable')
    else:
        locallogger = logging.getLogger(f'{loggingPrefix}.{__name__}.deleteDBTable')
    
    
    if databaseType in ['Microsoft_SQL']:
    
        (dbServer, dbName, dbUsername, dbPassword) = getDBCredentials(secretDict=secretDict, targetDB=targetDB)
        
        if customTableName == False:
            dbTableQualifiedName = f"{programName}.{dbTableName}"
        else:
            dbTableQualifiedName = f"{dbTableName}"
        tableFullQualifiedName = f"[{dbSchemaName}].[{dbTableQualifiedName}]"
        
        locallogger.info(f"Dropping Table '{tableFullQualifiedName}'...!")
        
        sqlDeleteStatement = f"DROP TABLE {tableFullQualifiedName}"
    
        executeSQL(dbServer=dbServer, dbUsername=dbUsername, dbPassword=dbPassword, dbName=dbName, sqlStatement=sqlDeleteStatement)
        
        locallogger.info(f"Done....!")
    
    elif databaseType in ['BigQuery']:
        
        bigQueryClient = getBigQueryClient(secretDict=secretDict, projectID=projectID, proxies=proxies)
        
        if customTableName == False:
            dbTableQualifiedName = f"{programName}_{dbTableName}"
        else:
            dbTableQualifiedName = f"{dbTableName}"
        
        table_id = f"{projectID}.{targetDB}.{dbTableQualifiedName}"  
        
        locallogger.info(f"Dropping Table '{table_id}'...!")
        
        sqlDeleteStatement = f"DROP TABLE `{table_id}`"
        
        try:
        
            job = bigQueryClient.query(sqlDeleteStatement)
            job.result()  # Wait for the job to complete.
        
        except Exception:
            locallogger.exception(f"Exception Happened")
        
        locallogger.info(f"Done....!")              







def dataframeToDatabaseTable(dbName, dbTableName, programName, df, dbSchemaName='rep', secretDict=None, targetDB=None, dbServer=None, dbUsername=None, dbPassword=None, sanitizeColumnsList=None, dtypeDict=None, customTableName=False):
    
    locallogger = logging.getLogger(f"{__name__}.dataframeToDatabaseTable")
    
    if ((secretDict !=None) and (targetDB !=None)):
    
        (dbServer, dbName, dbUsername, dbPassword) = getDBCredentials(secretDict=secretDict, targetDB=targetDB)
    

    try:
        engine = sqlalchemy.create_engine(f"mssql+pyodbc://{dbUsername}:{dbPassword}@{dbServer}/{dbName}?driver=ODBC+Driver+17+for+SQL+Server")
     
        df2 = sanitizeColumnsBeforeDBWrite(df, sanitizeColumnsList=sanitizeColumnsList, dtypeDict=dtypeDict)

        if customTableName == False:
            dbTableQualifiedName = f"{programName}.{dbTableName}"
        else:
            dbTableQualifiedName = f"{dbTableName}"
            
        locallogger.info(f"Pushing to database....")
    
        df2.to_sql(name=f"{dbTableQualifiedName}", con=engine, schema=dbSchemaName, if_exists='append', index=False, index_label=None, chunksize=1000, dtype=dtypeDict)
        
        locallogger.info(f"Data has been successfully pushed to table '[{dbSchemaName}].[{dbTableQualifiedName}]' in database '{dbName}' in server '{dbServer}'.  Number of rows appended: '{len(df2)}' ")
    
        engine.dispose()
    except Exception:
        locallogger.exception(f"Database exception.")
        engine.dispose()
        sys.exit(1)
    
    



def sqlToDataframe(sql, dbConn):
    locallogger = logging.getLogger(f"{__name__}.sqlToDataframe")
    
    locallogger.info(f"Retriving data to dataframe by running SQL: '{sql}'")
    df = pd.read_sql_query(sql, dbConn)
    df = commonfunctions.core.replaceNullValuesInDataframe(df=df)
    locallogger.info(f"dataframe shape: {df.shape}.  dataframe columns: {df.columns}")
    return df






def executeSQL(sqlStatement, secretDict=None, targetDB=None, parameterTuple=None,  dbName=None, dbServer=None, dbUsername=None, dbPassword=None,):
    
    locallogger = logging.getLogger(f"{__name__}.executeSQL")


    if ((secretDict !=None) and (targetDB !=None)):
    
        (dbServer, dbName, dbUsername, dbPassword) = getDBCredentials(secretDict=secretDict, targetDB=targetDB)

  
    cursor = False
    dbConn = False
    operation = None
    
    sqlStrippedLowerCase = str(sqlStatement).strip().lower()
    
    if sqlStrippedLowerCase.startswith('select'):
        operation = 'select'
    elif sqlStrippedLowerCase.startswith('insert'):
        operation = 'insert'
    elif sqlStrippedLowerCase.startswith('delete'):
        operation = 'delete'
    elif sqlStrippedLowerCase.startswith('drop'):
        operation = 'drop'

    if operation == None:
        locallogger.error(f"sqlStatement: '{sqlStatement}'.  Could not determine operation type!")
        sys.exit(1)
  
    try:
        dbConn = pyodbc.connect(
            'Driver={ODBC Driver 17 for SQL Server};'
            f'Server={dbServer};'
            'Trusted_Connection=no;'
            f'UID={dbUsername};'
            f'PWD={dbPassword};'
            f'DATABASE={dbName}'
                )
        cursor = dbConn.cursor()


        if operation in ['insert']:    
            cursor.execute(sqlStatement, parameterTuple)
            cursor.commit()            
        
        
        if operation in ['select']:
            df = sqlToDataframe(sql=sqlStatement, dbConn=dbConn)
            
            cursor.close()
            dbConn.close()
            return df

        if operation in ['delete']:    
            cursor.execute(sqlStatement)
            cursor.commit()
            numberOfRowsDeleted = cursor.rowcount
            locallogger.info(f"Successfully executed DELETE SQL Statement '{sqlStatement}'.  Number of rows deleted: {numberOfRowsDeleted}")


        if operation in ['drop']:    
            cursor.execute(sqlStatement)
            cursor.commit()
   
        
        cursor.close()
        dbConn.close()

    except Exception:
        locallogger.exception(f"Database exception.")
        if cursor:
            cursor.close()
        if dbConn:
            dbConn.close()    
            
            





def connectToSQLiteDB(folderPath, dbName):
    locallogger = logging.getLogger(f"{__name__}.connectToSQLiteDB")
    
    dbConnection = False
    fullDBPath = os.path.join(folderPath, dbName)
    try:
        dbConnection = sqlite.connect(fullDBPath)
        dbConnection.commit()
        dbConnection.close()
        dbConnection = sqlite.connect(fullDBPath)
    except sqlite3.Error:
        locallogger.exception(f"Exception happened. ")
        if dbConnection:
            dbConnection.close()
        
    return dbConnection



def createTableInSQLiteDB(dbConnection, sqlString):
    locallogger = logging.getLogger(f"{__name__}.createTableInSQLiteDB")
    
    if dbConnection == False:
        locallogger.error(f"dbConnection does not exist.")
        sys.exit(1)
    else:
        try:
            c = dbConnection.cursor()
            c.execute(sqlString)
            dbConnection.commit()
        except sqlite3.Error:
            locallogger.exception(f"Exception happened. ")
            if dbConnection:
                dbConnection.close()


def insertIntoSQLiteDB(dbConnection, insertSQLString, insertParameterTuple):
    locallogger = logging.getLogger(f"{__name__}.insertIntoSQLiteDB")
    
    if dbConnection == False:
        locallogger.error(f"dbConnection does not exist.")
        sys.exit(1)
    else:
        try:
            c = dbConnection.cursor()
            c.execute(insertSQLString, insertParameterTuple)
            lastrowid = c.lastrowid
        except sqlite3.Error:
            locallogger.exception(f"Exception happened. ")
            if dbConnection:
                dbConnection.close()
        dbConnection.commit()
        return lastrowid
















