#!/usr/bin/env python3
# SPDX-License-Identifier: CDDL-1.0
#
# Print out statistics for all zil stats. This information is
# available through the zil kstat.
#
# CDDL HEADER START
#
# The contents of this file are subject to the terms of the
# Common Development and Distribution License, Version 1.0 only
# (the "License").  You may not use this file except in compliance
# with the License.
#
# You can obtain a copy of the license at usr/src/OPENSOLARIS.LICENSE
# or https://opensource.org/licenses/CDDL-1.0.
# See the License for the specific language governing permissions
# and limitations under the License.
#
# When distributing Covered Code, include this CDDL HEADER in each
# file and include the License file at usr/src/OPENSOLARIS.LICENSE.
# If applicable, add the following below this CDDL HEADER, with the
# fields enclosed by brackets "[]" replaced with your own identifying
# information: Portions Copyright [yyyy] [name of copyright owner]
#
# This script must remain compatible with Python 3.6+.
#

import sys
import subprocess
import time
import copy
import os
import re
import signal
from collections import defaultdict
import argparse
from argparse import RawTextHelpFormatter

cols = {
	# hdr:       [size,      scale,      kstat name]
	"time":      [8,         -1,         "time"],
	"pool":      [12,        -1,         "pool"],
	"ds":        [12,        -1,         "dataset_name"],
	"obj":       [12,        -1,         "objset"],
	"cc":        [5,         1000,       "zil_commit_count"],
	"cwc":       [5,         1000,       "zil_commit_writer_count"],
	"cec":       [5,         1000,       "zil_commit_error_count"],
	"csc":       [5,         1000,       "zil_commit_stall_count"],
	"cSc":       [5,         1000,       "zil_commit_suspend_count"],
	"ic":        [5,         1000,       "zil_itx_count"],
	"iic":       [5,         1000,       "zil_itx_indirect_count"],
	"iib":       [5,         1024,       "zil_itx_indirect_bytes"],
	"icc":       [5,         1000,       "zil_itx_copied_count"],
	"icb":       [5,         1024,       "zil_itx_copied_bytes"],
	"inc":       [5,         1000,       "zil_itx_needcopy_count"],
	"inb":       [5,         1024,       "zil_itx_needcopy_bytes"],
	"idc":       [5,         1000,       "icc+inc"],
	"idb":       [5,         1024,       "icb+inb"],
	"iwc":       [5,         1000,       "iic+idc"],
	"iwb":       [5,         1024,       "iib+idb"],
	"imnc":      [6,         1000,       "zil_itx_metaslab_normal_count"],
	"imnb":      [6,         1024,       "zil_itx_metaslab_normal_bytes"],
	"imnw":      [6,         1024,       "zil_itx_metaslab_normal_write"],
	"imna":      [6,         1024,       "zil_itx_metaslab_normal_alloc"],
	"imsc":      [6,         1000,       "zil_itx_metaslab_slog_count"],
	"imsb":      [6,         1024,       "zil_itx_metaslab_slog_bytes"],
	"imsw":      [6,         1024,       "zil_itx_metaslab_slog_write"],
	"imsa":      [6,         1024,       "zil_itx_metaslab_slog_alloc"],
	"imc":       [5,         1000,       "imnc+imsc"],
	"imb":       [5,         1024,       "imnb+imsb"],
	"imw":       [5,         1024,       "imnw+imsw"],
	"ima":       [5,         1024,       "imna+imsa"],
	"se%":       [3,         100,        "imb/ima"],
	"sen%":      [4,         100,        "imnb/imna"],
	"ses%":      [4,         100,        "imsb/imsa"],
	"te%":       [3,         100,        "imb/imw"],
	"ten%":      [4,         100,        "imnb/imnw"],
	"tes%":      [4,         100,        "imsb/imsw"],
}

hdr = ["time", "ds", "cc", "ic", "idc", "idb", "iic", "iib",
	"imnc", "imnw", "imsc", "imsw"]

ghdr = ["time", "cc", "ic", "idc", "idb", "iic", "iib",
	"imnc", "imnw", "imsc", "imsw"]

cmd = ("Usage: zilstat [-hgdv] [-i interval] [-p pool_name]")

curr = {}
diff = {}
kstat = {}
ds_pairs = {}
pool_name = None
dataset_name = None
interval = 0
sep = "  "
gFlag = True
dsFlag = False

def prettynum(sz, scale, num=0):
	suffix = [' ', 'K', 'M', 'G', 'T', 'P', 'E', 'Z']
	index = 0
	save = 0

	if scale == -1:
		return "%*s" % (sz, num)

	# Rounding error, return 0
	elif 0 < num < 1:
		num = 0

	while num > scale and index < 5:
		save = num
		num = num / scale
		index += 1

	if index == 0:
		return "%*d" % (sz, num)

	if (save / scale) < 10:
		return "%*.1f%s" % (sz - 1, num, suffix[index])
	else:
		return "%*d%s" % (sz - 1, num, suffix[index])

def print_header():
	global hdr
	global sep
	for col in hdr:
		new_col = col
		if interval > 0 and cols[col][1] > 100:
			new_col += "/s"
		sys.stdout.write("%*s%s" % (cols[col][0], new_col, sep))
	sys.stdout.write("\n")

def print_values(v):
	global hdr
	global sep
	for col in hdr:
		val = v[cols[col][2]]
		if interval > 0 and cols[col][1] > 100:
			val = v[cols[col][2]] // interval
		sys.stdout.write("%s%s" % (
			prettynum(cols[col][0], cols[col][1], val), sep))
	sys.stdout.write("\n")

def print_dict(d):
	for pool in d:
		for objset in d[pool]:
			print_values(d[pool][objset])

def detailed_usage():
	sys.stderr.write("%s\n" % cmd)
	sys.stderr.write("Field definitions are as follows:\n")
	for key in cols:
		sys.stderr.write("%11s : %s\n" % (key, cols[key][2]))
	sys.stderr.write("\n")
	sys.exit(0)

def init():
	global pool_name
	global dataset_name
	global interval
	global hdr
	global curr
	global gFlag
	global sep

	curr = dict()

	parser = argparse.ArgumentParser(description='Program to print zilstats',
                                	 add_help=True,
					 formatter_class=RawTextHelpFormatter,
					 epilog="\nUsage Examples\n"\
				 		"Note: Global zilstats is shown by default,"\
						" if none of a|p|d option is not provided\n"\
				 		"\tzilstat -a\n"\
						'\tzilstat -v\n'\
						'\tzilstat -p tank\n'\
						'\tzilstat -d tank/d1,tank/d2,tank/zv1\n'\
						'\tzilstat -i 1\n'\
						'\tzilstat -s \"***\"\n'\
						'\tzilstat -f zcwc,zimnb,zimsb\n')

	parser.add_argument(
		"-v", "--verbose",
		action="store_true",
		help="List field headers and definitions"
	)

	pool_grp = parser.add_mutually_exclusive_group()

	pool_grp.add_argument(
		"-a", "--all",
		action="store_true",
		dest="all",
		help="Print all dataset stats"
	)

	pool_grp.add_argument(
		"-p", "--pool",
		type=str,
		help="Print stats for all datasets of a speicfied pool"
	)

	pool_grp.add_argument(
		"-d", "--dataset",
		type=str,
		help="Print given dataset(s) (Comma separated)"
	)

	parser.add_argument(
		"-f", "--columns",
		type=str,
		help="Specify specific fields to print (see -v)"
	)

	parser.add_argument(
		"-s", "--separator",
		type=str,
		help="Override default field separator with custom "
			 "character or string"
	)

	parser.add_argument(
		"-i", "--interval",
		type=int,
		dest="interval",
		help="Print stats between specified interval"
			 " (in seconds)"
	)

	parsed_args = parser.parse_args()

	if parsed_args.verbose:
		detailed_usage()

	if parsed_args.all:
		gFlag = False

	if parsed_args.interval:
		interval = parsed_args.interval

	if parsed_args.pool:
		pool_name = parsed_args.pool
		gFlag = False

	if parsed_args.dataset:
		dataset_name = parsed_args.dataset
		gFlag = False

	if parsed_args.separator:
		sep = parsed_args.separator

	if gFlag:
		hdr = ghdr

	if parsed_args.columns:
		hdr = parsed_args.columns.split(",")

		invalid = []
		for ele in hdr:
			if ele not in cols:
				invalid.append(ele)

		if len(invalid) > 0:
			sys.stderr.write("Invalid column definition! -- %s\n" % invalid)
			sys.exit(1)

	if pool_name and dataset_name:
		print ("Error: Can not filter both dataset and pool")
		sys.exit(1)

def FileCheck(fname):
	try:
		return (open(fname))
	except IOError:
		print ("Unable to open zilstat proc file: " + fname)
		sys.exit(1)

if sys.platform.startswith('freebsd'):
	# Requires py-sysctl on FreeBSD
	import sysctl

	def kstat_update(pool = None, objid = None):
		global kstat
		kstat = {}
		if not pool:
			file = "kstat.zfs.misc.zil"
			k = [ctl for ctl in sysctl.filter(file) \
				if ctl.type != sysctl.CTLTYPE_NODE]
			kstat_process_str(k, file, "GLOBAL", len(file + "."))
		elif objid:
			file = "kstat.zfs." + pool + ".dataset.objset-" + objid
			k = [ctl for ctl in sysctl.filter(file) if ctl.type \
				!= sysctl.CTLTYPE_NODE]
			kstat_process_str(k, file, objid, len(file + "."))
		else:
			file = "kstat.zfs." + pool + ".dataset"
			zil_start = len(file + ".")
			obj_start = len("kstat.zfs." + pool + ".")
			k = [ctl for ctl in sysctl.filter(file)
				if ctl.type != sysctl.CTLTYPE_NODE]
			for s in k:
				if not s or (s.name.find("zil") == -1 and \
					s.name.find("dataset_name") == -1):
					continue
				name, value = s.name, s.value
				objid = re.findall(r'0x[0-9A-F]+', \
					name[obj_start:], re.I)[0]
				if objid not in kstat:
					kstat[objid] = dict()
		                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                      