# Copyright (c) 2017 Renata Hodovan, Akos Kiss.
#
# Licensed under the BSD 3-Clause License
# <LICENSE.rst or https://opensource.org/licenses/BSD-3-Clause>.
# This file may not be copied, modified, or distributed except
# according to those terms.

import antlerinator
import autopep8
import logging
import re
import sys

from antlr4 import *
from argparse import ArgumentParser
from contextlib import contextmanager
from os.path import dirname, exists, join
from os import getcwd, makedirs

from .parser_builder import build_grammars
from .pkgdata import __version__, default_antlr_path
from .runtime.tree import *

logger = logging.getLogger('grammarinator')
logging.basicConfig(format='%(message)s')


class FuzzerGenerator(object):

    def __init__(self, lexer_root, parser_root, parser, actions):
        self.grammar_roots = [lexer_root, parser_root]
        self.parser = parser
        self.indent_level = 0
        self.actions = actions
        self.charset_idx = 0

        self.current_start_range = None
        self.token_start_ranges = dict()

        self.lexer_header = ''
        self.lexer_body = ''
        self.lexer_name = None
        self.parser_header = ''
        self.parser_body = ''
        self.parser_name = None

    @contextmanager
    def indent(self):
        self.indent_level += 4
        yield
        self.indent_level -= 4

    def line(self, src):
        return (' ' * self.indent_level) + src + '\n'

    def create_header(self, grammar_name, grammar_type):
        lexer = grammar_type == 'lexer'
        fuzzer_name = '{grammar_name}Un{grammar_type}'.format(
            grammar_name=grammar_name,
            grammar_type=grammar_type)

        src = self.line('# Generated by Grammarinator {version}\n'.format(version=__version__))
        src += self.line('from itertools import chain')
        src += self.line('from grammarinator.runtime import *\n')

        if lexer:
            self.lexer_header += src
        else:
            self.parser_header += src

        src = self.line('class {fuzzer_name}(Grammarinator):\n'.format(fuzzer_name=fuzzer_name))
        with self.indent():
            src += self.line('def __init__(self{args}):'.format(args='' if lexer else ', lexer'))

            with self.indent():
                src += self.line('super({fuzzer_name}, self).__init__()'.format(fuzzer_name=fuzzer_name))
                src += self.line('self.lexer = {lexer_ref}'.format(lexer_ref='self' if lexer else 'lexer'))
                src += self.line('self.set_options()\n')

        if lexer:
            self.lexer_body += src
        else:
            self.parser_body += src

        return fuzzer_name

    def init_fuzzer(self, grammar_name, lexer_fuzzer, parser_fuzzer):
        if lexer_fuzzer:
            self.lexer_name = self.create_header(grammar_name, 'lexer')
        if parser_fuzzer:
            self.parser_name = self.create_header(grammar_name, 'parser')

        if not (lexer_fuzzer or parser_fuzzer):
            self.lexer_name = self.create_header(grammar_name, 'lexer')
            self.parser_name = self.create_header(grammar_name, 'parser')

    def find_conditions(self, node):
        if not self.actions:
            return '1'

        if type(node) == str:
            return node

        action_block = getattr(node, 'actionBlock', None)
        if action_block:
            if action_block() and action_block().ACTION_CONTENT() and node.QUESTION():
                return ''.join([child.symbol.text for child in action_block().ACTION_CONTENT()])
            return '1'

        element = getattr(node, 'element', None) or getattr(node, 'lexerElement', None)
        if element:
            if not element():
                return '1'
            return self.find_conditions(element(0))

        child_ref = getattr(node, 'alternative', None) or getattr(node, 'lexerElements', None)
        return self.find_conditions(child_ref())

    def weight_array(self, nodes):
        return 'weights = [{content}]'.format(content=', '.join([self.find_conditions(node) for node in nodes]))

    def character_range_interval(self, node):
        start = node.characterRange().STRING_LITERAL(0).symbol.text[1:-1]
        end = node.characterRange().STRING_LITERAL(1).symbol.text[1:-1]

        return int(start.replace('\\u', '0x'), 16) if '\\u' in start else ord(start),\
               int(end.replace('\\u', '0x'), 16) if '\\u' in end else ord(end) + 1

    def lexer_charset_interval(self, src):
        elements = re.split('(\w-\w)', src)
        ranges = []
        for element in elements:
            if not element:
                continue

            # Convert character sequences like \n, \t, etc. into a single character.
            element = bytes(element, 'utf-8').decode('unicode_escape')
            if len(element) > 1:
                if element.strip('-').count('-') == 1:
                    start, end = element.split('-', 1)
                    ranges.append((ord(start), ord(end) + 1))
                else:
                    for char in element:
                        ranges.append((ord(char), ord(char) + 1))
            elif len(element) == 1:
                ranges.append((ord(element), ord(element) + 1))
        return ranges

    def chars_from_set(self, node):
        if node.characterRange():
            return [self.character_range_interval(node)]

        if node.LEXER_CHAR_SET():
            return self.lexer_charset_interval(node.LEXER_CHAR_SET().symbol.text[1:-1])

        if node.STRING_LITERAL():
            assert len(node.STRING_LITERAL().symbol.text) > 2, 'Negated string literal must not be empty.'
            first_char = ord(node.STRING_LITERAL().symbol.text[1])
            return [(first_char, first_char + 1)]

        if node.TOKEN_REF():
            src = node.TOKEN_REF().symbol.text
            assert src in self.token_start_ranges
            return self.token_start_ranges[src]

    def generate(self):
        for root in self.grammar_roots:
            self.generate_single(root)

    def generate_single(self, node):

        if isinstance(node, self.parser.GrammarSpecContext):
            self.init_fuzzer(node.identifier().TOKEN_REF().symbol.text.replace('Parser', '').replace('Lexer', ''),
                             node.grammarType().LEXER(), node.grammarType().PARSER())

            with self.indent():
                for child in node.children:
                    self.generate_single(child)
            return ''

        if isinstance(node, self.parser.OptionsSpecContext):
            options = []
            for option in node.option():
                ident = option.identifier()
                options.append('{name}="{value}"'.format(name=(ident.RULE_REF() or ident.TOKEN_REF()).symbol.text,
                                                         value=option.optionValue().getText()))

            set_options = self.line('def set_options(self):')
            with self.indent():
                set_options += self.line('self.options = dict({options})'.format(options=', '.join(options)))

            if self.lexer_body:
                self.lexer_body += set_options
            if self.parser_body:
                self.parser_body += set_options

        if isinstance(node, self.parser.ActionContext):
            if not self.actions:
                return ''

            scope_name = node.actionScopeName()
            if scope_name:
                action_scope = scope_name.LEXER() or scope_name.PARSER()
                assert action_scope, '{scope} scope not supported.'.format(scope=(scope_name.identifier().RULE_REF() or scope_name.identifier().TOKEN_REF()).symbol.text)
                action_scope = action_scope.symbol.text
            else:
                action_scope = 'parser'

            action_ident = node.identifier()
            action_type = (action_ident.RULE_REF() or action_ident.TOKEN_REF()).symbol.text
            raw_action_src = ''.join([child.symbol.text for child in node.actionBlock().ACTION_CONTENT()])

            if action_type == 'header':
                action_src = raw_action_src
            else:
                action_src = ''.join([self.line(line) for line in raw_action_src.splitlines()])

            # We simply append both member and header code chunks to the generated source.
            # It's the user's responsibility to define them in order.
            if action_scope == 'parser':
                # Both 'member' and 'members' keywords are accepted.
                if action_type.startswith('member'):
                    self.parser_body += action_src
                elif action_type == 'header':
                    self.parser_header += action_src
            elif action_scope == 'lexer':
                if action_type.startswith('member'):
                    self.lexer_body += action_src
            elif action_type == 'header':
                    self.lexer_header += action_src
            return ''

        if isinstance(node, (self.parser.ParserRuleSpecContext, self.parser.LexerRuleSpecContext)):
            parser_rule = isinstance(node, self.parser.ParserRuleSpecContext)
            node_type = UnparserRule if parser_rule else UnlexerRule
            rule_name = node.RULE_REF() if parser_rule else node.TOKEN_REF()

            # Mark that the next lexerAtom has to be saved as start range.
            if not parser_rule:
                self.current_start_range = []

            rule_header = self.line('def {rule_name}(self):'.format(rule_name=rule_name))
            with self.indent():
                local_ctx = self.line('local_ctx = dict()')
                rule_code = self.line('current = self.create_node({node_type}(name=\'{rule_name}\'))'.format(node_type=node_type.__name__,
                                                                                                             rule_name=rule_name))
                rule_block_name = 'ruleBlock' if parser_rule else 'lexerRuleBlock'
                rule_code += str(self.generate_single(getattr(node, rule_block_name)()))
                rule_code += self.line('return current\n')

            # local_ctx only has to be initialized if we have variable assignment.
            if 'local_ctx' in rule_code:
                rule_code = rule_header + local_ctx + rule_code
            else:
                rule_code = rule_header + rule_code

            if parser_rule:
                self.parser_body += rule_code
            else:
                self.lexer_body += rule_code

            if not parser_rule:
                self.current_start_range = None

            return rule_code

        if isinstance(node, (self.parser.RuleAltListContext, self.parser.AltListContext, self.parser.LexerAltListContext)):
            children = [child for child in node.children if isinstance(child, ParserRuleContext)]
            if len(children) == 1:
                return str(self.generate_single(children[0]))

            result = self.line(self.weight_array(children))
            result += self.line('choice = self.choice(weights)'.format(max=len(children)))
            for i, child in enumerate(children):
                result += self.line('{if_kw} choice == {idx}:'.format(if_kw='if' if i == 0 else 'elif', idx=i))
                with self.indent():
                    result += str(self.generate_single(child)) or self.line('pass')
            return result

        # Sequences.
        if isinstance(node, (self.parser.AlternativeContext, self.parser.LexerAltContext)):
            if not node.children:
                return self.line('current += UnlexerRule(src=\'\')')

            if isinstance(node, self.parser.AlternativeContext):
                children = node.element()
            elif isinstance(node, self.parser.LexerAltContext):
                children = node.lexerElements().lexerElement()
            else:
                children = []
            return ''.join([str(self.generate_single(child)) for child in children])

        if isinstance(node, (self.parser.ElementContext, self.parser.LexerElementContext)):
            if self.actions and node.actionBlock():
                # Conditions are handled at alternative processing.
                if node.QUESTION():
                    return ''

                action_src = ''.join([child.symbol.text for child in node.actionBlock().ACTION_CONTENT()])
                action_src = re.sub('\$(?P<var_name>\w+)', 'local_ctx[\'\g<var_name>\']', action_src)

                result = ''
                for line in action_src.splitlines():
                    result += self.line(line)
                return result

            suffix = None
            if node.ebnfSuffix():
                suffix = node.ebnfSuffix()
            elif hasattr(node, 'ebnf') and node.ebnf() and node.ebnf().blockSuffix():
                suffix = node.ebnf().blockSuffix().ebnfSuffix()

            if not suffix:
                return self.generate_single(node.children[0])

            suffix = suffix.children[0].symbol.text
            quant_type = {'?': 'zero_or_one', '*': 'zero_or_more', '+': 'one_or_more'}[suffix]
            result = self.line('for _ in self.{quant_type}():'.format(quant_type=quant_type))

            with self.indent():
                result += str(self.generate_single(node.children[0]))
            result += '\n'
            return result

        if isinstance(node, self.parser.LabeledElementContext):
            ident = node.identifier()
            name = ident.RULE_REF() or ident.TOKEN_REF()
            result = self.generate_single(node.atom() or node.block())
            result += self.line('local_ctx[\'{name}\'] = current.last_child'.format(name=name))
            return result

        if isinstance(node, self.parser.RulerefContext):
            return self.line('current += self.{rule_name}()'.format(rule_name=node.RULE_REF().symbol.text))

        if isinstance(node, (self.parser.LexerAtomContext, self.parser.AtomContext)):
            if node.characterRange():
                start, end = self.character_range_interval(node)
                if self.current_start_range is not None:
                    self.current_start_range.append(range(start, end))
                return self.line('current += self.create_node(UnlexerRule(src=self.char_from_list(range({start}, {end}))))'.format(start=start, end=end))

            if node.DOT():
                return self.line('current += UnlexerRule(src=self.any_char())')

            if node.notSet():
                if node.notSet().setElement():
                    options = self.chars_from_set(node.notSet().setElement())
                else:
                    options = []
                    for set_element in node.notSet().blockSet().setElement():
                        options.extend(self.chars_from_set(set_element))

                charset_var = 'charset_{idx}'.format(idx=self.charset_idx)
                self.charset_idx += 1
                self.lexer_header += '{charset_var} = list(chain(*multirange_diff(printable_unicode_ranges, [{charset}])))\n'.format(charset_var=charset_var, charset=','.join(['({start}, {end})'.format(start=chr_range[0], end=chr_range[1]) for chr_range in sorted(options, key=lambda x: x[0])]))
                return self.line('current += UnlexerRule(src=self.char_from_list({charset}))'.format(charset=charset_var))

            if isinstance(node, self.parser.LexerAtomContext) and node.LEXER_CHAR_SET():
                ranges = self.lexer_charset_interval(node.LEXER_CHAR_SET().symbol.text[1:-1])

                if self.current_start_range is not None:
                    self.current_start_range.extend(ranges)

                charset_var = 'charset_{idx}'.format(idx=self.charset_idx)
                self.charset_idx += 1
                self.lexer_header += '{charset_var} = list(chain({charset}))\n'.format(charset_var=charset_var, charset=', '.join(['range({start}, {end})'.format(start=chr_range[0], end=chr_range[1]) for chr_range in ranges]))
                return self.line('current += self.create_node(UnlexerRule(src=self.char_from_list({charset_var})))'.format(charset_var=charset_var))

            return ''.join([self.generate_single(child) for child in node.children])

        if isinstance(node, self.parser.TerminalContext):
            if node.TOKEN_REF():
                return self.line('current += self.lexer.{rule_name}()'.format(rule_name=node.TOKEN_REF().symbol.text))

            if node.STRING_LITERAL():
                src = node.STRING_LITERAL().symbol.text[1:-1]
                if self.current_start_range is not None:
                    self.current_start_range.append(src[0])
                return self.line('current += self.create_node(UnlexerRule(src=\'{src}\'))'.format(src=src))

        if isinstance(node, ParserRuleContext) and node.getChildCount():
            return ''.join([self.generate_single(child) for child in node.children])

        return ''


class FuzzerFactory(object):

    def __init__(self, work_dir, antlr):
        self.work_dir = work_dir

        antlr_dir = join(self.work_dir, 'antlr')
        makedirs(antlr_dir, exist_ok=True)
        # Add the path of the built grammars to the Python path to be available at parsing.
        sys.path.append(antlr_dir)

        self.lexer, self.parser, self.listener = build_grammars(antlr_dir, antlr=antlr)

    def generate_fuzzer(self, grammars, actions, out, pep8):
        lexer_root, parser_root, grammar_parser = None, None, None

        for grammar in grammars:
            root, grammar_parser = self.parse(grammar)
            # Lexer and/or combined grammars are processed first to evaluate TOKEN_REF-s.
            if root.grammarType().LEXER() or not root.grammarType().PARSER():
                lexer_root = root
            else:
                parser_root = root

        fuzzer_generator = FuzzerGenerator(lexer_root, parser_root, grammar_parser, actions)
        fuzzer_generator.generate()

        with open(join(out, fuzzer_generator.lexer_name + '.py'), 'w') as f:
            src = fuzzer_generator.lexer_header + '\n\n' + fuzzer_generator.lexer_body
            if pep8:
                src = autopep8.fix_code(src)
            f.write(src)

        with open(join(out, fuzzer_generator.parser_name + '.py'), 'w') as f:
            src = fuzzer_generator.parser_header + '\n\n' + fuzzer_generator.parser_body
            if pep8:
                src = autopep8.fix_code(src)
            f.write(src)

    def collect_imports(self, root, base_dir):
        imports = set()
        for prequel in root.prequelConstruct():
            if prequel.delegateGrammars():
                for delegate_grammar in prequel.delegateGrammars().delegateGrammar():
                    ident = delegate_grammar.identifier(0)
                    imports.add(join(base_dir, (ident.RULE_REF() or ident.TOKEN_REF()).symbol.text + '.g4'))
        return imports

    def parse_single(self, grammar):
        token_stream = CommonTokenStream(self.lexer(FileStream(grammar)))
        target_parser = self.parser(token_stream)
        root = target_parser.grammarSpec()
        # assert target_parser._syntaxErrors > 0, 'Parse error in ANTLR grammar.'
        return root, target_parser

    def parse(self, grammar):
        work_list = {grammar}
        root, target_parser = None, None

        while work_list:
            grammar = work_list.pop()
            current_root, current_parser = self.parse_single(grammar)

            # Save the 'outermost' grammar.
            if not root and not target_parser:
                root, target_parser = current_root, current_parser
            else:
                # Unite the rules of the imported grammar with the host grammar's rules.
                for rule in current_root.rules().ruleSpec():
                    root.rules().addChild(rule)

            work_list |= self.collect_imports(current_root, dirname(grammar))

        return root, target_parser


def execute():
    parser = ArgumentParser(description='Grammarinator: Processor')
    parser.add_argument('grammars', nargs='+', metavar='FILE',
                        help='ANTLR grammar files describing the expected format to generate.')
    parser.add_argument('--antlr', metavar='FILE', default=default_antlr_path,
                        help='path of the ANTLR jar file (default: %(default)s).')
    parser.add_argument('--no-actions', dest='actions', default=True, action='store_false',
                        help='do not process inline actions (default: %(default)s).')
    parser.add_argument('--pep8', default=False, action='store_true',
                        help='enable autopep8 to format the generated fuzzer.')
    parser.add_argument('-l', '--log-level', metavar='LEVEL', default=logging.INFO,
                        help='logging level')
    parser.add_argument('-o', '--out', metavar='DIR', default=getcwd(),
                        help='temporary working directory (default: .).')
    parser.add_argument('--version', action='version', version='%(prog)s {version}'.format(version=__version__))
    args = parser.parse_args()

    logger.setLevel(args.log_level)

    for grammar in args.grammars:
        if not exists(grammar):
            parser.error('{grammar} does not exist.'.format(grammar=grammar))

    if args.antlr == default_antlr_path:
        antlerinator.install(lazy=True)

    FuzzerFactory(args.out, args.antlr).generate_fuzzer(args.grammars, args.actions, args.out, args.pep8)


if __name__ == '__main__':
    execute()
