from nilmdb.utils.printf import *
import nilmdb.client
import sys

def setup(self, sub):
    cmd = sub.add_parser("extract", help="Extract data",
                         description="""
                         Extract data from a stream.
                         """)
    cmd.set_defaults(verify = cmd_extract_verify,
                     handler = cmd_extract)

    group = cmd.add_argument_group("Data selection")
    group.add_argument("path",
                       help="Path of stream, e.g. /foo/bar",
                       ).completer = self.complete.path
    group.add_argument("-s", "--start", required=True,
                       metavar="TIME", type=self.arg_time,
                       help="Starting timestamp (free-form, inclusive)",
                       ).completer = self.complete.time
    group.add_argument("-e", "--end", required=True,
                       metavar="TIME", type=self.arg_time,
                       help="Ending timestamp (free-form, noninclusive)",
                       ).completer = self.complete.time

    group = cmd.add_argument_group("Output format")
    group.add_argument("-B", "--binary", action="store_true",
                       help="Raw binary output")
    group.add_argument("-b", "--bare", action="store_true",
                       help="Exclude timestamps from output lines")
    group.add_argument("-a", "--annotate", action="store_true",
                       help="Include comments with some information "
                       "about the stream")
    group.add_argument("-m", "--markup", action="store_true",
                       help="Include comments with interval starts and ends")
    group.add_argument("-T", "--timestamp-raw", action="store_true",
                       help="Show raw timestamps in annotated information")
    group.add_argument("-c", "--count", action="store_true",
                       help="Just output a count of matched data points")
    return cmd

def cmd_extract_verify(self):
    if self.args.start > self.args.end:
        self.parser.error("start is after end")

    if self.args.binary:
        if (self.args.bare or self.args.annotate or self.args.markup or
            self.args.timestamp_raw or self.args.count):
            self.parser.error("--binary cannot be combined with other options")

def cmd_extract(self):
    streams = self.client.stream_list(self.args.path)
    if len(streams) != 1:
        self.die("error getting stream info for path %s", self.args.path)
    layout = streams[0][1]

    if self.args.timestamp_raw:
        time_string = nilmdb.utils.time.timestamp_to_string
    else:
        time_string = nilmdb.utils.time.timestamp_to_human

    if self.args.annotate:
        printf("# path: %s\n", self.args.path)
        printf("# layout: %s\n", layout)
        printf("# start: %s\n", time_string(self.args.start))
        printf("# end: %s\n", time_string(self.args.end))

    printed = False
    if self.args.binary:
        printer = sys.stdout.buffer.write
    else:
        printer = lambda x: print(x.decode('utf-8'))
    bare = self.args.bare
    count = self.args.count
    for dataline in self.client.stream_extract(self.args.path,
                                               self.args.start,
                                               self.args.end,
                                               self.args.count,
                                               self.args.markup,
                                               self.args.binary):
        if bare and not count:
            # Strip timestamp (first element).  Doesn't make sense
            # if we are only returning a count.
            dataline = b' '.join(dataline.split(b' ')[1:])
        printer(dataline)
        printed = True
    if not printed:
        if self.args.annotate:
            printf("# no data\n")
        return 2

    return 0