#!/usr/bin/env python3

# ----------------------------------------------------------------------
# File: eos-mdstat
# Author: Andreas-Joachim Peters - CERN
# ----------------------------------------------------------------------

# ************************************************************************
# * EOS - the CERN Disk Storage System                                   *
# * Copyright (C) 2024 CERN/Switzerland                                  *
# *                                                                      *
# * This program is free software: you can redistribute it and/or modify *
# * it under the terms of the GNU General Public License as published by *
# * the Free Software Foundation, either version 3 of the License, or    *
# * (at your option) any later version.                                  *
# *                                                                      *
# * This program is distributed in the hope that it will be useful,      *
# * but WITHOUT ANY WARRANTY; without even the implied warranty of       *
# * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the        *
# * GNU General Public License for more details.                         *
# *                                                                      *
# * You should have received a copy of the GNU General Public License    *
# * along with this program.  If not, see <http://www.gnu.org/licenses/>.*
# ************************************************************************

import sys
import time
import json
import os

mdcost1 = {}
mdcost2 = {}

shortnames = {"OpenRead":"open:r","OpenWrite":"open:w","OpenWriteTruncate":"open:t", "OpenWriteCreate":"open:c","OpenProc":"cli","Eosxd::prot::LS":"fuse:ls","Eosxd::prot::STAT":"fuse:stat","Eosxd::prot::SET":"fuse:set","Eosxd::ext::LS-Entry":"fuse:entry","IdMap":"generic:op","Delay":"delay","Stall":"stall"}

now = int(time.time())

def dict2prom(data_dict):
    prom = ""
    prefix = "eos"
    category = "md"
    subcategory = "user"
    metric = f"{prefix}_{category}_{subcategory}"

    prom += f"# HELP {metric} User meta-data service usage derived from eos ns stat\n"
    prom += f"# TYPE {metric} gauge\n"

    for key, value in data_dict.items():
        if key == "uid":
            continue

        for tag, val in value.items():
            if tag == "uid":
                continue

            prom += f"{prefix}_{category}_{subcategory}{{uid=\"{key}\", measure=\"{tag}\"}} {val}\n"

    return prom

def read_textfile(file_path, mdcost):
    with open(file_path, 'r') as file:
        for line in file:
            dict = {}
            parts = line.strip().split()
            for item in parts:
                key, value = item.split('=', 1)
                key = key.strip()
                value = value.strip()
                # Store key-value pair in the dictionary
                dict[key] = value
            if 'uid' in dict:
                if dict['uid'] == "all":
                    continue;

                if dict['uid'] not in mdcost:
                    mdcost[dict['uid']] = {}

                if dict['uid'] == "all":
                    if dict['gid'] == "all":
                        if 'cmd' in dict:
                            if dict['cmd'].startswith(("Stall::threads::","Delay::threads::")):
                                if dict['cmd'].endswith(("ms")):
                                    continue
                                n=dict['total']
                                elements = dict['cmd'].split("::")
                                user = elements[-1]
                                if user not in mdcost:
                                    mdcost[user] = {}
                                mdcost[user][elements[0]]=n
                    continue
                if 'cmd' in dict:
                    user=dict['uid']
                    if dict['cmd'].startswith(("OpenRead","OpenWrite","OpenWriteCreate","OpenProc","Eosxd::prot::LS","Eosxd::prot::STAT","Eosxd::prot::SET","Eosxd::ext::LS-Entry","IdMap")):
                        cmd=dict['cmd']
                        n=dict['total']
                        mdcost[user][cmd]=n
    return

arguments = sys.argv

f1 = "/var/eos/md/mdstats.1"
f2 = "/var/eos/md/mdstats.0"

tofile=""
tofiletmp=""
promfile=""
promfiletmp=""

if len(arguments) > 2:
    f1 = arguments[1]
    f2 = arguments[2]
else:
    tofile="/var/eos/md/md-report.json"
    tofiletmp = tofile + ".tmp"
    promfile="/var/eos/md/md-report.prom"
    promfiletmp= promfile + ".tmp"

    read_textfile(f1, mdcost1)
read_textfile(f2, mdcost2)

mddiff = {}
for user,cost in mdcost2.items():
    mddiff[user] = {}
    for key,value in cost.items():
        if user in mdcost1:
            if key in mdcost1[user]:
                if int(value) >= int(mdcost1[user][key]):
                   newvalue=int(value)-int(mdcost1[user][key])
            else:
                newvalue=int(value)
        else:
            newvalue=int(value)
        mddiff[user][shortnames[key]]=newvalue
    for item in shortnames:
        if shortnames[item] not in mddiff[user]:
            mddiff[user][shortnames[item]]=0
        mddiff[user]["unixtime"]=now

jsondump = json.dumps(mddiff, indent=4)
if len(tofile):
    with open(tofiletmp, "w") as file:
        file.write(jsondump)
        os.rename(tofiletmp,tofile)
else:
    print(jsondump)

if len(promfile):
    with open(promfiletmp, "w") as file:
        file.write(dict2prom(mddiff))
        os.rename(promfiletmp, promfile)
