applutil

#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import print_function

import sys

from urllib import quote, quote_plus

import json
try:
    from io import StringIO
except ImportError:
    from cStringIO import StringIO

from twisted.internet import protocol, reactor, defer
from twisted.web.client import Agent, ResponseDone
from twisted.web.http import Headers

agent = Agent(reactor, connectTimeout=3)

class ReadAll(protocol.Protocol):
    def __init__(self):
        self._S, self.defer = StringIO(), defer.Deferred()
    def dataReceived(self, raw):
        self._S.write(raw)
    def connectionLost(self, reason):
        if reason.check(ResponseDone):
            S = self._S.getvalue()
            try:
                J = json.loads(S)
            except:
                print(S)
                import traceback
                traceback.print_exc()
                self.defer.errback()
            else:
                self.defer.callback(J)
        else:
            self.defer.errback(reason)

class ProduceAll(object):
    def __init__(self, content):
        self.length, self.content = len(content), content
    def startProducing(self, consumer):
        consumer.write(self.content)
        return defer.succeed(None)
    def resumeProducing(self):
        pass
    def pauseProducing(self):
        pass
    def stopProcuding(self):
        pass

@defer.inlineCallbacks
def getURL(url):
    #print('GET',url,file=sys.stderr)
    R = yield agent.request('GET', str(url))

    if R.code!=200:
        print('oops',R.code,url,file=sys.stderr)
        raise RuntimeError("%s: %s"%(R.code, url))

    B = ReadAll()
    R.deliverBody(B)

    X = yield B.defer
    defer.returnValue(X)

@defer.inlineCallbacks
def addPV(pvs, I):
    H = Headers()
    H.setRawHeaders('Content-Type', ['application/json'])
    body = json.dumps([{'pv':pv} for pv in pvs])
    url = str('%s/archivePV'%I['mgmtURL'])
    R = yield agent.request('POST', url, headers=H,
                            bodyProducer=ProduceAll(body))

    if R.code!=200:
        print('oops',R.code,url,file=sys.stderr)

    B = ReadAll()
    R.deliverBody(B)

    X = yield B.defer
    if R.code!=200:
        raise RuntimeError("%s: %s"%(R.code, url))
    for E in X:
        print(E)

@defer.inlineCallbacks
def pausePV(pvs, I):
    while len(pvs):
        pv = pvs.pop()
        url = '%s/pauseArchivingPV?pv=%s'%(I['mgmtURL'],quote_plus(pv))
        R = yield getURL(url)
        print(pv,R,file=sys.stderr)

@defer.inlineCallbacks
def resumePV(pvs, I):
    while len(pvs):
        pv = pvs.pop()
        url = '%s/resumeArchivingPV?pv=%s'%(I['mgmtURL'],quote_plus(pv))
        R = yield getURL(url)
        print(pv,R,file=sys.stderr)

@defer.inlineCallbacks
def deletePV(pvs, I):
    while len(pvs):
        pv = pvs.pop()
        url = '%s/deletePV?pv=%s'%(I['mgmtURL'],quote_plus(pv))
        R = yield getURL(url)
        print(pv,R,file=sys.stderr)

@defer.inlineCallbacks
def abortPV(pvs, I):
    while len(pvs):
        pv = pvs.pop()
        url = '%s/abortArchivingPV?pv=%s'%(I['mgmtURL'],quote_plus(pv))
        R = yield getURL(url)
        print(pv,R,file=sys.stderr)

@defer.inlineCallbacks
def statusPV(pvs, I):
    while len(pvs):
        pv = pvs.pop()
        url = '%s/getPVStatus?pv=%s'%(I['mgmtURL'],quote_plus(pv))
        R = (yield getURL(url))[0]
        if R['status']=='Being archived':
            print('%s\t\t%s\t%s'%(pv,R['status'],R['lastEvent']))
        else:
            print('%s\t\t%s'%(pv,R['status']))

@defer.inlineCallbacks
def setRate(pvs, I, rate):
    period = 1.0/rate
    print('Set rate limit to %f Hz (%f sec)'%(rate, period))
    while len(pvs):
        pv = pvs.pop()
        url = '%s/changeArchivalParameters?pv=%s&samplingmethod=MONITOR&samplingperiod=%f'%(I['mgmtURL'],quote_plus(pv),period)
        print('ACT', pv, url)
        R = yield getURL(url)
        print(pv,R,file=sys.stderr)

@defer.inlineCallbacks
def doStuff(opts, args):
    print('Fetch Appliance info from',opts.url,file=sys.stderr)
    I = yield getURL(opts.url)
    print('Info',I,file=sys.stderr)

    if opts.pvs=='pattern':
        pvs=set()
        for pat in args:
            R = yield getURL('%s/getAllPVs?pv=%s'%(I['mgmtURL'], quote_plus(pat)))
            pvs.update(R)
        pvs=list(pvs)

    elif opts.pvs=='disconn':
        R = yield getURL('%s/getCurrentlyDisconnectedPVs'%(I['mgmtURL'],))

        pvs = [ent['pvName'] for ent in R]

    elif opts.pvs=='never':
        R = yield getURL('%s/getNeverConnectedPVs'%(I['mgmtURL'],))

        pvs = [ent['pvName'] for ent in R]

    elif opts.pvs=='pause':
        R = yield getURL('%s/getPausedPVsReport'%(I['mgmtURL'],))

        pvs = [ent['pvName'] for ent in R]

    elif opts.pvs=='all':
        pvs = yield getURL('%s/getAllPVs'%(I['mgmtURL'],))

    elif opts.pvs=='file':
        pvs = []
        for fname in args:
            if fname=='-':
                F = sys.stdin
            else:
                F = open(fname, 'r')
            [pvs.append(pv.split(None,1)[0].strip()) for pv in F]
            F.close()

    if opts.set_rate:
        yield setRate(pvs, I, opts.set_rate)
    elif opts.action is None:
        map(print, pvs)
    else:
        N = len(pvs)/opts.count
        if N<1:
            N=1
        Ds = [opts.action(pvs[i:i+N], I) for i in range(0, len(pvs), N)]
        print('Making',len(Ds),'Concurrent request(s)')
        yield defer.DeferredList(Ds, fireOnOneErrback=True)#, consumeErrors=True)

    print(len(pvs), 'PVs',file=sys.stderr)

if __name__=='__main__':
    from optparse import OptionParser
    P = OptionParser(usage='%prog [options] [<pvpattern(s)>|--disconn|--never|--paused]')
    P.add_option('-U','--url', help='Appliance info URL')
    P.add_option('--add', action='store_const', dest='action', const=addPV, default=None,
                 help='Start archiving')
    P.add_option('--pause', action='store_const', dest='action', const=pausePV,
                 help='Pause archiving')
    P.add_option('--resume', action='store_const', dest='action', const=resumePV,
                 help='Resume archiving')
    P.add_option('--delete', action='store_const', dest='action', const=deletePV,
                 help='Stop archiving and delete')
    P.add_option('--status', action='store_const', dest='action', const=statusPV,
                 help='Show PV status')
    P.add_option('--abort', action='store_const', dest='action', const=abortPV,
                 help='Cancel archiving request')
    P.add_option('--set-rate', type=float,
                 help='Set upper limit on PV update rate')

    P.add_option('--disconn', action='store_const', dest='pvs', const='disconn', default='pattern',
                 help='Operate on disconnected PVs')
    P.add_option('--never', action='store_const', dest='pvs', const='never',
                 help='Operate on never connected PVs')
    P.add_option('--paused', action='store_const', dest='pvs', const='pause',
                 help='Operate on paused PVs')
    P.add_option('--all', action='store_const', dest='pvs', const='all',
                 help='Operate on all PVs')
    P.add_option('-i', '--file', action='store_const', dest='pvs', const='file',
                 help='Operate on PVs from file ( - for stdin)')

    P.add_option('-C','--count', type=int, default=5,
                 help='Number of parallel operations')

    opts, args = P.parse_args()

    if opts.action is addPV and opts.count!=1:
        P.error('Always use --count 1 with --add')

    if len(args)==0 and opts.pvs=='pattern':
        P.error("one of PV name(s) or --disconn or --never is required")

    if not opts.url.startswith('http') and opts.url.find('/')==-1:
        opts.url='http://%s/mgmt/bpl/getApplianceInfo'%opts.url

    D = doStuff(opts, args)
    def error(E):
        if isinstance(E, defer.FirstError):
            E = E.subFailure
        E.printDetailedTraceback()
    D.addErrback(error)
    def stop(D):
        reactor.stop()
        return D
    D.addBoth(stop)
    reactor.run()
    print('Done',file=sys.stderr)