#!/usr/bin/env python

import sys
import getopt
import time
import pickle
import os
import re
import syslog

import isi.hw.fw as fw
import isi.hw.fw.config as fwConfig

VIT_NAME = 'Vit-7154' # from firmware.xml

def log(msg):
    '''Make it easy to do smarter logging later.'''
    print msg

class ClusterFirmwareVersions(object):
    '''Convenience wrapper for cluster-wide device records.'''

    def __init__(self, devices=None):
        # devices is in this format:
        # { int(nodeId) : [ Device() ] }
        self._devices = devices or {}
        self.fwPackage = None

    def getNodeVitesseVersions(self):
        '''Generate (nodeId, set(['vitesse version']))'''

        for (node, nodeDevices) in self._devices.items():
            vitesse = set([dev.version(VIT_NAME)
                for dev in nodeDevices if dev.hasFirmware(VIT_NAME)])

            yield (node, vitesse)

    def getVitesseVersions(self):
        '''Return a set of all vitesse versions on the cluster.'''
        allVersions = set()
        for (node, versions) in self.getNodeVitesseVersions():
            allVersions.update(versions)

        return allVersions

    def getVersionsByIdentity(self):
        '''Return {'node-unique device id': set(['versions across cluster'])}'''
        result = {}
        for device in sum(self._devices.values(), []):
            s = result.setdefault(device.identity(), set())
            s.add(device.version())

        return result

    def getDevicesByNode(self):
        return self._devices

    def _filteredVersions(self):
        # I don't care much about the format of the result of this, as long
        # as it compares properly, so a change in nodes, devices, or firmware
        # versions would make it unequal.
        result = {}
        for (node, devs) in self._devices.items():
            versions = [(dev.identity(), dev.version()) for dev in devs]
            result[node] = dict(versions)

        return result

    def __eq__(self, other):
        if self.fwPackage != other.fwPackage:
            return False
        return self._filteredVersions() == other._filteredVersions()

def myId():
    from isi.config import getNode
    return getNode()['node id']

def haveLowestId():
    from isi.fs.restripe.group import GroupInfo
    nodeIds = GroupInfo().getNodeList(ignoreSfFlag=False,
        ignoreReadOnlyFlag=False)

    return bool(nodeIds) and myId() == min(nodeIds)

# ##### Stuff for dealing with firmware records #####

_fwVersionsDir = '/ifs/.ifsvar/modules/firmware'
_clusterVersionsFile = os.path.join(_fwVersionsDir, 'fw-cluster.pickle')
# When the data format stored in the pickled files changes, increment this.
_formatVersion = 0

# Simple pickle wrappers
def pdump(data, f):
    pickle.dump((_formatVersion, data), f, pickle.HIGHEST_PROTOCOL)

def pload(f, default=None):
    try:
        unpickled = pickle.load(f)
    except (SystemExit, KeyboardInterrupt):
        raise
    except Exception, e:
        # So many things can go wrong unpickling, just report it and continue.
        log('Error loading data: %s', e)
        return default

    try:
        (format, data) = unpickled
    # Sadly, tuple unpacking can fail in at least two ways.
    except (TypeError, ValueError):
        return default

    # If the format is not current, we won't even bother with it.
    if format != _formatVersion:
        data = default

    return data

def createFwDir():
    try:
        os.makedirs(_fwVersionsDir)
    except os.error, e:
        import errno
        # Ignore failure due to it already existing, propagate other errors.
        if e.errno != errno.EEXIST:
            raise

def updateFwRecords():
    log('Reading and recording my firmware versions')

    devices = fw.probe()

    filename = 'fw-node-%d.pickle' % myId()
    f = open(os.path.join(_fwVersionsDir, filename), 'w')
    try:
        pdump(devices, f)
    finally:
        f.close()

def readPrevFwVersions():
    try:
        f = open(_clusterVersionsFile, 'r')
        try:
            return pload(f, default=ClusterFirmwareVersions())
        finally:
            f.close()
    except IOError:
        return ClusterFirmwareVersions()

def writeFwVersions(versions):
    f = open(_clusterVersionsFile, 'w')
    try:
        pdump(versions, f)
    finally:
        f.close()

_fwRecordMatcher = re.compile(r'fw-node-(\d+)\.pickle')

def fwRecordFiles():
    filenames = os.listdir(_fwVersionsDir)
    for filename in filenames:
        match = _fwRecordMatcher.match(filename)
        if match is None:
            continue
        full_name = os.path.join(_fwVersionsDir, filename)
        node = int(match.group(1))
        yield (node, full_name)

def readFwRecords():
    '''
    Read the record file for each node, then construct and
    return { int(nodeId) : [ Device() ] }
    '''
    records = {}
    for (node, filename) in fwRecordFiles():
        f = open(filename, 'r')
        try:
            records[node] = pload(f, default=[])
        finally:
            f.close()

    return records

def deleteFwRecords():
    for (node, filename) in fwRecordFiles():
        try:
            os.unlink(filename)
        except OSError, e:
            log('failed to delete %s (errno %d: %s)' %
                    (filename, e.errno, e.strerror))

# ##### Alerts fiddling #####

def raiseAlert():
    import isi.celog
    from isi.celog.eventids import eventid_map

    fwev = isi.celog.Event(
            event_id=eventid_map["SW_FIRMWARE_UPDATE_FAILED"],
            severity=isi.celog.EVENT_LEVEL_WARNING,
            message="A firmware update has not been completely applied")

    fwev.add()

def cancelAlert():
    # XXXTSU: this should use anonymous end
    import isi.celog.query as cq
    from isi.celog.eventids import eventid_map
    db = cq.Query()
    fwevs = db.get_events(
            eventids=[eventid_map["SW_FIRMWARE_UPDATE_FAILED"]])
    for ev in fwevs:
        ev.end()

# ##### Actual checks #####

def getVitesseMismatches(versions):
    '''
    Return a set of all vitesse versions that contribute to a 'local'
    mismatch. This may include versions from across the cluster, but
    only those that fail to match within some node-expander group.
    '''

    mismatchVersions = set()
    for (node, vitesseVers) in versions.getNodeVitesseVersions():
        if len(vitesseVers) > 1:
            mismatchVersions.update(vitesseVers)

    return mismatchVersions

def getFirmwareMismatches(versions):
    '''
    Return { 'DeviceName' : set(['version numbers']) } indicating all
    the versions for the device (which is unique per-node) across the
    network -- EXCEPT that all entries with a set size less than 2
    are removed, leaving only versions that vary across the cluster.
    '''

    versionsByIdent = versions.getVersionsByIdentity()
    filtered = [(k, v) for (k, v) in versionsByIdent.items() if len(v) > 1]
    return dict(filtered)

def doFirmwarePackageCheck(currentVersions, prevVersions):
    log('Checking device firmware against installed update package')
    # If neither the installed firmware package or the current versions
    # have changed, no nothing.
    if currentVersions == prevVersions:
        log('No firmware or firmware package changes -- skipping')
        return

    packageDevices = fwConfig.DeviceConfig().assemblies()
    # all devices in the cluster
    allDevices = sum(currentVersions.getDevicesByNode().values(), [])
    mismatchedVers = {} # {'user-facing dev id': set(['non-matching vers'])}
    unknownDevs = set()
    for device in allDevices:
        packageDevice = packageDevices.get(device.name())
        if packageDevice is None or not packageDevice.version():
            unknownDevs.add(device.userIdentity())
            continue

        if device.version() != packageDevice.version():
            s = mismatchedVers.setdefault(device.userIdentity(), set())
            s.add(device.version())

    if unknownDevs:
        log("I don't know any valid firmware versions for the the following:\n"
                + '\n'.join(['  ' + dev for dev in unknownDevs]))

    if mismatchedVers:
        for (identity, versions) in mismatchedVers.items():
            log('The following versions for %s do not match the installed\n'
                    'firmware package:\n%s'
                    % (identity, '\n'.join(['  ' + ver for ver in versions])))
        raiseAlert()
    else:
        cancelAlert()

def doFirmwareChecks():
    # If no fw package is installed, there's nothing to check.
    firmwarePackage = fwConfig.FirmwareConfig().toolVersion()
    if firmwarePackage is None:
        cancelAlert()
        deleteFwRecords()
        return

    currentVersions = ClusterFirmwareVersions(readFwRecords())
    currentVersions.fwPackage = firmwarePackage
    # Remove the records now that we have read them so we don't end up using
    # stale records next time around.
    deleteFwRecords()
    prevVersions = readPrevFwVersions()

    # Save current version records for next time around.
    writeFwVersions(currentVersions)

    # Call each check here
    doFirmwarePackageCheck(currentVersions, prevVersions)

# ##### main, etc. #####

def usage(error=None):
    if error:
        print 'Error: %s' % error.capitalize()

    print 'Usage:'
    print '  %s [-h|--help] [-c|--check] [-d<seconds>|--delay=<seconds>]' \
        % sys.argv[0]

def main(argv=None):
    if argv is None:
        argv = sys.argv

    try:
        (opts, args) = getopt.getopt(argv[1:], 'hcd:',
                ['help', 'check', 'delay='])
    except getopt.GetoptError, e:
        usage(str(e))
        return 2

    if args:
        usage('Unexpected arguments: %s' % ' '.join(args))
        return 2

    doCheck = False
    delaySeconds = 0
    for (opt, val) in opts:
        if opt in ('-h', '--help'):
            usage()
            return 0
        elif opt in ('-c', '--check'):
            doCheck = True
        elif opt in ('-d', '--delay'):
            try:
                delaySeconds = float(val)
            except ValueError, e:
                usage('Invalid delay time: %s' % val)
                return 2

    # Responsibility for checking all the firmware versions falls to the node
    # with the lowest id if not overridden with the --check option.
    if not doCheck:
        doCheck = haveLowestId()

    createFwDir()
    updateFwRecords()

    if not doCheck:
        return 0

    if delaySeconds > 0:
        log('Sleeping for %g seconds to allow nodes to report...' \
            % delaySeconds)
        time.sleep(delaySeconds)

    doFirmwareChecks()

    return 0


if __name__ == '__main__':
    sys.exit(main())

# vim: et sts=4 sw=4
