from abc import abstractmethod

from relationalai.early_access.dsl.core.exprs import contextStack
from relationalai.early_access.dsl.core.namespaces import Namespace
from relationalai.early_access.dsl.core.rules import Rule

def addrule(rul): return contextStack.root_context().addrule(rul)


# A @rule decorator that can be used to register a rule, optionally with one or more
# annotations.
#
# Examples:
#   @rule(Annotation.INNER_LOOP_NON_STRATIFIED, ...)
#   def my_rule():
#       ...
#
#   @rule()
#   def my_other_rule():
#       ...
def rule(*annotations):
    def decorator(func):
        rule_instance = addrule(func)
        for annotation in annotations:
            rule_instance.annotate(annotation)
        return func

    return decorator


class Relation:
    class NullContainer:
        def name(self): return ""

        def entity(self): return False

    # A Relation can serve as a context for declaring rules that
    # populate that relation and also for declaring constraints
    # that pertain only to this single relation -- e.g., functional
    # and mandatory constraints.
    #
    def __enter__(self):
        contextStack.push(self)
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        contextStack.pop()
        pass

    # A Relation can be "invoked" in the body of a rule definition to
    # declare what will be an Atom in the Schema generated by the rule
    # def.
    #
    # Note that because Schemas expect these Atoms to be defined on
    # RelationComponents, we have to make sure that such a component
    # exists in the Schema in context.
    #
    def __call__(self, *args):
        context = contextStack.root_context()
        if context is not None:
            return context.build_atom(self, args)
        return None

    def __init__(self, namespace, name, sig):
        self._relname = name
        self._container = namespace
        self._signature = sig
        self._rules = []

    # The "rule" to add is a function object with a number of parameters that should
    # equal the number of roles in this Relation.
    #
    def addrule(self, rule):
        if not isinstance(rule, Rule):
            # Assume that rule is a Python function that, when invoked can build a rule
            rule = Rule(self).elaborate(rule)

        self._rules.append(rule)
        return rule

    # Add a sequence of reading text snippets to use when verbalizing
    # this Relation.
    #
    def addverb(self, verb):
        self._verb = verb
        return self

    def arity(self):
        return self._signature.arity()

    def display(self):
        return self.qualified_name()

    def entityid(self):
        return hash(self.qualified_name())

    def function(self):
        self._signature.declare_functional()

    # Generates a new Rule that populates this relation, returning a tuple whose
    # first component is the Rule object and whose next n-1 components are the
    # free variables of this rule in order. Allows one to construct new rules that
    # populate a relation and get a handle on all free variables individually using
    # one method call, e.g.:
    #
    #   (rule, x, y, z) = relation.genrule()
    #
    def genrule(self):
        rule = Rule(self)
        self._rules.append(rule)
        tup = [rule]
        for v in rule.freevars():
            tup.append(v)
        return tuple(tup)

    def pprint(self):
        return self.display() + self._signature.display()

    # Returns the fully-qualified name that will be used in queries to
    # refer to this Relation.
    #
    def qualified_name(self):
        if not isinstance(self._container, Namespace):
            raise Exception(f"Cannot find containing namespace for Relation named {self._relname}")

        if self._container == Namespace.top:
            return self._relname
        else:
            # TODO: [VAMI] We changed `::` to `__` to spead up E2E testing. We should change it back to `::` before release.
            return f"{self._container.qualified_name()}__{self._relname}"

class AbstractRelationSignature:

    def __init__(self, *types):
        self._types = types

    def __call__(self, *args, **kwargs):
        seq = [self]
        for i in range(len(args)):
            seq.append(args[i])
        context = contextStack.root_context()
        if context is not None:
            return context.build_relation_variable(seq, kwargs)
        return None

    def arity(self):
        return len(self._types)

    def entityid(self):
        relhash = hash([r.entityid() for r in self._types])
        if self.functional():
            return hash((relhash, True))
        else:
            return hash((relhash, False))

    def relational(self) -> bool:
        return True

    # An n-ary AbstractRelationSignature is functional if its n-th column
    # is functionally dependent on the first n-1 columns.
    #
    @abstractmethod
    def functional(self) -> bool:
        """Return True when relation is functional otherwise False."""

    def domainsig(self):
        if self.functional():
            a = self.arity()
            return self._types[0:a - 1]
        else:
            return self._types

    def rangesig(self):
        if self.functional():
            return self._types[self.arity() - 1]
        else:
            return None

    def display(self):
        dsp = [r.display() for r in self.domainsig() if r is not None]
        dom = "(" + ", ".join(dsp) + ")"
        range_sig = self.rangesig()
        if range_sig is not None:
            return f"{dom} -> {range_sig.display()}"
        else:
            return dom

    # An AbstractRelationSignature is a subtype of another if the two have the
    # same number of types and in each corresponding type pair (t1, t2),
    # the type that plays t1 is a subtype of the type that plays t2.
    #
    def subtype_of(self, sig2):
        if self.arity() != sig2.arity():
            return False

        for i in range(self.arity()):
            if not self._types[i].subtype_of(sig2._types[i]):
                return False

        return True

class RelationSignature(AbstractRelationSignature):

    def __init__(self, *types):
        super().__init__(*types)
        self._is_functional = False

    def declare_functional(self):
        self._is_functional = True

    def functional(self): return self._is_functional
