from relationalai.semantics.metamodel import Relation, Annotation, Update, Node, Literal, Sequence, Loop
from relationalai.semantics.metamodel.builtins import builtins as bt
from relationalai.semantics.metamodel.rewriter import Rewriter

__all__ = [
    "LoopyInstructionAdder", "LoopyGlobalAdder",
    "init_relation", "init_annotation", "mk_init",
    "global_relation", "global_annotation", "mk_global",
    "assign_relation", "assign_annotation", "mk_assign",
    "upsert_relation", "upsert_annotation", "mk_upsert",
    "monoid_relation", "monoid_annotation", "mk_monoid",
    "monus_relation", "monus_annotation", "mk_monus",
    "algorithm_relation", "algorithm_annotation", "mk_algorithm",
    "script_relation", "script_annotation", "mk_script",
    "while_relation", "while_annotation", "mk_while",
]

# Loopy annotations

# LoopyIR annotation rewriters
class LoopyInstructionAdder(Rewriter):
    def __init__(self, anno: Annotation):
        self.anno = anno
        super().__init__()

    def apply_update(self, node: Node):
        return self(node)

    def update(self, node: Update):
        new_annos = tuple(node.annotations) + (self.anno,)
        return node.mut(annotations=new_annos)

class LoopyGlobalAdder(Rewriter):
    def __init__(self, globals: list[Relation]):
        self.globals = globals
        super().__init__()

    def apply_update(self, node: Node):
        return self(node)

    def update(self, node: Update):
        if node.relation in self.globals:
            new_annos = tuple(node.annotations) + (global_annotation(),)
            return node.mut(annotations=new_annos)
        else:
            return node

# Init

init_relation = Relation("init")
def init_annotation():
    return Annotation(relation=init_relation)

def mk_init(i: Node):
    return LoopyInstructionAdder(init_annotation()).apply_update(i)

# Empty

empty_relation = Relation("empty")
def empty_annotation():
    return Annotation(relation=empty_relation)

def mk_empty(i: Node):
    return LoopyInstructionAdder(empty_annotation()).apply_update(i)

# Global

global_relation = Relation("global")
def global_annotation():
    return Annotation(relation=global_relation)

def mk_global(globals: list[Relation], i: Node):
    return LoopyGlobalAdder(globals).apply_update(i)

# Assign

assign_relation = Relation("assign")
def assign_annotation():
    return Annotation(assign_relation)

def mk_assign(i: Node):
    return LoopyInstructionAdder(assign_annotation()).apply_update(i)

# Upsert

upsert_relation = Relation("upsert")
def upsert_annotation(arity: int):
    arity_val = Literal(bt.core.Integer, arity)
    return Annotation(relation=upsert_relation, args=(arity_val,))

def mk_upsert(i: Node, arity: int):
    return LoopyInstructionAdder(upsert_annotation(arity)).apply_update(i)

# Monoid

monoid_relation = Relation("monoid")
def monoid_annotation(arity: int, monoid_type, monoid_op):
    arity_val = Literal(bt.core.Integer, arity)
    return Annotation(relation=monoid_relation, args=(arity_val, monoid_type, monoid_op))

def mk_monoid(i: Node, arity: int, monoid_type, monoid_op):
    return LoopyInstructionAdder(monoid_annotation(arity, monoid_type, monoid_op)).apply_update(i)

# Monus

monus_relation = Relation("monus")
def monus_annotation(arity: int, monoid_type, monoid_op):
    arity_val = Literal(bt.core.Integer, arity)
    return Annotation(relation=monus_relation, args=(arity_val, monoid_type, monoid_op))

def mk_monus(i: Node, arity: int, monoid_type, monoid_op):
    return LoopyInstructionAdder(monus_annotation(arity, monoid_type, monoid_op)).apply_update(i)

# Algorithm

algorithm_relation = Relation("algorithm", ())
def algorithm_annotation():
    return Annotation(relation=algorithm_relation)

def mk_algorithm(s: Sequence):
    new_annos = tuple(s.annotations) + (algorithm_annotation(),)
    return s.mut(annotations=new_annos)

# Script

script_relation = Relation("script", ())
def script_annotation():
    return Annotation(relation=script_relation)

def mk_script(s: Sequence):
    new_annos = tuple(s.annotations) + (script_annotation(),)
    return s.mut(annotations=new_annos)

# While

while_relation = Relation("while", ())
def while_annotation():
    return Annotation(relation=while_relation)

def mk_while(s: Sequence | Loop):
    new_annos = tuple(s.annotations) + (while_annotation(),)
    return s.mut(annotations=new_annos)

# Loopy debug annotations

# recursion_content
debug_relation = Relation("recursion_logging", ())

def debug_annotation(type: str, iter: int):
    type_lit = Literal(bt.core.String, type)
    iter_lit = Literal(bt.core.Integer, iter)
    return Annotation(relation=debug_relation, args=(type_lit, iter_lit))

def recursion_content_debug_annotation(iter: int):
    return debug_annotation("relation_content", iter)
