from __future__ import absolute_import
from nilmdb.printf import *
import nilmdb.client
import nilmdb.layout
import sys

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

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

    group = cmd.add_argument_group("Output format")
    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("-c", "--count", action="store_true",
                       help="Just output a count of matched data points")

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.annotate:
        printf("# path: %s\n", self.args.path)
        printf("# layout: %s\n", layout)
        printf("# start: %s\n", self.time_string(self.args.start))
        printf("# end: %s\n", self.time_string(self.args.end))

    printed = False
    for dataline in self.client.stream_extract(self.args.path,
                                               self.args.start,
                                               self.args.end,
                                               self.args.count):
        if self.args.bare and not self.args.count:
            # Strip timestamp (first element).  Doesn't make sense
            # if we are only returning a count.
            dataline = ' '.join(dataline.split(' ')[1:])
        print dataline
        printed = True
    if not printed:
        if self.args.annotate:
            printf("# no data\n")
        return 2

    return 0
