#!/usr/bin/env python3
# Copyright (c) 2021 Arista Networks, Inc.  All rights reserved.
# Arista Networks, Inc. Confidential and Proprietary.

# pylint: disable=redefined-outer-name
# pylint: disable=multiple-statements
# pylint: disable=bad-indentation

# wfc: wait-for-cpu 
#
# Wait for a process to stabilize CPU usage

from __future__ import absolute_import, division, print_function
import os, sys, time, argparse, datetime

def top(pids):
    data = {}
    cmd  = 'top -H -b -n 1'
    # pylint: disable-next=anomalous-backslash-in-string
    cmd += '| grep "%s"' % '\|'.join(['^\s*'+pid for pid in pids])
    out = os.popen(cmd).read().strip().split('\n')
    for line in out:
        if not line: continue
        parts = line.split()
        cpu = float(parts[-4])
        pid = parts[0]
        data[pid] = [cpu]
    #end
    return data
#end

def process_cpu_usage(pid):
    data = CPU[pid]
    if len(data) < 5: return
    stable = all( cpu < args.threshold for cpu in data )
    PST[pid] = [] if not stable else ( PST[pid] + [current] )
    CPU[pid] = []
#end

def check_stable():

    if args.debug:
       print()
       for pid in PST: print( 'pid:', pid, 'stable:', len( PST[ pid ] ) )
       print()
    #end

    stable = all( len(PST[pid]) > 4 for pid in PST )
    if stable:
       last = max([ PST[pid][0] for pid in PST ])
       print( int( ( last - start ).total_seconds() - 5 ) )
       sys.exit(0)
    #end
#end

def check_timeout():
    elapsed = (current - start).total_seconds()
    if elapsed > args.timeout:
       sys.stderr.write('timeout\n')
       sys.exit(1)
    #end
#end

#------------------------------------------------------------------------------

ap = argparse.ArgumentParser(
   description='wfc: wait-for-cpu'
)

ap.add_argument( '-t', 
   action='store', 
   dest='threshold',
   type=int,
   default=7, 
   help='CPU threshold (%%)'
)

ap.add_argument( '-T', 
   action='store', 
   dest='timeout', 
   type=int,
   default=3600,
   help='Timeout (secs)'
)

ap.add_argument( '-d', 
   action='store_true', 
   dest='debug',
   default=False, 
   help='Show debug output'
)

ap.add_argument( 'proc', 
   nargs='+',
   action='store', 
   help='Process name'
)

args = ap.parse_args()

pids = []

# get the pids for the process tree
# XXX: we do not really want child processes
# XXX: just all threads within a process
for proc in args.proc:
    # pylint: disable-next=consider-using-f-string
    ppids = os.popen('pidof %s' % proc).read().strip().split()
    for pid in ppids:
       # pylint: disable-next=consider-using-f-string,anomalous-backslash-in-string
       cmd = 'pstree -lp %s | grep -oe "\([0-9]*\)"' % pid
       out = os.popen(cmd).read().strip().split('\n')
       for line in out: pids += [line]
    #end
#end

if not pids:
   sys.stderr.write('invalid proc\n')
   sys.exit(1)
#end

data = top(pids)

pids = set( data )

PST = { pid : [] for pid in pids }
CPU = { pid : [] for pid in pids }

start = datetime.datetime.now()

while 1:
     
    time.sleep( 1 )

    data = top(pids)

    current = datetime.datetime.now()

    for pid in pids: CPU[pid] += ( data[pid] if pid in data else [-1] )
    
    for pid in CPU: process_cpu_usage(pid)

    check_stable()

    check_timeout()

#end
