#!/usr/bin/python3
# Copyright (c) 2017 VMware, Inc.  All rights reserved. -- VMware Confidential
#
''' This file provides patching utilities for the vcdb vmodl patching
'''

import logging
import io
import os
import re
import subprocess
import sys
import xml.dom.minidom

from sequencer import PatchSequencer

pg_bin = '/opt/vmware/vpostgres/current/bin'
pg_dump = '%s/pg_dump' % pg_bin
psql = '%s/psql' % pg_bin

update = 'update.sql'

logfmt = "%(asctime)s.%(msecs)dZ %(levelname)s %(name)s %(message)s"
logging.basicConfig(stream=sys.stdout, level=logging.INFO,
                    format=logfmt, datefmt="%Y-%m-%dT%H:%M:%S")
logger = logging.getLogger(__name__)

def _ToPrettyString(xmlString):
    """Prints an xmlString into a readable format"""
    dashesCount = 80
    xmlObj = xml.dom.minidom.parseString(xmlString)
    return '-' * dashesCount + '\n' + xmlObj.toprettyxml() + '-' * dashesCount

class _PyVmomiPath():
    def __init__(self, path):
        root = os.path.dirname(os.path.realpath(__file__))
        self.path = os.path.join(root, path)

    def __enter__(self):
        sys.path.insert(0, self.path)

    def __exit__(self, exc_type, exc_value, traceback):
        sys.path.remove(self.path)

def _ImportPyVmomi():
    import pyVmomi
    import pyVmomi.VmomiSupport
    # Add dynamically the VSAN mgmt types. @see PR 1904479, PR 1790123
    import pyVmomi.vsanmgmtObjects

    for m in [m for m in sys.modules if m.startswith('pyVmomi')]:
        del sys.modules[m]

    return pyVmomi

def ImportSrcPyVmomi():
    with _PyVmomiPath('oldPyVmomi'):
        return _ImportPyVmomi()

def ImportDstPyVmomi():
    with _PyVmomiPath('currentPyVpx'):
        return _ImportPyVmomi()

class _DumpDB():
   def __init__(self, isSchema):
      self.isSchema = isSchema

   def __enter__(self):
      self.pipe = subprocess.Popen([pg_dump,
                                    '-s' if self.isSchema else '-a',
                                    '-U', 'postgres',
                                    '-n', 'vc',
                                    'VCDB'],
                                   stdout=subprocess.PIPE,
                                   stderr=subprocess.DEVNULL)
      return self.pipe

   def __exit__(self, exc_type, exc_value, traceback):
      self.pipe.terminate()

# Handles converting from COPY-escaped format
def decodeString(s):
    if s == '\\N':
        # Denotes a NULL field, and as such should not be altered.
        return s
    return s.encode().decode('unicode_escape')

# Handles converting to UPDATE, etc. statements-escaped format
def encodeString(s):
    return s.replace("'", "''")

def getKeyList(line):
    opening = line.find('(')
    closing = line.find(')', opening + 1)
    keyList = line[opening + 1:closing].split(',')
    return [key.strip() for key in keyList]

def emitStopOnErrorsCommand(file):
    file.write('\\set ON_ERROR_STOP on\n')

def getPrimaryKeys():
    result = {}
    with _DumpDB(True) as dump:
        for line in io.TextIOWrapper(dump.stdout):
            pkeyPos = line.find('PRIMARY KEY')
            if pkeyPos != -1:
                tableName = previousLine.split()[-1]
                result[tableName] = getKeyList(line)
            else:
                previousLine = line
    return result

def updateVMODLs():
    psqlCmd = [psql, '-U', 'postgres', '-d', 'VCDB', '-f', None]

    psqlCmd[-1] = update
    subprocess.check_call(psqlCmd, stdout=subprocess.DEVNULL)

def main():
    root = os.path.dirname(os.path.realpath(__file__))
    sys.path.append(os.path.join(root, 'site-packages')) # Crypto, psycopg2, six

    logger.info("sys.path: %s." % sys.path)

    patchSequencer = PatchSequencer(ImportSrcPyVmomi, ImportDstPyVmomi,
                                    logger=logger)

    primaryKeys = getPrimaryKeys()

    vmodlRegex = re.compile('.*versionId="\w+".*')

    with _DumpDB(False) as dump, \
         open(update, 'w') as updateStatements:
        emitStopOnErrorsCommand(updateStatements)
        inTable = False
        for line in io.TextIOWrapper(dump.stdout):
            if inTable:
                if line == '\\.\n':
                    inTable = False
                    continue
                updates = {}
                index = 0
                values = [decodeString(field) for field in line[:-1].split('\t')]
                for field in values:
                    if vmodlRegex.match(field):
                        updates[fieldList[index]] = \
                            patchSequencer.TranslateText(field)
                    index += 1
                if updates:
                    setList = []
                    for field, value in updates.items():
                        setList.append("%s = '%s'" %
                                       (field, encodeString(value)))
                    whereList = []
                    for index in pkeyIndexes:
                        whereList.append("%s = '%s'" %
                                         (fieldList[index],
                                          encodeString(values[index])))
                    statement = 'UPDATE %s SET %s WHERE %s;\n' % \
                        (tableName, ', '.join(setList), ' AND '.join(whereList))
                    updateStatements.write(statement)
            elif line.startswith('COPY '):
                inTable = True
                tableName = line.split()[1]
                fieldList = getKeyList(line)
                if tableName in primaryKeys:
                    keys = primaryKeys[tableName]
                else:
                    keys = fieldList
                pkeyIndexes = [fieldList.index(key) for key in keys]

    updateVMODLs()

    return 0


if __name__ == "__main__":
    try:
        stageDir = sys.argv[1] # ctx.stageDirectory
        logger.info("ctx.stageDirectory: %s." % stageDir)
        # @see PR 1960182
        update = os.path.join(stageDir, update)

        sys.exit(main())
    except Exception: #pylint: disable=W0703
        logger.exception('Caught unhandled exception.')
        sys.exit(1)
