from relationalai.semantics.frontend.base import Field, Library
from relationalai.semantics import metamodel as mm
from relationalai.semantics.metamodel.builtins import builtins as bt

library = Library("lqp")

Symbol = library.Type("Symbol")
mm_Symbol = bt.lqp.Symbol
assert isinstance(mm_Symbol, mm.TypeNode)

mm_Integer = bt.core.Integer
assert isinstance(mm_Integer, mm.TypeNode)

export = library.Relation("export", [])
external = library.Relation("external", [])
from_cdc = library.Relation("from_cdc", [])
track = library.Relation("track", [Field.input("library", Symbol), Field.input("relation", Symbol)])

#------------------------------------------------------
# Internal annotations required during compilation
#------------------------------------------------------

# @declare_constraint annotation marks Require nodes that should be emitted as LQP constraint declarations
_declare_constraint_relation = mm.Relation("declare_constraint", ())
declare_constraint = mm.Annotation(relation=_declare_constraint_relation, args=())

# @discharge annotation marks Require/Check nodes that should be removed during code generation
_discharged_relation = mm.Relation("discharged", ())
discharged = mm.Annotation(relation=_discharged_relation, args=())

#------------------------------------------------------
# Internal annotations required for configuring execution in the engine
#------------------------------------------------------

# @recursion_config(:require_seminaive) annotation for ensuring that recursion is safe. See
# RecursionConfig pass for more details.
_recursion_config_relation = mm.Relation("recursion_config", (mm.Field("recursion_type", mm_Symbol),))
require_seminaive = mm.Annotation(relation=_recursion_config_relation, args=(mm.Literal(mm_Symbol, "require_seminaive"),))

# @function(:checked) annotation indicates this relation has a functional dependency (k=1)
_function_relation = mm.Relation("function", (mm.Field("code", mm_Symbol, input=True),))
function_checked = mm.Annotation(relation=_function_relation, args=(mm.Literal(mm_Symbol, "checked"),))

# @function(:checked, k) annotation indicates this relation has a k-functional dependency
_function_ranked_relation = mm.Relation("function", (
    mm.Field("code", mm_Symbol, input=True),
    mm.Field("rank", mm_Integer, input=True),
))
def function_ranked_checked(k: int) -> mm.Annotation:
    """Create a @function(:checked, k) annotation for k-functional dependencies."""
    assert isinstance(mm_Symbol, mm.TypeNode)
    return mm.Annotation(relation=_function_ranked_relation, args=(mm.Literal(mm_Symbol, "checked"), mm.Literal(mm_Integer, k)))
