from relationalai.semantics.metamodel.compiler import Pass
from relationalai.semantics.metamodel import ir, builtins as rel_builtins, factory as f, visitor
from relationalai.semantics.metamodel import types

# Generate date arithmetic expressions, such as
# `rel_primitive_date_add(:day, [date] delta, res_2)` by finding the period
# expression for the delta and adding the period type to the date arithmetic expression.
#
# date_add and it's kin are generated by a period expression, e.g.,
# `day(delta, res_1)`
# followed by the date arithmetic expression using the period
# `date_add([date] res_1 res_2)`
class PeriodMath(Pass):
    def rewrite(self, model: ir.Model, options:dict={}) -> ir.Model:
        period_rewriter = self.PeriodRewriter()
        model = period_rewriter.walk(model)
        period_math_rewriter = self.PeriodMathRewriter(period_rewriter.period_vars)
        model = period_math_rewriter.walk(model)
        return model

    # Find all period builtins. We need to make them safe for the emitter (either by
    # translating to a cast, or removing) and store the variable and period type for use
    # in the date/datetime add/subtract expressions.
    class PeriodRewriter(visitor.Rewriter):
        def __init__(self):
            super().__init__()
            self.period_vars: dict[ir.Var, str] = {}

        def handle_lookup(self, node: ir.Lookup, parent: ir.Node) -> ir.Lookup:
            if not rel_builtins.is_builtin(node.relation):
                return node

            if node.relation.name not in {
                "year", "month", "week", "day", "hour", "minute", "second", "millisecond", "microsecond", "nanosecond"
            }:
                return node

            assert len(node.args) == 2, "Expect 2 arguments for period builtins"
            assert isinstance(node.args[1], ir.Var), "Expect result to be a variable"
            period = node.relation.name
            result_var = node.args[1]
            self.period_vars[result_var] = period

            # Ideally we could now remove the unused and unhandled period type construction
            # but we may also need to cast the original variable to an Int64 for use by the
            # date/datetime add/subtract expressions.
            # TODO: Remove the node entirely where possible and update uses of the result
            return f.lookup(rel_builtins.cast, [types.Int64, node.args[0], result_var])

    # Update date/datetime add/subtract expressions with period information.
    class PeriodMathRewriter(visitor.Rewriter):
        def __init__(self, period_vars: dict[ir.Var, str]):
            super().__init__()
            self.period_vars: dict[ir.Var, str] = period_vars

        def handle_lookup(self, node: ir.Lookup, parent: ir.Node) -> ir.Lookup:
            if not rel_builtins.is_builtin(node.relation):
                return node

            if node.relation.name in { "date_diff", "datetime_diff" }:
                # ensure the part argument is a symbol literal for the emitter
                part = node.args[0]
                assert isinstance(part, ir.Literal)
                return f.lookup(node.relation, [f.literal(part.value, types.Symbol), node.args[1], node.args[2], node.args[3]])

            if node.relation.name not in {
                "date_add", "date_subtract", "datetime_add", "datetime_subtract"
            }:
                return node

            if len(node.args) == 4:
                # We've already visited this lookup
                return node

            assert isinstance(node.args[1], ir.Var), "Expect period to be a variable"
            period_var = node.args[1]
            assert period_var in self.period_vars, "datemath found, but no vars to insert"

            period = self.period_vars[period_var]

            new_args = [f.literal(period, types.Symbol)] + [arg for arg in node.args]

            return f.lookup(node.relation, new_args)
