#!/bin/env python3

#
# query - parse and operate on BOINC job logs
#

# joblog format is 13 fields per line. 6 of those are text field names.
#
#  00
#  log timestamp in Unix epoch format
#  01  02
# 'ue' estimated run time of the work unit
#  03  04
# 'ct' actual CPU time
#  05  06
# 'fe' estimated flops
#  07  08
# 'nm' work unit name
#  09  10
# 'et' elapsed time (work unit wall clock time)
#  11  12
# 'es' status (rc)

import argparse
import glob
import json
import math
import os
import socket
import sys
import time
import xml.etree.ElementTree as ET

import homefarm

def report(d):
    if args.json:
        print(json.dumps(d))
    else:
        # print host banner
        print(("-" * (79 - len(d['host']))) + " " + d['host'])
        # report errors if they exist
        if d['err'] is not None:
            print(d['err'])
            sys.exit(2)

        # print formatted report bits
        #
        # first the WU count and total elapsed time. if there are no
        # matches, of if we've been called with '-c', we're done.
        if args.wutype != None:
            print(f"WUs matching '{d['wutype']}' for {data['proj']} in past {d['span']} hours: {data['matches']}")
        else:
            print(f"WUs for {data['proj']} in past {d['span']} hours: {data['matches']}")
        print(f"Total CPU time used:  {homefarm.strtime(d['cputime'])}")
        if d['matches'] == 0 or args.count:
            sys.exit(0)
        # we're still here, so print the mix/max/avg runtimes
        print(f"\tMin runtime: {homefarm.strtime(d['times'][0])}")
        print(f"\tMax runtime: {homefarm.strtime(d['times'][1])}")
        print(f"\tAvg runtime: {homefarm.strtime(d['times'][2])}")
        if d['matches'] < 5:
            sys.exit(0)
        # and we have more than 5 results, so show the quintile
        # bucket data
        print("WUs by quintile:")
        for i in range(5):
            print(f"\t<= {homefarm.strtime(d['quints'][i][0])}\t {d['quints'][i][1]}\t ({d['quints'][i][2]:04.1f}%)")
        # finally, a blank line for good output spacing
        print()
    sys.exit(0)


# set up and handle arguments
parser = argparse.ArgumentParser(description='parse BOINC job logs')
parser.add_argument('-p', '--project', metavar='PROJECT_NAME', dest="proj", default="ALL",
                    help="partial name of the joblog to be parsed. must be unique")
parser.add_argument('-t', '--type', metavar='WU_TYPE', dest="wutype", default=None,
                    help="the workunit type/name to filter by (e.g. MCM1)")
parser.add_argument('-s', '--timespan', metavar='HOURS', dest="span", type=float, default=24,
                    help="how many hours of logs to parse (default: 24)")
parser.add_argument('-c', '--count-only', dest="count", action="store_true",
                    help="show WU counts only; do not print min/max/quintile times")
parser.add_argument('-j', '--json', dest="json", action="store_true",
                    help="output data as JSON")
args = parser.parse_args()


# initialize the JSON data
data = { 'host': socket.gethostname(), 'proj': args.proj,
         'wutype':args.wutype, 'span': args.span, 'matches': 0,
         'cputime': 0, 'times': [], 'quints': [], 'err': None }
# calculate our timespan for log matches
timelimit = math.floor(time.mktime(time.localtime()) - (3600 * args.span))

# turn args.proj into a set of joblog filepaths
joblogs = []
if args.proj == "ALL":
    joblogs = glob.glob("/var/lib/boinc/job_log_*.txt")
else:
    globpath = "/var/lib/boinc/job_log_*{}*.txt".format(args.proj)
    joblogs = glob.glob(globpath)
    if len(joblogs) == 0:
        data['err'] = f'no matching attached projects: {args.proj} (maybe no WUs completed yet?)'
        report(data)
    if len(joblogs) > 1:
        data['err'] = f"multiple projects matching {args.proj}; must have exactly one"
        report(data)


# if 'ALL' wasn't requested then get project info, so we can use
# proper, repeatable project names rather than the string we were
# initially handed
if args.proj != "ALL":
    key=""
    with open('/var/lib/boinc/gui_rpc_auth.cfg', 'r') as f:
        key=f.readline()
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        try:
            s.connect((data['host'], 31416))
        except:
            data['err'] = "could not connect with key '{key}'"
            report(data)
        # try to auth
        authed = homefarm.auth_conn(s, key)
        if not authed:
            data['err'] = "could not auth with key '{key}'"
            report(data)
        req = "<boinc_gui_rpc_request><get_project_status/></boinc_gui_rpc_request>\003"
        s.sendall(req.encode('utf-8'))
        reply = ET.fromstring(homefarm.recv_msg(s))
        for p in reply.iter('project'):
            url = p.find('master_url').text
            if args.proj in url:
                data['proj'] = p.find('project_name').text
            else:
                continue



matches = []
cputime = 0
f = None

for joblog in joblogs:
    try:
        f = open(joblog, 'r')
    except:
        if args.json:
            data['err'] = f"couldn't open joblog {r}"
            report(data)

    for line in f:
        line = line.rstrip("\n")
        # split line into fields: see top comment for record layout
        fields = line.split()
        # skip this line if it's not within our timespan
        if int(fields[0]) < timelimit:
            continue
        # skip line if it's not the type/name we're looking for
        if args.wutype != None:
            if args.wutype not in fields[8]:
                continue
        # match! turn the elapsed time into a rounded integer
        time = round(float(fields[4]))
        # add the result to our list of matches (we care about the
        # timings for future calculations, not the WU itself)
        matches.append(time)
        # and accumulate the time into a total
        cputime += time

# capture the number of matches and total elapsed time
data['matches'] = len(matches)
data['cputime'] = cputime
if data['matches'] > 0:
    # calculate min/max/avg times
    mintime = min(matches)
    maxtime = max(matches)
    avgtime = sum(matches)/len(matches)
    data['times'] = [ mintime, maxtime, avgtime ]
    # calculate the data for quintuple timing buckets
    # qspan is the width/duration of a bucket, in seconds
    qspan = int((max(matches) -  min(matches)) / 5)
    # q[] just pre-calculates the timings of the individual quintiles,
    # since those are used in more than one place in the code that
    # follows
    q = [ min(matches) + qspan, min(matches) + qspan * 2, min(matches) + qspan * 3,
          min(matches) + qspan * 4, max(matches) ]
    # qcounts is the count per quintile. counting happens right after
    qcounts = [0, 0, 0, 0, 0]
    for m in matches:
        if m >= min(matches) and m < q[0]:
            qcounts[0] += 1
        elif m >= q[0] and m < q[1]:
            qcounts[1] += 1
        elif m >= q[1] and m < q[2]:
            qcounts[2] += 1
        elif m >= q[2] and m < q[3]:
            qcounts[3] += 1
        else:
            qcounts[4] += 1
    # now we store the timing, count, and percentage of total for each quint
    for i in range(5):
        data['quints'].append( [ q[i], qcounts[i],  qcounts[i] / len(matches) * 100 ] )

report(data)
