#!/bin/python

#
# Copyright (c) 2021 Vasile Vilvoiu <vasi@vilvoiu.ro>
#
# This is free software; you can redistribute it and/or modify
# it under the terms of the MIT license. See LICENSE for details.
#

import argparse
import re
import json
import s2util
import sys
import munch


class ProjectRuleConcatAction(argparse.Action):
    def __call__(self, parser, namespace, values, option_string=None):
        current = getattr(namespace, self.dest)
        new = {'action': 'project', 'expression': values}
        if current is None:
            current = [new]
        else:
            current.append(new)
        setattr(namespace, self.dest, current)


class RejectRuleConcatAction(argparse.Action):
    def __call__(self, parser, namespace, values, option_string=None):
        current = getattr(namespace, self.dest)
        new = {'action': 'reject', 'expression': values}
        if current is None:
            current = [new]
        else:
            current.append(new)
        setattr(namespace, self.dest, current)


class ProjectionIterator(s2util.ObjectIterator):
    def __init__(self, provider, rules, filter):
        super().__init__(provider)
        self.rules = rules
        self.filter = filter

    def project(self, obj):
        nobject = {}
        for key in obj.keys():
            project = True
            for rule in self.rules:
                kpass = re.search(rule['compiled'], key)
                if kpass:
                    if rule['action'] == 'reject':
                        project = False
                    elif rule['action'] == 'project':
                        project = True
                    else:
                        raise Exception('unknown action {}'.format(rule['action']))
            if project:
                nobject[key] = obj[key]
        return nobject

    def process(self, obj, archive_url, object_no):
        try:
            pred_eval = eval(self.filter, {'re': re}, {'r': munch.Munch.fromDict(obj)})
        except Exception as e:
            print('Predicate evaluation failed: {}'.format(e), file=sys.stderr)
            pred_eval = False
        if pred_eval:
            print(json.dumps(self.project(obj)))


def main():
    # Parse command line arguments
    parser = argparse.ArgumentParser(description='Filter-projector for semantic scholar entries')
    parser.add_argument('corpus', type=str, help='corpus directory URL (where manifest file is found)')
    parser.add_argument('--archive-regex', type=str, help='regex to match against archive filenames before processing',
                        default='s2[-]corpus[-][1-9]?[0-9]{3}[.]gz')
    parser.add_argument('--index', type=str, default=None, help='path to index used for iteration')
    parser.add_argument('-p', '--project', action=ProjectRuleConcatAction,
                        help='field projection rule (may be repeated)', dest='prules')
    parser.add_argument('-r', '--reject', action=RejectRuleConcatAction,
                        help='field rejection rule (may be repeated)', dest='prules')
    parser.add_argument('-f', '--filter', default='True',
                        help='filter expression (written in Python, the object is named "r", dot-accessible via munch)')
    args = parser.parse_args()

    # Load infex
    if args.index is not None:
        provider = s2util.IndexObjectProvider(s2util.SearchableIndex(args.index), args.corpus)
    else:
        archive_urls = s2util.fetch_archive_urls(args.corpus, args.archive_regex)
        provider = s2util.SequentialObjectProvider(archive_urls)

    # Compile projection regexp list
    projection_rules = []
    if args.prules is not None:
        for r in args.prules:
            rule = r.copy()
            rule['compiled'] = re.compile(rule['expression'])
            projection_rules.append(rule)

    # Compile filter expression
    filter_expr = compile(args.filter, 'expr', 'eval')

    # Filter and print to stdout
    ProjectionIterator(provider, projection_rules, filter_expr).iterate()


if __name__ == '__main__':
    main()
