#!/usr/bin/env python
# coding: utf-8

import os, sys, re, json, argparse
import collections, itertools

def dump_path(args):
    path = getattr(args, 'core', None)
    if path is not None:
        return path

    path = os.getenv('SHOEBOX')
    if path is not None:
        return path

    return './sydcore'

def dump_json(obj, fmt = None):
    if fmt is not None:
        try:
            formatted = fmt.format(**obj)
        except (KeyError, AttributeError, TypeError):
            formatted = ''
        sys.stdout.write(formatted + "\n")
    else:
        json.dump(obj, sys.stdout, sort_keys = True,
                  indent = 4, separators = (',', ': '))
        sys.stdout.write('\n')

def format_json(objs, fmt = None):
    if fmt is None:
        return objs
    p = []
    for obj in objs:
        try:
            formatted = fmt.format(**obj)
        except (KeyError, AttributeError, TypeError):
            formatted = "{id}".format(**obj)
        p.append(formatted)

    return p

def match_event(event, pattern = None, match_format = None):
    if pattern is None:
        return True
    try:
        return pattern.match(match_format.format(**event))
    except (KeyError, AttributeError, TypeError):
        return False

def match_any(patterns, string, flags = 0):
    for p in patterns:
        if p.match(string) is not None:
            return True
    return False

class ShoeBox:
    FORMATS_SUPPORTED = (1,)

    def __init__(self, dump = 'sydcore'):
        self.dump  = dump

        self.fmt  = None
        self.head = None

    def __enter__(self):
        self.fd = os.open(self.dump, os.O_RDONLY|os.O_NOFOLLOW|os.O_NOATIME)
        self.fp = os.fdopen(self.fd, 'r')
        self.check_format()

        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.fp.close()
        if exc_type is not None:
            return False # Raise the exception
        return True

    @property
    def path(self):
        return os.path.abspath(self.dump)

    @property
    def events(self):
        return self.read_events()

    @property
    def child(self):
        return self.match("startup:[0-9]+", "{event_name}:{pid}", 1)

    @property
    def exit(self):
        return self.match("exit", "{event_name}")

    @property
    def child_pid(self):
        c = self.child
        if c:
            return c['pid']

    def last(self, limit = 1, fmt = None):
        if limit > 0:
            limit = -limit
        self.rewind()
        return list(format_json(self.read_events(limit = limit), fmt))

    def _parse_pid(self, pid):
        if pid == 'child':
            return self.child_pid
        else:
            return int(pid)

    def search(self, **kwargs):
        self.rewind()
        events = list(self.events)

        if 'pid' in kwargs:
            pids = kwargs['pid']

            if isinstance(pids, str):
                pids = [pids,]

            pids = set(map(self._parse_pid, pids))

            def _filter_pid(event):
                if 'pid' not in event:
                    return False
                if event['pid'] not in pids:
                    return False
                return True
            events = filter(_filter_pid, events)
        if kwargs.get('sort', True):
            events = sorted(events, key = lambda event: event['id'])

        return tuple(events)

    def oops(self):
        return self.match("assert", "{event_name}")

    def match(self, pattern, match_format, limit = 0):
        self.rewind()

        p = re.compile(pattern, re.I|re.U)
        events = []

        for event in list(self.events):
            try:
                m = p.match(match_format.format(**event))
            except (KeyError, AttributeError, TypeError):
                continue
            if m:
                events.append(event)
            if limit > 0 and len(events) >= limit:
                break

        if not events:
            return None
        if limit == 1:
            return events.pop(0)
        return events

    def check_format(self):
        line = self.fp.readline()
        obj  = json.loads(line)

        if 'id' not in obj:
            self.fp.close()
            raise NotImplementedError("missing id attribute")
        elif obj['id'] != 0:
            self.fp.close()
            raise NotImplementedError("invalid id attribute `%r' for format check" % obj['id'])
        elif 'shoebox' not in obj:
            self.fp.close()
            raise NotImplementedError("missing shoebox attribute")
        elif obj['shoebox'] not in ShoeBox.FORMATS_SUPPORTED:
            self.fp.close()
            raise NotImplementedError("unsupported shoebox format `%r'" % obj['shoebox'])

        self.fmt  = obj['shoebox']
        self.head = self.fp.tell()

    def rewind(self):
        self.fp.seek(self.head, os.SEEK_SET)

    def readlines(self, limit = 0):
        if limit == 0:
            for line in self.fp:
                yield line
        elif limit > 0:
            for line in self.fp:
                yield line
                limit -= 1
                if limit == 0:
                    break
        else: # limit < 0
            self.fp.seek(0, os.SEEK_END)
            fpos = self.fp.tell()
            self.fp.seek(max(fpos - (2048 * abs(limit)), self.head), os.SEEK_SET)
            for line in self.fp:
                yield line
                limit += 1
                if limit == 0:
                    break

    def read_events(self, limit = 0):
        for json_line in self.readlines(limit):
            if not json_line or not json_line.startswith('{'):
                continue
            try:
                obj = json.loads(json_line)
            except TypeError as err:
                sys.stderr.write("WTF? %r\n" % json_line)
                raise
            except ValueError as err:
                sys.stderr.write("Unable to parse JSON: %r\n" % err)
                match = re.search('char (?P<char>[0-9]+)', ' '.join(err.args))
                if match is not None:
                    char = int(match.group('char'))
                    bh = max([char - 10, 0])
                    eh = min([char + 10, len(json_line)])
                    hl = json_line[:bh] + '\033[1m\033[91m' + json_line[bh:eh] + '\033[0m' + json_line[eh:]
                else:
                    hl = json_line
                sys.stderr.write(hl)
                raise
            yield obj

    def tree(self, pid, pattern, match_format, quick = False):
        if pid < 0:
            pid = self.child_pid

        events = []
        parents = set()
        for event in self.events:
            if 'pid' not in event:
                continue
            if event['pid'] != pid:
                continue
            events.append(event)

            if 'process' in event:
                if 'stat' in event['process']:
                    if event['process']['stat'] is not None:
                        if 'errno' in event['process']['stat']:
                            pass # TODO: warn
                        else:
                            parents.add(event['process']['stat']['ppid'])
                    elif 'syd' in event['process']:
                        if event['process']['syd'] is None:
                            pass # TODO: warn
                        else:
                            parents.add(event['process']['syd']['ppid'])

        for ppid in parents:
            self.rewind()
            events += self.tree(ppid, None, None, True)

        if quick:
            return events

        events_out = [event for event in events if match_event(event, pattern, match_format)]
        return sorted(events_out, key = lambda event: event['id'])

class syd:
    class process:
        def count(events):
            p = set()
            for e in events:
                if 'pid' in e:
                    p.add(e['pid'])
            return len(p)

        def exit_status(events, pid):
            event = syd.match.first(events, "%d:wait:EXIT:[0-9]+" % pid, '{pid}:{event_name}:{ptrace[name]}:{ptrace[msg]}')
            if event is None:
                return -1
            return event['ptrace']['msg']

        def sydbox_child(events):
            event = syd.match.first(events, 'true:[0-9]+', '{process[syd][flag_SYDBOX_CHILD]}:{pid}')
            if event is None:
                return -1
            return event['pid']

def command_debug(args):
    from IPython.config.loader import Config
    try:
        get_ipython
    except NameError:
        nested = 0
        cfg = Config()
        prompt_config = cfg.PromptManager
        prompt_config.in_template = 'In <\\#>: '
        prompt_config.in2_template = '   .\\D.: '
        prompt_config.out_template = 'Out<\\#>: '
    else:
        print("Running nested copies of IPython.")
        print("The prompts for the nested copy have been modified")
        cfg = Config()
        nested = 1

    from IPython.terminal.embed import InteractiveShellEmbed

    ipshell = InteractiveShellEmbed(config=cfg,
                                    banner1 = 'Where else would you hide your money?',
                                    exit_msg = 'Leaving debugger, back to occupying the world.')

    sydcore = dump_path(args)
    if not os.path.exists(sydcore):
        print("'%s' does not exist (specify path with shoebox --core /path/to/sydcore debug)" % sydcore);
        sys.exit(1)

    with ShoeBox(sydcore) as shoebox:
        ipshell()

def command_exec(args):
    code = sys.stdin.read()

    print(dump_path(args))
    print(os.getenv("SHOEBOX"))
    if args.scope == 'line':
        with ShoeBox(dump_path(args)) as sb:
            for event in sb.read_events():
                exec(code, {}, {'event': event, 'syd': syd})
    else:
        events = []
        with ShoeBox(dump_path(args)) as sb:
            for event in sb.read_events():
                events.append(event)
        exec(code, {}, {'events': events, 'syd': syd})

def command_show(args):
    if args.pattern is None:
        pattern = None
        match_format = None
    else:
        pattern = re.compile(args.pattern, re.I|re.U)
        match_format = args.match

    limit  = abs(args.limit_match)
    events = list() # TODO: use a set + frozenset(event.items())
    with ShoeBox(dump_path(args)) as sb:
        for event in sb.read_events(args.limit_event):
            if match_event(event, pattern, match_format):
                events.append(event)
                if limit > 0: # limit == 0 means no limit!
                    limit -= 1
                    if limit == 0:
                        break


    # Only integer types are supported: {id}, {pid} etc.
    # It's a bit "clever", though: {id}{pid} also works...
    if args.sort: # Only integer types are supported
        events = sorted(events, key = lambda event: int(args.sort.format(**event)), reverse = args.reverse)
    for event in events:
        dump_json(event, args.format)

def command_tree(args):
    if args.pattern is None:
        pattern = None
        match_format = None
    else:
        pattern = re.compile(args.pattern, re.I|re.U)
        match_format = args.match

    with ShoeBox(dump_path(args)) as sb:
        events = sb.tree(args.pid, pattern, match_format)
        for event in events:
            dump_json(event, args.format)

def main():
    parser = argparse.ArgumentParser(prog='shoebox',
                                     description='Pink hiding in a shoe box',
                                     formatter_class=argparse.RawDescriptionHelpFormatter,
                                     # usage='%(prog)s [options] {command [arg...]}',
                                     epilog='''
Hey you, out there on the road,
Always doing what you're told,
Can you help me?

Send bug reports to "alip@exherbo.org"
Attaching poems encourages consideration tremendously.''')
    parser.add_argument('-V', '--version', action = 'version', version = '%(prog)s-@SHOEBOX_VERSION@')
    parser.add_argument('-c', '--core', default = None, metavar = 'SYDCORE', help = 'Path to sydbox core dump')

    subparser = parser.add_subparsers(help = 'command help')

    parser_debug = subparser.add_parser('debug', help = 'Fire debug shell (requires IPython-2)')
    parser_debug.set_defaults(func = command_debug)

    parser_exec = subparser.add_parser('exec', help = 'Execute python code on standart input')
    parser_exec.add_argument('-s', '--scope', default = 'line',
        help = 'Scope: "line" execute code for each "event" (default), "all" execute code once for all "events"')
    parser_exec.set_defaults(func = command_exec)

    parser_show = subparser.add_parser('show', help = 'Show dump')
    parser_show.add_argument('-m', '--match', default = '{pid}', help = 'Match format')
    parser_show.add_argument('-p', '--pattern', help = 'Match pattern (regex)')
    parser_show.add_argument('-f', '--format', default = None, help = 'Format string')
    parser_show.add_argument('-l', '--limit-match', default = 0, type = int, help = 'Limit matches')
    parser_show.add_argument('-L', '--limit-event', default = 0, type = int, help = 'Limit events')
    parser_show.add_argument('-s', '--sort', default = '{id}', help = 'Sort events by an integer value (id, pid etc.)')
    parser_show.add_argument('-r', '--reverse', action='store_true', default = False, help = 'Sort in reverse')
    parser_show.set_defaults(func = command_show)

    parser_tree = subparser.add_parser('tree', help = 'Show process tree')
    parser_tree.add_argument('-m', '--match', default = '{event_name}', help = 'Match format')
    parser_tree.add_argument('-p', '--pattern', help = 'Match pattern (regex)')
    parser_tree.add_argument('-f', '--format', default = None, help = 'Format string')
    parser_tree.add_argument('pid', type = int, metavar = 'PID',
                             default = -1, nargs = '?',
                             help = 'PID to match, default: detect eldest child')
    parser_tree.set_defaults(func = command_tree)

    args = parser.parse_args()
    if getattr(args, 'func', None):
        return args.func(args)
    else:
        parser.print_help()

if __name__ == '__main__':
    sys.exit(main())
