#! /usr/bin/env python3
# -*- mode: Python; -*-


"""
vrt-rm-structs

Remove from VRT structures matching given criteria (annotation values) either
completely or by replacing their content with fixed strings.
"""


# TODO:
# - Support multiple attribute conditions, conjunctive or disjunctive,
#   or with full Boolean operations
# - Make attribute condition more flexible as regards value quoting
# - Retain values of specified positional attributes (requires processing
#   the positional-attributes comment)
# - Replace different attributes with different values


import re
import sys

from collections import defaultdict
from itertools import chain

from vrtargsoolib import InputProcessor


class StructureRemover(InputProcessor):

    DESCRIPTION = """
    Remove from VRT structures matching given criteria (annotation
    values) either completely or by replacing their content with fixed
    strings.
    """

    ARGSPECS = [
        ('--structure-name=STRUCT "text" -> struct_name',
         'remove structures STRUCT matching the given criteria'),
        ('--attribute-condition=ATTR_REGEXP -> attr_regex',
         """remove structures whose annotation attributes match regular
         expression ATTR_REGEXP; note that attribute values in ATTR_REGEXP
         need to be enclosed in double quotes and that an unescaped "."
         is replaced with '[^"]', so that the regular expression matches
         within a single annotation attribute value
         """,
         dict(required=True)),
        ('--replace|hide',
         """replace the content of the matching structures with fixed
         strings, instead of removing them completely
         """),
        ('--positional-attribute-replacement|posattr=STR "_" -> repl_posattr',
         """with --replace, replace positional attribute values of the matching
         structures with STR
         """),
        (('--structural-attribute-replacement|structattr=STR "removed" ->'
          ' repl_structattr'),
         """with --replace, replace structural attribute (annotation) values of
         the matching structures with STR
         """),
        ('--keep-structural-attributes=ATTRLIST -> keep_structattr',
         """with --replace, keep the values of structural attributes
         (annotations) listed in ATTRLIST intact instead of replacing
         them with a fixed string; ATTRLIST is a space-separated list
         of attribute specifications that are either of the form
         STRUCT_ATTRNAME or STRUCT:ATTRNAMELIST where STRUCT is the
         structure (such as "text"), ATTRNAME is a an annotation
         attribute of STRUCT and ATTRNAMELIST is a comma-separated
         list of them"""),
        ('--add-attribute=ATTR_VALUE -> add_attr',
         """with --replace, add annotation attribute ATTR_VALUE to
         structures whose content has been replaced with fixed
         strings; ATTR_VALUE is of the form ATTRNAME="VALUE"
         """),
        ('--verbose',
         """print information about the removed structures to stderr"""),
    ]

    def __init__(self):
        super().__init__()

    def main(self, args, inf, ouf):

        def make_attrcond_re(attr_regex):
            # Replace unescaped "." with [^"]

            def repl_fullstop(mo):
                # If the full stop is preceded by an even number of
                # backslashes, replace it
                start = mo.start()
                if (start - len(mo.string[:start].rstrip('\\'))) % 2 == 0:
                    return '[^"]'
                else:
                    return mo.group()

            attr_regex = re.sub('\.', repl_fullstop, attr_regex)
            # print(attr_regex, file=sys.stderr)
            return re.compile(' ' + attr_regex)

        def make_keep_structattr(attrlist):
            result = defaultdict(set)
            for attrspec in attrlist.split():
                if ':' in attrspec:
                    struct, _, attrnames = attrspec.partition(':')
                    result[struct] |= set(attrnames.split(','))
                elif '_' in attrspec:
                    struct, _, attrname = attrspec.partition('_')
                    result[struct].add(attrname)
                else:
                    self.error_exit(
                        'Attribute specification "' + attrspec
                        + '" is not of the form STRUCT:ATTR[,ATTR...] nor'
                        + ' STRUCT_ATTR')
            return result

        def replace_structattrs(line_s, struct_name):

            def attrsubst_repl_func(mo):
                name, value, quote, vbar = mo.group(
                    'name', 'value', 'quote', 'vbar')
                if name in keep_structattr[struct_name]:
                    repl = value
                else:
                    repl = vbar + repl_structattr + vbar
                return ' ' + name + '=' + quote + repl + quote

            if struct_name not in keep_structattr:
                return attrsubst_re.sub(attrsubst_repl, line_s)
            else:
                return attrsubst_re.sub(attrsubst_repl_func, line_s)

        LESS_THAN = '<'.encode()[0]
        SLASH = '/'.encode()[0]
        EXCLAM = '!'.encode()[0]
        VBAR = '|'.encode()[0]
        struct_name = args.struct_name
        struct_begin = ('<' + struct_name + ' ').encode()
        struct_end = ('</' + struct_name + '>').encode()
        remove = not args.replace
        verbose = args.verbose
        attrcond_re = make_attrcond_re(args.attr_regex)
        # This considers all attribute values beginning with a vertical bar as
        # feature-set-valued attributes, even though they should also have a
        # trailing vertical bar; however, this is easier to implement with a
        # replacement expression (as opposed to a function). Exactness is not
        # so important here.
        attrsubst_re = re.compile(
            r'''\s (?P<name> \w+) = (?P<quote> ["\'])
                                    (?P<value> (?P<vbar> \|?) (?:.*?) )
                                    (?P=quote)''',
            re.VERBOSE)
        repl_structattr = args.repl_structattr
        attrsubst_repl = (r' \g<name>=\g<quote>\g<vbar>'
                          + args.repl_structattr + r'\g<vbar>\g<quote>')
        repl_posattr = args.repl_posattr.encode()
        repl_posattr_tab = (args.repl_posattr + '\t').encode()
        keep_structattr = make_keep_structattr(args.keep_structattr or '')
        if args.add_attr:
            name, val = args.add_attr.split('=', 1)
            if val[0] in '"\'' and val[0] == val[-1] and len(val) > 1:
                val = val[1:-1]
            add_attr = ' ' + name + '="' + val + '">\n'
        else:
            add_attr = None
        remove_counts = defaultdict(int)
        matched = False
        linenr = 0
        for line in inf:
            linenr += 1
            if line[0] == LESS_THAN:
                line_s = line.decode()
                if matched:
                    if line.startswith(struct_end):
                        matched = False
                    elif line[1] not in [SLASH, EXCLAM]:
                        this_struct_name = line_s[1:].partition(' ')[0]
                        line_s = replace_structattrs(line_s, this_struct_name)
                        if verbose:
                            remove_counts[this_struct_name] += 1
                        line = line_s.encode()
                    if remove:
                        continue
                elif (line.startswith(struct_begin)
                      and attrcond_re.search(line_s)):
                    remove_counts[struct_name] += 1
                    if verbose:
                        print(('{verb} structure {struct} beginning on line'
                               ' {linenr}: {line}').format(
                                   verb=('Removing' if remove
                                         else 'Hiding the content of'),
                                   struct=struct_name, linenr=linenr,
                                   line=line_s[:-1]),
                               file=sys.stderr)
                    matched = True
                    if remove:
                        continue
                    line_s = replace_structattrs(line_s, struct_name)
                    if add_attr:
                        line_s = line_s.rstrip()[:-1] + add_attr
                    line = line_s.encode()
            elif matched and line != b'\n':
                remove_counts['TOKEN'] += 1
                if remove:
                    continue
                if b'\t|' in line:
                    # If the line contains an attribute value beginning with a
                    # vertical bar, replace such values with a value wrapped in
                    # vertical bars. For speed, only test the first character.
                    line = b'\t'.join(
                        (b'|' + repl_posattr + b'|' if posattr[0] == VBAR
                         else repl_posattr)
                        for posattr in line[:-1].split(b'\t')) + b'\n'
                else:
                    line = ((repl_posattr_tab * (line.count(b'\t') + 1))[:-1]
                            + b'\n')
            ouf.write(line)
        if verbose:
            if remove_counts:
                print('Removed:' if remove else 'Hid the content of:',
                      file=sys.stderr)
                for structname, count in remove_counts.items():
                    if structname != 'TOKEN':
                        print('  {count:8d} {struct} structures'
                              .format(count=count, struct=structname),
                              file=sys.stderr)
                print('  {count:8d} tokens'
                      .format(count=remove_counts['TOKEN']),
                      file=sys.stderr)
            else:
                print('No structures matching the condition were found',
                      file=sys.stderr)


if __name__ == '__main__':
    StructureRemover().run()
