from functools import cache
from typing import Any, cast, overload

import cvc5.pythonic as cvc5
import z3

from common.z3_utils.expr_utils import Z3ExprDict, is_empty_and, is_empty_or
from common.z3_utils.rewriter import clone_vars

CVC5_REQUIRES_AT_LEAST_2 = {"And", "Or"}

_CONVERSION_CACHE = Z3ExprDict[z3.ExprRef, Any]()


@cache
def _convert_sort(sort: z3.SortRef) -> cvc5.SortRef:
    if sort.eq(z3.BoolSort()):
        return cvc5.BoolSort()
    elif sort.eq(z3.IntSort()):
        return cvc5.IntSort()
    else:
        return cvc5.DeclareSort(sort.name())


@cache
def _convert_const(const: z3.Const) -> cvc5.Const:
    return cvc5.Const(str(const.decl()), _convert_sort(const.sort()))


@cache
def _convert_fun_decl(decl: z3.FuncDeclRef) -> cvc5.FuncDeclRef:
    sorts = [_convert_sort(decl.domain(i)) for i in range(decl.arity())]
    sorts.append(_convert_sort(decl.range()))

    return cvc5.Function(str(decl), *sorts)


def _convert_quantifier(
    quantifier: z3.QuantifierRef,
    children: list[cvc5.ExprRef] | None = None,
) -> cvc5.QuantifierRef:
    if quantifier in _CONVERSION_CACHE:
        return cast(cvc5.QuantifierRef, _CONVERSION_CACHE[quantifier])

    old_vars = clone_vars(quantifier)
    new_vars = [_convert_const(var) for var in old_vars]

    if children is None:
        body = z3.substitute_vars(
            quantifier.body(), *reversed(old_vars)
        )  # Z3 uses vars in reverse order
        new_body = convert(body)
    else:
        assert len(children) == 1
        new_body = cast(cvc5.BoolRef, children[0])

    if quantifier.is_forall():
        result = cvc5.ForAll(new_vars, new_body)
    else:
        result = cvc5.Exists(new_vars, new_body)

    _CONVERSION_CACHE[quantifier] = result
    return result


def _convert_children(
    expression: z3.ExprRef, children: list[cvc5.ExprRef] | None = None
) -> list[cvc5.ExprRef]:
    if children is not None:
        return children
    if expression in _CONVERSION_CACHE:
        return cast(list[cvc5.ExprRef], _CONVERSION_CACHE[expression])
    result = [convert(child) for child in expression.children()]
    _CONVERSION_CACHE[expression] = result
    return result


@overload
def convert(
    expression: z3.BoolRef, children: list[cvc5.ExprRef] | None = None
) -> cvc5.BoolRef: ...


@overload
def convert(
    expression: z3.ExprRef, children: list[cvc5.ExprRef] | None = None
) -> cvc5.ExprRef: ...


def convert(
    expression: z3.ExprRef, children: list[cvc5.ExprRef] | None = None
) -> cvc5.ExprRef:
    if expression in _CONVERSION_CACHE:
        return cast(cvc5.ExprRef, _CONVERSION_CACHE[expression])
    result = do_convert(expression, children)
    _CONVERSION_CACHE[expression] = result
    return result


def do_convert(
    expression: z3.ExprRef, children: list[cvc5.ExprRef] | None = None
) -> cvc5.ExprRef:
    if z3.is_true(expression) or is_empty_and(expression):
        assert children is None
        return cvc5.BoolVal(True)
    elif z3.is_false(expression) or is_empty_or(expression):
        assert children is None
        return cvc5.BoolVal(False)
    elif z3.is_int_value(expression):
        assert children is None
        return cvc5.IntVal(expression.as_long())
    elif z3.is_lt(expression):
        left, right = cast(list[cvc5.ArithRef], _convert_children(expression, children))

        return left < right
    elif z3.is_le(expression):
        left, right = cast(list[cvc5.ArithRef], _convert_children(expression, children))

        return left <= right
    elif z3.is_ge(expression):
        left, right = cast(list[cvc5.ArithRef], _convert_children(expression, children))

        return left >= right
    elif z3.is_gt(expression):
        left, right = cast(list[cvc5.ArithRef], _convert_children(expression, children))

        return left > right
    elif z3.is_add(expression):
        left, right = cast(list[cvc5.ArithRef], _convert_children(expression, children))

        return left + right
    elif z3.is_mod(expression):
        left, right = cast(list[cvc5.ArithRef], _convert_children(expression, children))

        return left % right
    elif z3.is_sub(expression):
        left, right = cast(list[cvc5.ArithRef], _convert_children(expression, children))

        return left - right
    elif z3.is_app_of(expression, z3.Z3_OP_UMINUS):
        (sub,) = _convert_children(expression, children)
        return -cast(cvc5.ArithRef, sub)
    elif z3.is_const(expression):
        assert children is None
        return _convert_const(expression)
    elif z3.is_eq(expression):
        left_child, right_child = _convert_children(expression, children)
        return left_child == right_child
    elif z3.is_distinct(expression):
        left_child, right_child = _convert_children(expression, children)
        return left_child != right_child
    elif z3.is_quantifier(expression):
        return _convert_quantifier(expression, children)
    elif expression.decl().kind() == z3.Z3_OP_UNINTERPRETED:
        fun = _convert_fun_decl(expression.decl())
        children = _convert_children(expression, children)

        return fun(*children)
    else:
        form = str(expression.decl())
        children = _convert_children(expression, children)

        if len(children) == 1 and form in CVC5_REQUIRES_AT_LEAST_2:
            return children[0]

        return cast(cvc5.ExprRef, getattr(cvc5, form)(*children))


def cvc5_literal_arith(expression: cvc5.ArithRef) -> int | None:
    if expression.is_int():
        return cvc5.IntNumRef(expression.ast, expression.ctx).as_long()
    else:
        return None
