import argparse
import os
import logging
import sys

from scripts.tools.trace import trace_utils
from data.chum import chumpy

MAX_DURATION_SEC = 2

def uri_splitter(uri, max_duration_sec):
    """Parse the uri and split it into smaller chunks of time

    Chrome cannot handle large amounts of ftrace data, so we split the desired
    range into several ftrace sections of a maximum size"""
    uri_proto = chumpy.parseChumUriToProto(uri)
    store, range_ = chumpy.parseChumUri(uri)
    start = range_.start_time
    end = range_.end_time
    max_duration_ns = int(1e9 * max_duration_sec)
    for chunk_start in range(start, end, max_duration_ns):
        uri_proto.start_time_ns = chunk_start
        # All chunks but the last one should have max_duration_ns
        uri_proto.end_offset_ns = min(max_duration_ns, end - chunk_start)
        yield uri_proto

NS = 1000000000

def main(uri, out_dir=None, include_bpf=True, include_user=True, machine='pcu', pid=0, max_chunk_duration=MAX_DURATION_SEC):
    pid_map = trace_utils.getPidMap(uri, machine)
    for uri_proto in uri_splitter(uri, max_chunk_duration):
        uri = chumpy.renderChumUri(uri_proto)
        ftrace, _events = trace_utils.getFtrace(
            uri,
            include_bpf,
            False, # include_ftrace
            include_user,
            machine,
            pid_map,
            pid)
        if not ftrace:
            logging.warning('No data for uri: %s', uri)
            continue
        if out_dir:
            s = uri_proto.start_time_ns // NS
            ns = uri_proto.start_time_ns % NS
            filename = os.path.join(out_dir, 'trace_{:d}_{:09d}.log'.format(s, ns))
            open(filename, 'w').write(''.join(ftrace))
            logging.info('Wrote %s (%d bytes)', filename, sum(map(len, ftrace)))
        else:
            sys.stdout.write(''.join(ftrace))


if __name__ == '__main__':
    TRUE = ['true', 'on', 'yes']
    FALSE = ['false', 'off','no']
    BOOL = TRUE + FALSE

    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument('uri',
        help='Chum URI of range to generate')
    parser.add_argument('--out_dir',
        help='Dir to save files')
    parser.add_argument('--bpf',
        choices=BOOL, default='true',
        help='Include bpf (scheduler) events')
    parser.add_argument('--user',
        choices=BOOL, default='true',
        help='Include user space trace events')
    parser.add_argument('--machine',
        help='Machine to run trace on',
        default='pcu')
    parser.add_argument('--max_chunk_duration',
        type=float, default=MAX_DURATION_SEC,
        help='Maximum length of time per single ftrace file in seconds (default {}s)'.format(MAX_DURATION_SEC))
    parser.add_argument('--pid',
        type=int, default=0,
        help='filter by process ID')

    logging.basicConfig(level=logging.INFO)

    args = parser.parse_args()
    args = vars(args)
    include_bpf = args.pop('bpf') in TRUE
    include_user = args.pop('user') in TRUE

    try:
        main(**args,
            include_bpf=include_bpf,
            include_user=include_user)
    except IOError as e:
        if e.errno == 32: # Broken pipe
            pass
        else:
            raise
