# SPDX-License-Identifier: GPL-2.0

"""
Record monitored data access patterns.
"""

import argparse
import os
import signal

import _damo_fmt_str
import _damo_records
import _damo_subproc
import _damo_sysinfo
import _damon
import _damon_args
import _damon_modules

class DataForCleanup:
    kdamonds_idxs = None
    orig_kdamonds = None
    record_handle = None

data_for_cleanup = DataForCleanup()

cleaning = False

def cleanup_exit(exit_code):
    global cleaning
    if cleaning == True:
        return
    cleaning = True
    if data_for_cleanup.kdamonds_idxs != None:
        # ignore returning error, as kdamonds may already finished
        _damon.turn_damon_off(data_for_cleanup.kdamonds_idxs)
        err = _damon.stage_kdamonds(data_for_cleanup.orig_kdamonds)
        if err:
            print('failed restoring previous kdamonds setup (%s)' % err)

    if data_for_cleanup.record_handle:
        _damo_records.finish_recording(data_for_cleanup.record_handle)

    exit(exit_code)

def sighandler(signum, frame):
    print('\nsignal %s received' % signum)
    cleanup_exit(signum)

def handle_args(args):
    if not args.deducible_target:
        if _damon.any_kdamond_running() or \
                _damo_records.can_record_from_damon_stat():
            args.deducible_target = 'ongoing'

    args.output_permission, err = _damo_records.parse_file_permission_str(
            args.output_permission)
    if err != None:
        print('wrong --output permission (%s) (%s)' %
                (args.output_permission, err))
        exit(1)

    # backup duplicate output file
    if os.path.isfile(args.out):
        os.rename(args.out, args.out + '.old')


    if 'mem_footprint' in args.do_record:
        footprint_file_path = '%s.mem_footprint' % args.out
        if os.path.isfile(footprint_file_path):
            os.rename(footprint_file_path, footprint_file_path + '.old')

def get_ongoing_kdamonds(sysfs_kdamonds):
    if _damon.any_kdamond_running():
        return sysfs_kdamonds

    kdamonds, err = _damon_modules.damon_stat_kdamonds()
    if err is None:
        return kdamonds

    print('DAMON is not turned on via sysfs.')
    print('DAMON_STAT is also unavailable (%s).' % err)
    print('No way to proceed for ongoing kdamonds recording.')
    exit(1)

def tracepoints_from_args(args):
    if args.snapshot is not None:
        return None
    do_record = args.do_record
    if not 'access' in do_record and not 'damon_trace_events' in do_record:
        return None
    if 'damon_trace_events' in args.do_record:
        return _damo_sysinfo.available_damon_tracepoints()

    if args.schemes_target_regions is True:
        return [_damo_records.perf_event_damos_before_apply]
    return [_damo_records.perf_event_damon_aggregated,
            _damo_records.perf_event_damon_monitor_intervals_tune]

def snapshot_requests_from_args(args, kdamonds):
    ''' Returns snapshot_request and error '''
    if not 'access' in args.do_record or args.snapshot is None:
        return None, None

    tried_regions_of = None
    if args.schemes_target_regions:
        tried_regions_of = []
        for kidx, kd in enumerate(kdamonds):
            for cidx, ctx in enumerate(kd.contexts):
                for sidx, scheme in enumerate(ctx.schemes):
                    tried_regions_of.append([kidx, cidx, sidx])

    dfilters, err = _damon_args.damos_options_to_filters(
            args.snapshot_damos_filter)
    if err is not None:
        return None, 'wrong --snapshot_damos_filters (%s)' % err

    record_filter, err = _damo_records.args_to_filter(args)
    if err is not None:
        return None, err

    return _damo_records.SnapshotRequest(
            tried_regions_of=tried_regions_of, snapshot_damos_filters=dfilters,
            record_filter=record_filter, total_sz_only=False,
            dont_merge_regions=False), None

def mk_handle(args, kdamonds, monitoring_intervals):
    tracepoints = tracepoints_from_args(args)
    snapshot_request, err = snapshot_requests_from_args(args, kdamonds)
    if err is not None:
        print(err)
        cleanup_exit(1)
    if snapshot_request is not None:
        snapshot_interval_sec = _damo_fmt_str.text_to_sec(
                args.snapshot[0])
        snapshot_count = _damo_fmt_str.text_to_nr(args.snapshot[1])
    else:
        snapshot_interval_sec = None
        snapshot_count = None

    output_flush_sec = _damo_fmt_str.text_to_sec(args.output_flush_sec)

    do_records = args.do_record
    damon_tracer = args.damon_tracer
    if damon_tracer is not None:
        tracer = damon_tracer
        if tracer == 'perf' and args.perf_path is not None:
            tracer = args.perf_path
        if not _damo_subproc.avail_cmd(tracer):
            print('--damon_tracer (%s) is unavailable' % tracer)
            cleanup_exit(1)
    else:
        perf_path = 'perf'
        if args.perf_path is not None:
            perf_path = args.perf_path
        if _damo_subproc.avail_cmd(perf_path):
            damon_tracer = 'perf'
        elif _damo_subproc.avail_cmd('trace-cmd'):
            damon_tracer = 'trace-cmd'
            if 'cpu_profile' in do_records:
                do_records.remove('cpu_profile')
        else:
            print('Please install trace-cmd or perf and retry')
            cleanup_exit(1)

    handle = _damo_records.RecordingHandle(
            # for access pattern monitoring
            damon_tracer=damon_tracer,
            tracepoints=tracepoints, file_path=args.out,
            file_format=args.output_type,
            file_permission=args.output_permission,
            monitoring_intervals=monitoring_intervals,
            # for perf profile
            do_profile='cpu_profile' in args.do_record,
            # for children processes recording and memory footprint
            kdamonds=kdamonds, add_child_tasks=args.include_child_tasks,
            record_mem_footprint='mem_footprint' in args.do_record,
            record_vmas='vmas' in args.do_record,
            record_proc_stats='proc_stats' in args.do_record,
            timeout=args.timeout, snapshot_request=snapshot_request,
            snapshot_interval_sec=snapshot_interval_sec,
            snapshot_count=snapshot_count,
            max_seconds_per_file=output_flush_sec)
    if args.perf_path is not None:
        err = handle.set_perf_path(args.perf_path)
        if err is not None:
            print('--perf_path handling fail (%s)' % err)
            cleanup_exit(1)
    if 'damon_trace_events' in args.do_record:
        handle.leave_damon_trace_rawfile = True

    return handle

def main(args):
    global data_for_cleanup

    _damon.ensure_root_and_initialized(args)

    handle_args(args)

    # Setup for cleanup
    data_for_cleanup.orig_kdamonds = _damon.current_kdamonds()
    signal.signal(signal.SIGINT, sighandler)
    signal.signal(signal.SIGTERM, sighandler)

    # Now the real works
    if _damon_args.is_ongoing_target(args):
        kdamonds = get_ongoing_kdamonds(data_for_cleanup.orig_kdamonds)
        # TODO: Support multiple kdamonds, multiple contexts
        monitoring_intervals = kdamonds[0].contexts[0].intervals
    else:
        err, kdamonds = _damon_args.turn_damon_on(args)
        if err:
            print('could not turn DAMON on (%s)' % err)
            cleanup_exit(-2)
        data_for_cleanup.kdamonds_idxs = ['%d' % idx
                for idx, k in enumerate(kdamonds)]
        # TODO: Support multiple kdamonds, multiple contexts
        monitoring_intervals = kdamonds[0].contexts[0].intervals
        now_kdamonds = _damon.current_kdamonds()
        kdamonds[0].pid = now_kdamonds[0].pid

    record_handle = mk_handle(args, kdamonds, monitoring_intervals)

    data_for_cleanup.record_handle = record_handle

    if record_handle.will_take_awhile():
        print('Press Ctrl+C to stop')
    _damo_records.start_recording(record_handle)
    cleanup_exit(0)

def set_argparser(parser):
    _damon_args.set_argparser(parser, add_record_options=True, min_help=True)
    parser.add_argument('--output_type',
                        choices=_damo_records.file_types,
                        default=_damo_records.file_type_json_compressed,
                        help='output file\'s type')
    parser.add_argument('--output_permission', type=str, default='600',
                        help='permission of the output file')
    parser.add_argument(
            '--output_flush_sec', type=str, default='3600',
            help='intermediate output files flush duration in seconds')
    parser.add_argument('--perf_path', type=str, help='path of perf tool')
    parser.add_argument('--exclude_child_tasks', action='store_false',
                        dest='include_child_tasks',
                        help='do not record access of child processes')
    parser.add_argument('--include_child_tasks', action='store_true',
                        help='record accesses of child processes')
    parser.add_argument('--schemes_target_regions', action='store_true',
                        help='record schemes tried to be applied regions')
    parser.add_argument('--snapshot', metavar=('<delay>', '<count>'), nargs=2,
                        help='record accesses as snapshots')
    parser.add_argument('--timeout', type=float, metavar='<seconds>',
                        help='stop recording after the given seconds')
    parser.add_argument('--do_record', nargs='+',
                        default=['access', 'cpu_profile', 'mem_footprint',
                                 'vmas', 'proc_stats'],
                        choices=['access', 'cpu_profile', 'mem_footprint',
                                 'vmas', 'proc_stats', 'damon_trace_events'],
                        help='what to do record')
    parser.add_argument('--damon_tracer', metavar='cmd',
                        choices=['perf', 'trace-cmd'],
                        # tracer to use.  Hide this as this is an experimental
                        # option.
                        help=argparse.SUPPRESS)
    _damo_records.set_snapshot_damos_filters_option(parser)
    _damo_records.set_filter_argparser(parser)
    return parser
