#!/usr/bin/python3
#
# Copyright (c) 2017-2019 NVIDIA CORPORATION. All rights reserved.
# This file is part of webloader (see TBD).
# See the LICENSE file for licensing terms (BSD-style).
#

import argparse
import os
import sys

from tarproclib import reader

parser = argparse.ArgumentParser("Dump the first matching file from a tar file.")
parser.add_argument("-S", "--select", default="True",
                    help="expression that must evaluate to true before saving")
parser.add_argument("-v", "--verbose", action="store_true",
                    help="output more information")
parser.add_argument("-f", "--field", default="img;image;ppm;pgm;pbm;jpeg;jpg;png",
                    help="field to be selected")
parser.add_argument("input", default="-", nargs="?")
args = parser.parse_args()


class Record(object):
    def __init__(self, d):
        self.__dict__ = d

    def __getitem__(self, attr):
        return self.__dict__[attr]

    def __setitem__(self, attr, value):
        self.__dict__[attr] = value


def dprint(*args, **kw):
    print(*args, file=sys.stderr, **kw)


def getfirst(a, keys, default=None):
    if isinstance(keys, str):
        keys = keys.split(";")
    for k in keys:
        result = a.get(k)
        if result is not None:
            return result, k
    return default


def parse_field_spec(fields):
    if isinstance(fields, str):
        fields = fields.split()
    return [field.split(";") for field in fields]


for sample in reader.TarIterator(args.input):
    assert isinstance(sample, dict), sample
    if args.verbose:
        print("#", sample.get("__key__"), args.field, sample.keys(), file=sys.stderr)
    _ = Record(sample)
    if not eval(args.select):
        continue
    result, key = getfirst(sample, args.field)
    if result is not None:
        dprint(sample["__key__"] + "." + key)
        assert isinstance(result, bytes)
        with os.fdopen(sys.stdout.fileno(), "wb", closefd=False) as stdout:
            stdout.write(result)
            stdout.flush()
        break
