# This file is part of A Boneless Datastructure Language
# Copyright (C) 2020  Soni L.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import collections.abc
import re
from abdl import predicates
from abdl import exceptions

class PatternElement:
    def on_not_in_key(self, frame, path, defs):
        raise NotImplementedError

    def on_in_key(self, frame, path, defs):
        raise NotImplementedError

    def collect_params(self, res: list):
        pass

class Arrow(PatternElement):
    def on_not_in_key(self, frame, path, defs):
        assert not path[-1].empty
        path.append(Holder(key=None, value=None, name=None, parent=path[-1].value, empty=True))
        return False

class StringKey(PatternElement):
    def __init__(self, toks):
        self.key = toks[0]
        self.skippable = toks[1] == '?'

    def on_in_key(self, frame, path, defs):
        return self.on_not_in_key(frame, path, defs)

    def on_not_in_key(self, frame, path, defs):
        path[-1].iterator = self.extract(path[-1].parent)
        path[-1].empty = False
        return True

    def extract(self, obj):
        try:
            yield (self.key, obj[self.key])
        except (TypeError, IndexError, KeyError):
            if not self.skippable:
                raise exceptions.ValidationError

class RegexKey(PatternElement):
    def __init__(self, toks):
        self.key = toks[0]
        self.compiled = re.compile(self.key)
        self.skippable = toks[1] == '?'

    def on_in_key(self, frame, path, defs):
        return self.on_not_in_key(frame, path, defs)

    def on_not_in_key(self, frame, path, defs):
        filtered_iterator = self.filter(path[-1].iterator)
        del path[-1].iterator
        path[-1].iterator = filtered_iterator
        del filtered_iterator
        path[-1].empty = False
        return True

    def filter(self, iter_):
        for el in iter_:
            try:
                if self.compiled.search(el[0]):
                    yield el
                elif not self.skippable:
                    raise exceptions.ValidationError
            except TypeError:
                if not self.skippable:
                    raise exceptions.ValidationError

class KeySubtree(PatternElement):
    def __init__(self, toks):
        self.key = toks[0]
        self.skippable = toks[1] == '?'

    def on_not_in_key(self, frame, path, defs):
        path[-1].subtree = True
        filtered_iterator = self.filter(path[-1].iterator, defs)
        del path[-1].iterator
        path[-1].iterator = filtered_iterator
        del filtered_iterator
        path[-1].empty = False
        return True

    def filter(self, iter_, defs):
        for x in iter_:
            for y in match_helper(self.key, defs, x[0]):
                yield (y, x[1])

    def collect_params(self, res: list):
        for sub in self.key:
            sub.collect_params(res)

class ValueSubtree(PatternElement):
    def __init__(self, toks):
        self.key = toks[0]
        self.skippable = toks[1] == '?'

    def on_not_in_key(self, frame, path, defs):
        assert not path[-1].empty
        path.append(Holder(key=None, value=None, name=None, parent=path[-1].value, empty=False, subtree=True))
        path[-1].iterator = self.filter(path[-1].parent, defs)
        return True

    def filter(self, parent, defs):
        for x in match_helper(self.key, defs, parent):
            yield (x, parent)

    def collect_params(self, res: list):
        for sub in self.key:
            sub.collect_params(res)

class Ident(PatternElement):
    def __init__(self, toks):
        self.key = toks[0]

    def on_not_in_key(self, frame, path, defs):
        path[-1].name = self.key
        path[-1].empty = False
        return True

class Param(PatternElement):
    def __init__(self, toks):
        assert isinstance(toks[1], Ident)
        self.skippable = toks[0] == '?'
        self.key = toks[1].key

    def on_in_key(self, frame, path, defs):
        return self.on_not_in_key(frame, path, defs)

    def on_not_in_key(self, frame, path, defs):
        path[-1].iterator = self.extract(path[-1].parent, defs[self.key])
        path[-1].empty = False
        return True

    def extract(self, obj, key):
        try:
            yield (key, obj[key])
        except (TypeError, IndexError, KeyError):
            if not self.skippable:
                raise exceptions.ValidationError

    def collect_params(self, res: list):
        res.append(self.key)

    def get_value(self, defs):
        return defs[self.key]

class ApplyPredicate(PatternElement):
    def __init__(self, toks):
        assert isinstance(toks[1], Ident)
        self.skippable = toks[0] == '?'
        self.key = toks[1].key

    def on_in_key(self, frame, path, defs):
        filtered_iterator = self.filter(path[-1].iterator, defs)
        del path[-1].iterator
        path[-1].iterator = filtered_iterator
        del filtered_iterator
        path[-1].empty = False
        return True

    def check(self, defs, obj):
        if predicates._to_predicate(defs[self.key]).accept(obj):
            return True
        if self.skippable:
            return False
        raise exceptions.ValidationError

    def on_not_in_key(self, frame, path, defs):
        assert len(path) == 1
        if not self.check(defs, path[-1].value):
            path.clear()
        return False

    def filter(self, iter_, defs):
        for el in iter_:
            if self.check(defs, el[1]):
                yield el

    def collect_params(self, res: list):
        res.append(self.key)

class End(PatternElement):
    def on_in_key(self, frame, path, defs):
        try:
            path[-1].next()
            return False
        except StopIteration:
            path.pop()
            while frame.prev() and not isinstance(frame.current_op, End):
                pass
            if not frame.prev():
                # FIXME?
                path.clear()
        return True # FIXME?

def _pairs(obj):
    if isinstance(obj, collections.abc.Mapping):
        return iter(obj.items())
    elif isinstance(obj, collections.abc.Sequence):
        return iter(enumerate(obj, 0))
    elif isinstance(obj, collections.abc.Set):
        return iter(((e, e) for e in obj))
    else:
        # maybe there's more stuff I can implement later
        raise TypeError

class Holder:
    def __init__(self, key, value, name, parent=None, iterator=None, empty=False, subtree=False):
        self.name = name
        self.key = key
        self.value = value
        self.empty = empty
        self._iterator = iterator
        self.parent = parent
        self.subtree = subtree

    @property
    def iterator(self):
        if self._iterator is None:
            self._iterator = _pairs(self.parent)
        return self._iterator

    @iterator.setter
    def iterator(self, value):
        assert self._iterator is None
        self._iterator = value

    @iterator.deleter
    def iterator(self):
        self._iterator = None

    def next(self):
        self.key, self.value = next(self.iterator)

class Frame:
    def __init__(self, ops):
        self.ops = ops
        self.pc = -1

    def next(self):
        pc = self.pc + 1
        if pc >= len(self.ops):
            return False
        self.pc = pc
        return True

    @property
    def current_op(self):
        return self.ops[self.pc]

    def prev(self):
        pc = self.pc - 1
        if pc < 0:
            return False
        self.pc = pc
        return True

def match_helper(ops, defs, tree):
    frame = Frame(ops)

    path = [Holder(key=None, value=tree, parent=None, iterator=iter(()), name=None)]
    in_key = False
    while path:
        if not frame.next():
            assert not path[-1].empty
            res = {}
            for h in path:
                if h.subtree:
                    for name, kv in h.key.items():
                        res[name] = kv
                elif h.name is not None:
                    res[h.name] = (h.key, h.value)
            yield res
            assert len(path) == 1 or isinstance(frame.current_op, End)
            if not frame.prev():
                return
            in_key = True
        else:
            if in_key:
                in_key = frame.current_op.on_in_key(frame, path, defs)
            else:
                in_key = frame.current_op.on_not_in_key(frame, path, defs)