#!/usr/bin/env python3

import copy
import argparse
import json
import sys
import os
import re
import shutil

CONNECTION_JSON_FILE_PATH = "/var/lib/nfs/syno_connection_stat.json"
CONNECTION_JSON_OLD_FILE_PATH = CONNECTION_JSON_FILE_PATH + ".old"

ERRORS_JSON_FILE_PATH = "/var/lib/nfs/syno_errors_stat.json"
ERRORS_JSON_OLD_FILE_PATH = ERRORS_JSON_FILE_PATH + ".old"
ERRORS_JSON_TMP_FILE_PATH = ERRORS_JSON_FILE_PATH + ".tmp"

LATENCY_HISTOGRAM_JSON_FILE_PATH = "/var/lib/nfs/syno_latency_histogram.json"
LATENCY_HISTOGRAM_JSON_OLD_FILE_PATH = LATENCY_HISTOGRAM_JSON_FILE_PATH + ".old"
LATENCY_HISTOGRAM_JSON_TMP_FILE_PATH = LATENCY_HISTOGRAM_JSON_FILE_PATH + ".tmp"

TOTAL_CONNECTION_KEY = "total_conn"
MAX_CONNECTION_KEY = "max_conn"
PROC_TOTAL_CONNECTION_STAT = "/proc/fs/nfsd/syno_total_connection_stat"
PROC_MAX_CONNECTION = "/proc/fs/nfsd/syno_max_connection"
PROC_TOTAL_CONNECTION_RESET = "/proc/fs/nfsd/syno_total_connection_reset"
PROC_TOTAL_ERRORS = "/proc/fs/nfsd/syno_total_error"
#TODO: use historgram
PROC_LATENCY_HISTOGRAM = "/proc/fs/nfsd/syno_latency_histogram"
LATENCY_HISTOGRAM_SIZE_IN_ONE_ROW = 18 # 3 ( version ) * 3 ( stage ) * 2 ( r, w )

########################################## GENERIC ###########################################

def read_proc(path):
    try:
        with open(path, 'r') as F:
            return F.readlines()
    except:
        return []

def write_proc(path, val):
    if not os.path.isfile(path):
        return
    try:
        with open(path, 'w') as F:
            F.write(val)
    except Exception as e:
        print('Failed to write {} into {}, err {}'.format(val, path, e))

def read_json_file(path):
    try:
        with open(path, 'r') as F:
            return json.load(F)
    except:
        return {}

def write_json_file(path, content):
    try:
        json_object = json.dumps(content, indent=4)
        with open(path, 'w') as F:
            F.write(json_object)
    except Exception as e:
        print('Failed to write json file {}'.format(e))

######################################### CONNECTION #########################################

def update_connection_stat(stat, total_conn, max_conn):
    if MAX_CONNECTION_KEY not in stat:
        stat[MAX_CONNECTION_KEY] = max_conn
    else:
        stat[MAX_CONNECTION_KEY] = max(max_conn, stat[MAX_CONNECTION_KEY])

    if TOTAL_CONNECTION_KEY not in stat:
        stat[TOTAL_CONNECTION_KEY] = {}

    for key, vers in total_conn.items():
        if key not in stat[TOTAL_CONNECTION_KEY]:
            stat[TOTAL_CONNECTION_KEY][key] = { "v2": 0, "v3": 0, "v4": 0 }
        stat[TOTAL_CONNECTION_KEY][key]["v2"] = max(vers[0], stat[TOTAL_CONNECTION_KEY][key]["v2"])
        stat[TOTAL_CONNECTION_KEY][key]["v3"] = max(vers[1], stat[TOTAL_CONNECTION_KEY][key]["v3"])
        stat[TOTAL_CONNECTION_KEY][key]["v4"] = max(vers[2], stat[TOTAL_CONNECTION_KEY][key]["v4"])

    return stat

def read_total_connection():
    stat = {}
    lines = read_proc(PROC_TOTAL_CONNECTION_STAT)
    for line in lines:
        host = re.search(r'^\"(.+)\":', line).group(1)
        values = [int (n) for n in line.replace('\"' + host + '\":', '').split()]
        stat[host] = values

    return stat

def read_max_connection():
    lines = read_proc(PROC_MAX_CONNECTION)
    if len(lines) == 1:
        return int(lines[0])
    return 0

def save_connection_stat():
    total_connection = read_total_connection()
    max_connection = read_max_connection()
    stat = read_json_file(CONNECTION_JSON_FILE_PATH)
    new_stat = update_connection_stat(stat, total_connection, max_connection)
    write_json_file(CONNECTION_JSON_FILE_PATH, new_stat)

def reset_connection_stat():
    write_proc(PROC_TOTAL_CONNECTION_RESET, "1")
    write_proc(PROC_MAX_CONNECTION, "0")

def dump_connection_stat():
    connection = read_json_file(CONNECTION_JSON_FILE_PATH)
    print(json.dumps(connection))

######################################## TOTAL ERRORS ########################################

def read_total_errors():
    stat = {"v2": {}, "v3": {}, "v4": {}}
    lines = read_proc(PROC_TOTAL_ERRORS)
    for line in lines:
        no = re.search(r'^err (\d+):', line).group(1)
        values = [int (n) for n in line.replace('err ' + no + ':', '').split()]
        vers = 2
        for val in values:
            stat["v" + str(vers)][no] = val
            vers += 1
    return stat

def init_errors_stat(src, stat):
    for vers, errors in src.items():
        if vers not in stat:
            stat[vers] = {}
        for err, cnt in errors.items():
            if err in stat[vers]:
                continue
            stat[vers][err] = 0
    return stat

def update_errors_stat(stat, diff):
    # init stat with missing key from diff.
    stat = init_errors_stat(diff, stat)
    # do update
    for vers, errors in stat.items():
        for err, cnt in errors.items():
            if vers in diff and err in diff[vers]:
                errors[err] += diff[vers][err]
    return stat;

def calc_errors_diff(curr, prev):
    for vers, errors in curr.items():
        if vers not in prev:
            continue
        for err, cnt in errors.items():
            if err not in prev[vers]:
                continue
            errors[err] = max(0, errors[err] - prev[vers][err])
    return curr;

def dump_errors_stat():
    prev = read_json_file(ERRORS_JSON_OLD_FILE_PATH)
    stat = read_json_file(ERRORS_JSON_FILE_PATH)
    errors = calc_errors_diff(stat, prev)
    print(json.dumps(errors))

def save_errors_stat():
    errors = read_total_errors()
    prev = read_json_file(ERRORS_JSON_TMP_FILE_PATH)
    # You should save errors first before caculating diff
    write_json_file(ERRORS_JSON_TMP_FILE_PATH, errors)
    diff = calc_errors_diff(errors, prev)
    old_stat = read_json_file(ERRORS_JSON_FILE_PATH)
    new_stat = update_errors_stat(old_stat, diff)
    write_json_file(ERRORS_JSON_FILE_PATH, new_stat)

##################################### LATENCY HISTOGRAM ######################################

def read_latency_histogram():
    rws = ["r", "w"]
    stages = ["total", "vfs", "rpc"]
    vers = ["v2", "v3", "v4"]
    rw_stat = {type: {} for type in rws}
    stage_stat = {}
    for stage in stages:
        stage_stat[stage] = copy.deepcopy(rw_stat)
    stat = {}
    for v in vers:
        stat[v] = copy.deepcopy(stage_stat)
    lines = read_proc(PROC_LATENCY_HISTOGRAM)
    time = 1
    for line in lines:
        values = [int (n) for n in line.split()]
        if len(values) != LATENCY_HISTOGRAM_SIZE_IN_ONE_ROW:
            return {}
        idx = 0
        for v in vers:
            for stage in stages:
                for type in rws:
                    stat[v][stage][type][str(time)] = values[idx]
                    idx += 1
        time *= 2;
    return stat

def init_latency_stat(src, stat):
    for vers, stages in src.items():
        if vers not in stat:
            stat[vers] = {}
        for stage, types in stages.items():
            if stage not in stat[vers]:
                stat[vers][stage] = {}
            for type, hist in types.items():
                if type not in stat[vers][stage]:
                    stat[vers][stage][type] = {}
                for time, cnt in hist.items():
                    if time in stat[vers][stage][type]:
                        continue
                    stat[vers][stage][type][time] = 0
    return stat

def update_latency_histogram_stat(stat, diff):
    # init stat with missing key from diff
    stat = init_latency_stat(diff, stat)
    # update with diff
    for vers, stages in stat.items():
        for stage, types in stages.items():
            for type, hist in types.items():
                for time, cnt in hist.items():
                    if vers in diff and stage in diff[vers] and type in diff[vers][stage] and time in diff[vers][stage][type]:
                        hist[time] += diff[vers][stage][type][time]
    return stat

def calc_latency_diff(stat, prev):
    for vers, stages in stat.items():
        for stage, types in stages.items():
            for type, hist in types.items():
                for time, cnt in hist.items():
                    if vers in prev and stage in prev[vers] and type in prev[vers][stage] and time in prev[vers][stage][type]:
                        hist[time] = max(0, hist[time] - prev[vers][stage][type][time])
    return stat

def save_latency_histogram_stat():
    latency = read_latency_histogram()
    prev = read_json_file(LATENCY_HISTOGRAM_JSON_TMP_FILE_PATH)
    # You should save latency first before caculating diff
    write_json_file(LATENCY_HISTOGRAM_JSON_TMP_FILE_PATH, latency)
    diff = calc_latency_diff(latency, prev)
    old_stat = read_json_file(LATENCY_HISTOGRAM_JSON_FILE_PATH)
    new_stat = update_latency_histogram_stat(old_stat, diff)
    write_json_file(LATENCY_HISTOGRAM_JSON_FILE_PATH, new_stat)

def calc_percentile(counts, target, times):
    if target == 0:
        return 0
    for i in range(len(counts)):
        target -= min(target, counts[i])
        if target == 0:
            if len(times) <= i:
                return 0
            return times[i]
    return 0

def dump_latency_histogram():
    prev = read_json_file(LATENCY_HISTOGRAM_JSON_OLD_FILE_PATH)
    curr = read_json_file(LATENCY_HISTOGRAM_JSON_FILE_PATH)
    perc = [50.0, 60.0, 70.0, 80.0, 90.0, 95.0, 99.0, 99.9]
    dump_stat = {}
    stat = calc_latency_diff(curr, prev)
    times = []

    for vers, stages in stat.items():
        dump_stat[vers] = {}
        for stage, types in stages.items():
            dump_stat[vers][stage] = {}
            for type, hist in types.items():
                dump_stat[vers][stage][type] = {}
                total_count = 0
                counts = []
                perc_stat = {}
                times = [int (n) for n in hist.keys()]
                times.sort()
                for t in times:
                    time = str(t)
                    if time not in hist:
                        return {}
                    total_count += hist[time]
                    counts.append(hist[time])
                for p in perc:
                    perc_stat[str(p)] = calc_percentile(counts, total_count * p / 100.0,  times)
                dump_stat[vers][stage][type]["raw"] = hist
                dump_stat[vers][stage][type]["perc"] = perc_stat
    print(json.dumps(dump_stat))

####################################### MAIN FUNCTIONS #######################################

def save_udc_stat():
    save_connection_stat()
    save_errors_stat()
    save_latency_histogram_stat()

def reset_udc_stat():
    reset_connection_stat()
    if os.path.isfile(CONNECTION_JSON_FILE_PATH):
        os.rename(CONNECTION_JSON_FILE_PATH, CONNECTION_JSON_OLD_FILE_PATH)
    if os.path.isfile(ERRORS_JSON_FILE_PATH):
        shutil.copyfile(ERRORS_JSON_FILE_PATH, ERRORS_JSON_OLD_FILE_PATH)
    if os.path.isfile(LATENCY_HISTOGRAM_JSON_FILE_PATH):
        shutil.copyfile(LATENCY_HISTOGRAM_JSON_FILE_PATH, LATENCY_HISTOGRAM_JSON_OLD_FILE_PATH)

def init_udc_stat():
    if os.path.isfile(ERRORS_JSON_TMP_FILE_PATH):
        os.unlink(ERRORS_JSON_TMP_FILE_PATH)
    if os.path.isfile(LATENCY_HISTOGRAM_JSON_TMP_FILE_PATH):
        os.unlink(LATENCY_HISTOGRAM_JSON_TMP_FILE_PATH)


def main():
    parser = argparse.ArgumentParser(description='NFS: collect stat for udc')
    group = parser.add_mutually_exclusive_group()
    group.add_argument('--dump_conn', help='Dump connection stat from json format', action='store_true')
    group.add_argument('--dump_err', help='Dump errors stat from json format', action='store_true')
    group.add_argument('--dump_lat', help='Dump latency histogram  from json format', action='store_true')
    group.add_argument('-s', '--save', help='Save stat into json file', action='store_true')
    group.add_argument('-r', '--reset', help='Reset all stat json file', action='store_true')
    group.add_argument('-i', '--init', help='Re-init stat when nfsd service stop', action='store_true')

    args = parser.parse_args()

    if args.dump_conn is True:
        dump_connection_stat()
    elif args.dump_err is True:
        dump_errors_stat()
    elif args.dump_lat is True:
        dump_latency_histogram()
    elif args.save is True:
        save_udc_stat()
    elif args.init is True:
        init_udc_stat()
    elif args.reset is True:
        reset_udc_stat()

if __name__ == '__main__':
    main()