# SPDX-License-Identifier: GPL-2.0

"""
Show status of DAMON.
"""

import argparse
import collections
import json
import os

import _damo_fmt_str
import _damo_yaml
import _damon
import _damon_args

def read_kdamonds_from_file(input_file):
    # returns read kdamonds and error
    kdamonds, err = _damon_args.kdamonds_from_json_arg(input_file)
    if err is not None:
        return None, 'reading %s failed (%s)' % (input_file, err)
    return kdamonds, None

def pr_damon_parameters(input_file, format, raw_nr, omit_defaults):
    if input_file is None:
        kdamonds = _damon.current_kdamonds()
    else:
        kdamonds, err = read_kdamonds_from_file(input_file)
        if err is not None:
            print(err)
            exit(1)

    for k in kdamonds:
        for c in k.contexts:
            for s in c.schemes:
                if omit_defaults:
                    if s.access_pattern == _damon.DamosAccessPattern():
                        s.access_pattern = None
                    if s.quotas == _damon.DamosQuotas():
                        s.quotas = None
                    if s.watermarks == _damon.DamosWatermarks():
                        s.watermarks = None
                s.stats = None
                s.tried_regions = None

    pr_kdamonds(kdamonds, format, raw_nr, show_cpu=False,
                params_only=True)

def update_pr_schemes_stats(input_file, format, raw_nr, damos_stat_fields):
    if input_file is None:
        err = _damon.update_schemes_stats()
        if err:
            print(err)
            exit(1)
        kdamonds = _damon.current_kdamonds()
    else:
        kdamonds, err = read_kdamonds_from_file(input_file)
        if err is not None:
            print(err)
            exit(1)

    stats = []
    for kd_idx, kdamond in enumerate(kdamonds):
        for ctx_idx, ctx in enumerate(kdamond.contexts):
            for scheme_idx, scheme in enumerate(ctx.schemes):
                indices = {
                        'kdamond':kd_idx, 'context': ctx_idx, 'scheme': scheme_idx}
                stat_kvpair = scheme.stats.to_kvpairs(raw_nr)
                if damos_stat_fields:
                    filtered_stat_kvpair = collections.OrderedDict()
                    for k in damos_stat_fields:
                        filtered_stat_kvpair[k] = stat_kvpair[k]
                    stat_kvpair = filtered_stat_kvpair
                stats.append([indices, stat_kvpair])

    if format == 'json':
        print(json.dumps(stats, indent=4))
        return
    elif format == 'yaml':
        text, err = _damo_yaml.dump(stats)
        if err is not None:
            print('yaml dump fail (%s)' % err)
            exit(1)
        print(text)
        return

    for indices, stat_kvpair in stats:
        if len(stats) > 1:
            print('kdamond %d / context %d / scheme %d' % (indices['kdamond'],
                indices['context'], indices['scheme']))
        for k, v in stat_kvpair.items():
            if len(stat_kvpair.keys()) > 1:
                print('%s: %s' % (k, v))
            else:
                print('%s' % v)
        if len(stats) > 1:
            print()

def pr_kdamonds_summary(input_file, format, raw_nr, show_cpu):
    if input_file is None:
        kdamonds = _damon.current_kdamonds()
    else:
        kdamonds, err = read_kdamonds_from_file(input_file)
        if err is None:
            print(err)
            exit(1)
    summary = [k.summary_str(show_cpu, raw_number=raw_nr) for k in kdamonds]
    if format == 'json':
        print(json.dumps(summary, indent=4))
        return
    elif format == 'yaml':
        text, err = _damo_yaml.dump(summary)
        if err is not None:
            print('yaml dump fail (%s)' % err)
            exit(1)
        print(text)
        return
    for idx, line in enumerate(summary):
        print('%d\t%s' % (idx, line))

def pr_kdamonds(kdamonds, format, raw_nr, show_cpu, params_only=False,
                omit_damos_tried_regions=False, show_modules=False):
    if format == 'json':
        print(json.dumps([k.to_kvpairs(raw_nr) for k in kdamonds], indent=4))
        return
    elif format == 'yaml':
        text, err = _damo_yaml.dump([k.to_kvpairs(raw_nr) for k in kdamonds])
        if err is not None:
            print('yaml dump fail (%s)' % err)
            exit(1)
        print(text)
        return
    else:
        for idx, k in enumerate(kdamonds):
            print('kdamond %d' % idx)
            print(_damo_fmt_str.indent_lines(
                k.to_str(raw_nr, show_cpu, params_only,
                         omit_damos_tried_regions), 4))

    if not show_modules:
        return

    modules_dir = '/sys/module'
    if not os.path.isdir(modules_dir):
        return
    module_state = {}
    for module in os.listdir(modules_dir):
        if not module.startswith('damon_'):
            continue
        if module.startswith('damon_sample_'):
            continue
        params_dir = os.path.join(modules_dir, module, 'parameters')
        if not os.path.isdir(params_dir):
            continue
        for param in os.listdir(params_dir):
            if not param in ['enabled', 'enable']:
                continue
            with open(os.path.join(params_dir, param), 'r') as f:
                val = f.read().strip()
            if val == 'Y':
                state = 'on'
            else:
                state = 'off'
            module_state[module] = state
    for module in sorted(module_state.keys()):
        print('%s: %s' % (module, module_state[module]))

def main(args):
    _damon.ensure_root_and_initialized(args)

    if args.json and args.format is None:
        args.format = 'json'
    if args.format is None:
        args.format = 'report'

    if args.damon_params or args.damon_params_omit_defaults:
        return pr_damon_parameters(args.input_file, args.format, args.raw,
                                   args.damon_params_omit_defaults)

    if args.kdamonds_summary:
        return pr_kdamonds_summary(args.input_file, args.format, args.raw,
                                   args.show_cpu_usage)

    if args.damos_stats:
        return update_pr_schemes_stats(args.input_file, args.format, args.raw,
                                       args.damos_stat_fields)

    if args.input_file is None:
        update_tried_regions = True
        # tried regions are not printed when --omit_damos_tried_regions or
        # --json is set.  When --json is set, the output is formatted via
        # Damos.to_kvpairs(), which doesn't include tried regions information.
        # Updating tried regions for such cases makes no sense but just incur
        # system overhead and user delay, so don't do the update in the case.
        if args.omit_damos_tried_regions or args.format:
            update_tried_regions = False

        kdamonds, err = _damon.update_read_kdamonds(
                nr_retries=5, update_stats=True,
                update_tried_regions=update_tried_regions,
                update_quota_effective_bytes=True,
                # todo: set update_tuned_intervals as False if not needed
                do_update_tuned_intervals=True)
        if err != None:
            print('cannot update and read kdamonds: %s' % err)
            exit(1)
    else:
        kdamonds, err = read_kdamonds_from_file(args.input_file)
        if err is not None:
            print(err)
            exit(1)
    pr_kdamonds(kdamonds, args.format, args.raw, args.show_cpu_usage,
                args.omit_damos_tried_regions, show_modules=True)

def set_argparser(parser):
    parser.add_argument(
            '--format', choices=['json', 'yaml', 'report'],
            help='format of the output')
    parser.add_argument('--json', action='store_true', default=False,
                        # suppress help message in favor of --format.
                        # help='print output in json format')
                        help=argparse.SUPPRESS)
    parser.add_argument('--raw', action='store_true', default=False,
            help='print raw numbers')
    parser.add_argument('--show_cpu_usage', action='store_true', default=False,
            help='show cpu usage for kdamond')
    parser.add_argument('--kdamonds_summary', action='store_true',
            help='print kdamond summary only')
    parser.add_argument('--damos_stats', action='store_true',
            help='print DAMOS scheme stats only')
    parser.add_argument('--damos_stat_fields', metavar='<stat field name>',
            choices=['nr_tried', 'sz_tried', 'nr_applied', 'sz_applied',
                'qt_exceeds'], nargs='+',
            help='DAMOS stat fields to print')
    parser.add_argument('--omit_damos_tried_regions', action='store_true',
                        help='omit DAMOS tried regions')
    parser.add_argument('--damon_params', action='store_true',
            help='print entered DAMON parameters only')
    parser.add_argument('--damon_params_omit_defaults', action='store_true',
            help='print entered DAMON parameters only, omitting defaults')
    parser.add_argument(
            '--input_file', metavar='<file>', help=' '.join([
                'A json file containing the status of kdamonds to show.',
                'If this is not given, capture and show status of',
                'current kdamonds.']))
    _damon_args.add_interface_debug_args(parser)
    return parser
