"""Constant folding of IR values.

For example, 3 + 5 can be constant folded into 8.
"""

from __future__ import annotations

from typing import Union
from typing_extensions import Final

from mypy.nodes import Expression, IntExpr, MemberExpr, NameExpr, OpExpr, StrExpr, UnaryExpr, Var
from mypyc.irbuild.builder import IRBuilder

# All possible result types of constant folding
ConstantValue = Union[int, str]
CONST_TYPES: Final = (int, str)


def constant_fold_expr(builder: IRBuilder, expr: Expression) -> ConstantValue | None:
    """Return the constant value of an expression for supported operations.

    Return None otherwise.
    """
    if isinstance(expr, IntExpr):
        return expr.value
    if isinstance(expr, StrExpr):
        return expr.value
    elif isinstance(expr, NameExpr):
        node = expr.node
        if isinstance(node, Var) and node.is_final:
            value = node.final_value
            if isinstance(value, (CONST_TYPES)):
                return value
    elif isinstance(expr, MemberExpr):
        final = builder.get_final_ref(expr)
        if final is not None:
            fn, final_var, native = final
            if final_var.is_final:
                value = final_var.final_value
                if isinstance(value, (CONST_TYPES)):
                    return value
    elif isinstance(expr, OpExpr):
        left = constant_fold_expr(builder, expr.left)
        right = constant_fold_expr(builder, expr.right)
        if isinstance(left, int) and isinstance(right, int):
            return constant_fold_binary_int_op(expr.op, left, right)
        elif isinstance(left, str) and isinstance(right, str):
            return constant_fold_binary_str_op(expr.op, left, right)
    elif isinstance(expr, UnaryExpr):
        value = constant_fold_expr(builder, expr.expr)
        if isinstance(value, int):
            return constant_fold_unary_int_op(expr.op, value)
    return None


def constant_fold_binary_int_op(op: str, left: int, right: int) -> int | None:
    if op == "+":
        return left + right
    if op == "-":
        return left - right
    elif op == "*":
        return left * right
    elif op == "//":
        if right != 0:
            return left // right
    elif op == "%":
        if right != 0:
            return left % right
    elif op == "&":
        return left & right
    elif op == "|":
        return left | right
    elif op == "^":
        return left ^ right
    elif op == "<<":
        if right >= 0:
            return left << right
    elif op == ">>":
        if right >= 0:
            return left >> right
    elif op == "**":
        if right >= 0:
            return left**right
    return None


def constant_fold_unary_int_op(op: str, value: int) -> int | None:
    if op == "-":
        return -value
    elif op == "~":
        return ~value
    elif op == "+":
        return value
    return None


def constant_fold_binary_str_op(op: str, left: str, right: str) -> str | None:
    if op == "+":
        return left + right
    return None