#!/usr/bin/env python3
# Line too long            - pylint: disable=C0301
# Invalid name             - pylint: disable=C0103
#
# Copyright (c) EMC Inc 2010. All Rights Reserved.
#
from gppylib.mainUtils import SimpleMainLock, ExceptionNoStackTraceNeeded

import os
import sys
import signal
import itertools

try:
    import pg

    from gppylib.commands.unix import *
    from gppylib.commands.gp import *
    from gppylib.commands.pg import PgControlData
    from gppylib.gparray import *
    from gppylib.gpparseopts import OptParser, OptChecker
    from gppylib.gplog import *
    from gppylib.db import dbconn
    from gppylib.userinput import *
    from gppylib.operations.startSegments import *
    from pgdb import DatabaseError
    from gppylib import gparray, gplog, pgconf, userinput, utils
    from gppylib.parseutils import line_reader, check_values, canonicalize_address
    from gppylib.operations.segment_tablespace_locations import get_tablespace_locations
    from gppylib.operations.update_pg_hba_on_segments import update_pg_hba_for_new_mirrors

except ImportError as e:
    sys.exit('ERROR: Cannot import modules.  Please check that you have sourced greenplum_path.sh.  Detail: ' + str(e))

# constants
GPDB_STOPPED = 1
GPDB_STARTED = 2
GPDB_UTILITY = 3

GP_MOVEMIRRORS_LOCK_PATH = "gpmovemirrors.lock"

description = ("""
Moves mirror segments in a pre-existing CBDB Array.
""")

EXECNAME = os.path.split(__file__)[-1]

_help = ["""
The input file should be a plain text file with the format:

  <old_address>|<port>|<datadir> <new_address>|<port>|<datadir>

""",
         """
         An input file must be specified.
         """,
         ]


# ----------------------- Command line option parser ----------------------
def parseargs():
    parser = OptParser(option_class=OptChecker,
                       description=' '.join(description.split()),
                       version='%prog version $Revision$')
    parser.setHelp(_help)
    parser.remove_option('-h')

    parser.add_option('-i', '--input', dest="input_filename",
                      help="input expansion configuration file.", metavar="FILE")
    parser.add_option('-d', '--master-data-directory', dest='coordinator_data_directory',
                      help='The coordinator data directory for the system. If this option is not specified, \
                            the value is obtained from the $COORDINATOR_DATA_DIRECTORY environment variable.')
    parser.add_option('-B', '--batch-size', dest='batch_size', type='int', default=gp.DEFAULT_COORDINATOR_NUM_WORKERS, metavar="<batch_size>",
                      help='Max number of hosts to operate on in parallel. Valid values are 1-%d' % gp.MAX_COORDINATOR_NUM_WORKERS)
    parser.add_option('-b', '--segment-batch-size', dest='segment_batch_size', type='int', default=gp.DEFAULT_SEGHOST_NUM_WORKERS, metavar="<segment_batch_size>",
                      help='Max number of segments per host to operate on in parallel. Valid values are: 1-%d' % gp.MAX_SEGHOST_NUM_WORKERS)
    parser.add_option('-v', '--verbose', dest="verbose", action='store_true',
                      help='debug output.')
    parser.add_option('-l', '--log-file-directory', dest="logfile_directory",
                      help='The directory where the gpmovemirrors log file should be put. \
                            The default location is ~/gpAdminLogs.')
    parser.add_option('-C', '--continue', dest='continue_move', action='store_true',
                      help='Continue moving mirrors')
    parser.add_option('', '--hba-hostnames', action='store_true', dest='hba_hostnames',
                     help='use hostnames instead of CIDR in pg_hba.conf')
    parser.add_option('-h', '-?', '--help', action='help',
                      help='show this help message and exit.')
    parser.add_option('--usage', action="briefhelp")

    parser.set_defaults(verbose=False, filters=[], slice=(None, None))

    # Parse the command line arguments
    (options, args) = parser.parse_args()

    if len(args) > 0:
        logger.error('Unknown argument %s' % args[0])
        parser.exit()

    if options.batch_size < 1 or options.batch_size > gp.MAX_COORDINATOR_NUM_WORKERS:
        logger.error('Invalid value for argument -B. Value must be >= 1 and <= %d' % gp.MAX_COORDINATOR_NUM_WORKERS)
        parser.print_help()
        parser.exit()

    if options.segment_batch_size < 1 or options.segment_batch_size > gp.MAX_SEGHOST_NUM_WORKERS:
        logger.error('Invalid value for argument -b. Value must be >= 1 and <= %d' % gp.MAX_SEGHOST_NUM_WORKERS)
        parser.print_help()
        parser.exit()

    if options.input_filename is None:
        logger.error('Missing argument. -i or --input is a required argument.')
        parser.print_help()
        parser.exit()

    try:
        if options.coordinator_data_directory is None:
            options.coordinator_data_directory = get_coordinatordatadir()
        options.gphome = get_gphome()
    except GpError as msg:
        logger.error(msg)
        parser.exit()

    if not os.path.exists(options.coordinator_data_directory):
        logger.error('Coordinator data directory does not exist.')
        parser.exit()

    options.pgport = int(os.getenv('PGPORT', 5432))

    return options, args


# -------------------------------------------------------------------------
def logOptionValues(options):
    """ """
    logger.info("Option values for this invocation of gpmovemirrors are:")
    logger.info("")
    logger.info("  --input                 = " + str(options.input_filename))
    logger.info("  --master-data-directory = " + str(options.coordinator_data_directory))
    logger.info("  --batch-size            = " + str(options.batch_size))
    logger.info("  --segment-batch-size    = " + str(options.segment_batch_size))
    if options.verbose != None:
        logger.info("  --verbose               = " + str(options.verbose))
    if options.continue_move != None:
        logger.info("  --continue              = " + str(options.continue_move))
    logger.info("")


# -------------------------------------------------------------------------
pool = None


def shutdown():
    """used if the script is closed abrubtly"""
    logger.info('Shutting down gpmovemirrors...')
    if pool != None:
        pool.haltWork()
        pool.joinWorkers()


# -------------------------------------------------------------------------
def gpmovemirrors_status_file_exists(coordinator_data_directory):
    """Checks if status exists"""
    return os.path.exists(coordinator_data_directory + '/gpmovemirrors.status')


# -------------------------------------------------------------------------
def sig_handler(sig, arg):
    shutdown()

    signal.signal(signal.SIGTERM, signal.SIG_DFL)
    signal.signal(signal.SIGHUP, signal.SIG_DFL)

    # raise sig
    os.kill(os.getpid(), sig)


# -------------------------------------------------------------------------
def lookupGpdb(address, port, dataDirectory):
    """ Look up the segment gpdb by address, port, and dataDirectory """

    gpdbList = gpArrayInstance.getDbList()
    for gpdb in gpdbList:
        if address == str(gpdb.getSegmentAddress()) and port == str(gpdb.getSegmentPort()) and dataDirectory == str(
                gpdb.getSegmentDataDirectory()):
            return gpdb
    return None

# -------------------------------------------------------------------------
# -------------------------------------------------------------------------
class Mirror:
    """ This class represents information about a mirror. """

    def __init__(self, address, port, dataDirectory, inPlace=False):
        self.address = address
        self.port = port
        self.dataDirectory = dataDirectory
        self.inPlace = inPlace

    def __str__(self):
        tempStr = "address = " + str(self.address) + '\n'
        tempStr = tempStr + "port = " + str(self.port) + '\n'
        tempStr = tempStr + "dataDirectory   = " + str(self.dataDirectory) + '\n'
        return tempStr


# -------------------------------------------------------------------------
# -------------------------------------------------------------------------
class Configuration:
    """ This class represents the mirrors. """

    def __init__(self):
        self.inputFile = None
        self.fileData = None
        self.oldMirrorList = []
        self.newMirrorList = []

    def _getParsedRow(self, lineno, line):
        groups = line.split() # NOT line.split(' ') due to MPP-15675
        if len(groups) != 2:
            raise ExceptionNoStackTraceNeeded("expected 2 groups but found %d" % len(groups))
        oldParts = groups[0].split('|')
        newParts = groups[1].split('|')
        if len(oldParts) != 3 or len(newParts) != 3:
            msg = "expected 3 parts in old and new groups, got %d and %d" % (oldParts, newParts)
            raise ExceptionNoStackTraceNeeded(msg)
        oldAddress, oldPort, oldDataDirectory = oldParts
        newAddress, newPort, newDataDirectory = newParts
        check_values(lineno, address=oldAddress, port=oldPort, datadir=oldDataDirectory)
        check_values(lineno, address=newAddress, port=newPort, datadir=newDataDirectory)

        return (oldAddress, oldPort, oldDataDirectory), (newAddress, newPort, newDataDirectory)

    def read_input_file(self, inputFile):
        self.inputFile = inputFile

        with open(inputFile) as f:
            for lineno, line in line_reader(f):
                try:
                    (oldAddress, oldPort, oldDataDirectory), (newAddress, newPort, newDataDirectory) = \
                        self._getParsedRow(lineno, line)

                    oldGpdb = lookupGpdb(oldAddress, oldPort, oldDataDirectory)
                    if oldGpdb is None or oldGpdb.getSegmentRole() != ROLE_MIRROR:
                        if oldGpdb is None:
                            raise Exception(
                                "Old mirror segment does not exist with given information: address = %s, port = %s, segment data directory = %s"
                                % (oldAddress, str(oldPort), oldDataDirectory))
                        else:
                            raise Exception(
                                "Old mirror segment is not currently in a mirror role: address = %s, port = %s, segment data directory = %s"
                                % (oldAddress, str(oldPort), oldDataDirectory))
                    inPlace = (oldAddress == newAddress and oldDataDirectory == newDataDirectory)
                    if inPlace and oldPort == newPort:
                        logger.warn("input file %s contained request to move a mirror with identical attributes (%s, %s, %s)"
                                    % (inputFile, newAddress, str(newPort), newDataDirectory))
                    oldMirror = Mirror(address=oldAddress
                                       , port=oldPort
                                       , dataDirectory=oldDataDirectory
                                       , inPlace=inPlace
                                       )
                    newMirror = Mirror(address=newAddress
                                       , port=newPort
                                       , dataDirectory=newDataDirectory
                                       )
                    self.oldMirrorList.append(oldMirror)
                    self.newMirrorList.append(newMirror)
                except ValueError:
                    raise ValueError('Missing or invalid value on line %d of file: %s.' % (lineno, inputFile))
                except Exception as e:
                    raise Exception('Invalid input file on line %d of file %s: %s' % (lineno, inputFile, str(e)))

    def write_output_file(self, filename, oldToNew=True):
        """ Write out the configuration in a format appropriate for use with the -i option. """
        fp = open(filename, 'w')

        if oldToNew == True:
            firstList = self.oldMirrorList
            secondList = self.newMirrorList
        else:
            firstList = self.newMirrorList
            secondList = self.oldMirrorList
        for i in range(0, len(firstList)):
            oldEntry = firstList[i]
            newEntry = secondList[i]
            line = canonicalize_address(oldEntry.address) + \
                   "|" + str(oldEntry.port) + "|" + oldEntry.dataDirectory
            line = line + " "
            line = line + canonicalize_address(newEntry.address) + \
                   "|" + str(newEntry.port) + "|" + newEntry.dataDirectory
            line = line + '\n'
            fp.write(line)


# -------------------------------------------------------------------------
# --------------------------------- main ----------------------------------
# -------------------------------------------------------------------------
gp_movemirrors = None
remove_pid = True
options = None
args = None
pidfilepid = None  # pid of the process which has the lock
locktorelease = None
sml = None  # sml (simple main lock)
mirror_tablespace_locations = {}   #map of mirror and its tablespace locations

logger = get_default_logger()

try:

    # setup signal handlers so we can clean up correctly
    signal.signal(signal.SIGTERM, sig_handler)
    signal.signal(signal.SIGHUP, sig_handler)

    options, args = parseargs()

    setup_tool_logging(appName=EXECNAME
                       , hostname=getLocalHostname()
                       , userName=getUserName()
                       , logdir=options.logfile_directory
                       )

    enable_verbose_logging()

    mainOptions = {"pidlockpath": GP_MOVEMIRRORS_LOCK_PATH}
    sml = SimpleMainLock(mainOptions)
    otherpid = sml.acquire()
    if otherpid is not None:
        logger.error("Lockfile %s indicates that an instance of %s is already running with PID %s" % (sml.ppath, "gpmovemirrors", otherpid))
        logger.error("If this is not the case, remove the lockfile directory at %s" % (sml.ppath))
        sys.exit(1)

    logger.info("Invocation of gpmovemirrors mirrors")
    logOptionValues(options)

    dburl = dbconn.DbURL()

    # Get array configuration
    try:
        gpArrayInstance = GpArray.initFromCatalog(dburl, utility=True)
    except DatabaseError as ex:
        logger.error('Failed to connect to database.  Make sure the')
        logger.error('Cloudberry instance is running, and that')
        logger.error('your environment is correct, then rerun')
        logger.error('gpmovemirrors ' + ' '.join(sys.argv[1:]))
        sys.exit(1)

    """ Get the old to new mirror configuration. """
    newConfig = Configuration()
    newConfig.read_input_file(options.input_filename)

    PgHbaEntriesToUpdate = []
    """ Do some sanity checks on the input. """
    for oldMirror, newMirror in zip(newConfig.oldMirrorList, newConfig.newMirrorList):
        seg = lookupGpdb(oldMirror.address, oldMirror.port, oldMirror.dataDirectory)
        if seg is None:
            raise Exception(
                "Old mirror segment does not exist with given information: address = %s, port = %s, segment data directory = %s"
                % (oldMirror.address, str(oldMirror.port), oldMirror.dataDirectory))
        if seg.getSegmentContentId() == COORDINATOR_CONTENT_ID:
            raise Exception(
                "Cannot move coordinator or coordinator mirror segments: address = %s, port = %s, segment data directory = %s"
                % (oldMirror.address, str(oldMirror.port), oldMirror.dataDirectory))
        if seg.getSegmentRole() != ROLE_MIRROR:
            raise Exception(
                "Old mirror segment is not currently in a mirror role: address = %s, port = %s, segment data directory = %s"
                % (oldMirror.address, str(oldMirror.port), oldMirror.dataDirectory))
        # exclude any in place mirrors, i.e if the host on which the mirror reside still stays same, in that
        # case we need not update the entries
        if not oldMirror.inPlace:
            for segment in gpArrayInstance.getDbList():
                # identify the corresponding primary pair of the mirror being moved
                if segment.getSegmentContentId() == seg.getSegmentContentId() and segment.getSegmentRole() != seg.getSegmentRole():
                    PgHbaEntriesToUpdate.append((segment.getSegmentDataDirectory(), segment.getSegmentHostName(), newMirror.address))


    """ Prepare common execution steps for running commands on segments """
    mirrorsToDelete = [mirror for mirror in newConfig.oldMirrorList if not mirror.inPlace]
    totalDirsToDelete = len(mirrorsToDelete)
    numberOfWorkers = min(totalDirsToDelete, options.batch_size)

    """ Fetch tablespace locations before mirror move """
    for mirror in mirrorsToDelete:
        mirror_tablespace_locations[mirror.dataDirectory] = get_tablespace_locations(False, mirror.dataDirectory)

    """ Update pg_hba.conf on primary segments with new mirror information """
    update_pg_hba_for_new_mirrors(PgHbaEntriesToUpdate, options.hba_hostnames, options.batch_size)

    """ Create a backout file for the user if they choose to go back to the original configuration. """
    backout_filename = options.input_filename + "_backout_" + str(time.strftime("%m%d%y%H%M%S"))
    newConfig.write_output_file(backout_filename, False)

    """ Start gprecoverseg. """
    recoversegOptions = "-i " + newConfig.inputFile + " -B " + str(options.batch_size) + \
                        " -b " + str(options.segment_batch_size) + " -v -a -d " + options.coordinator_data_directory
    if options.logfile_directory != None:
        recoversegOptions = recoversegOptions + " -l " + str(options.logfile_directory)
    logger.info('About to run gprecoverseg with options: ' + recoversegOptions)
    cmd = GpRecoverSeg("Running gprecoverseg", options=recoversegOptions)
    cmd.run(validateAfter=True)

    """ Reinitialize gparray after the new mirrors are in place. """
    gpArrayInstance = GpArray.initFromCatalog(dburl, utility=True)
    mirrorDirectories = set((seg.getSegmentHostName(),seg.getSegmentContentId()) for seg in gpArrayInstance.getDbList() if seg.isSegmentMirror())

    """ Delete old mirror directories. """
    if numberOfWorkers > 0:
        pool = WorkerPool(numWorkers=numberOfWorkers)
        for mirror in mirrorsToDelete:
            logger.info("About to remove old mirror segment directories: " + mirror.dataDirectory)

            if mirror_tablespace_locations[mirror.dataDirectory]:
                for host, content, tablespace_dir in mirror_tablespace_locations[mirror.dataDirectory]:
                    if (host, content) not in mirrorDirectories:
                        cmd = RemoveDirectory("removing old mirror tablespace directory", tablespace_dir, ctxt=REMOTE,
                                              remoteHost=mirror.address)
                        pool.addCommand(cmd)

            cmd = RemoveDirectory("remove old mirror segment directories", mirror.dataDirectory, ctxt=REMOTE,
                                  remoteHost=mirror.address)
            pool.addCommand(cmd)

        # Wait for the segments to finish
        try:
            pool.join()
        except:
            pool.haltWork()
            pool.joinWorkers()

        failure = False
        results = []
        for cmd in pool.getCompletedItems():
            r = cmd.get_results()
            if not cmd.was_successful():
                logging.error("Unable to remove old mirror segment directory: " + str(r))
                failure = True

        pool.haltWork()
        pool.joinWorkers()
        if failure:
            logging.error("Although the system is in the process of recovering the mirrors to their new location,")
            logging.error("There was an issue removing the old mirror segment directories.")
            raise Exception("Unable to complete removal of old mirror segment directories.")

except Exception as e:
    if options is not None and options.verbose:
        logger.exception("gpmovemirrors failed. exiting...")
    else:
        logger.error("gpmovemirrors failed: %s \n\nExiting..." % e)
    sys.exit(3)

except KeyboardInterrupt:
    # Disable SIGINT while we shutdown.
    signal.signal(signal.SIGINT, signal.SIG_IGN)

    shutdown()

    # Re-enabled SIGINT
    signal.signal(signal.SIGINT, signal.default_int_handler)

    sys.exit('\nUser Interrupted')

finally:
    if sml is not None:
        sml.release()
