from dataclasses import dataclass
from functools import cached_property
from typing import cast

import z3

from common.z3_utils.expr_utils import transform_quantifier


@dataclass
class Simplifier:
    enabled: bool = True

    @cached_property
    def false(self) -> z3.BoolRef:
        return z3.BoolVal(False)

    @cached_property
    def true(self) -> z3.BoolRef:
        return z3.BoolVal(True)

    def simplified_not(self, formula: z3.BoolRef) -> z3.BoolRef:
        if not self.enabled:
            return z3.Not(formula)

        if z3.is_true(formula):
            return self.false
        elif z3.is_false(formula):
            return self.true
        else:
            return z3.Not(formula)

    def simplified_and(self, *args: z3.BoolRef) -> z3.BoolRef:
        if not self.enabled:
            if len(args) == 0:
                return self.true
            elif len(args) == 1:
                return args[0]
            else:
                return z3.And(*args)

        non_true_args = [arg for arg in args if not z3.is_true(arg)]
        if len(non_true_args) == 0:
            return self.true
        elif len(non_true_args) == 1:
            return non_true_args[0]
        elif any(z3.is_false(child) for child in non_true_args):
            return self.false
        else:
            return z3.And(*non_true_args)

    def simplified_or(self, *args: z3.BoolRef) -> z3.BoolRef:
        if not self.enabled:
            if len(args) == 0:
                return self.false
            elif len(args) == 1:
                return args[0]
            else:
                return z3.Or(*args)

        non_false_args = [arg for arg in args if not z3.is_false(arg)]

        if len(non_false_args) == 0:
            return self.false
        elif len(non_false_args) == 1:
            return non_false_args[0]
        elif any(z3.is_true(child) for child in non_false_args):
            return self.true
        else:
            return z3.Or(*non_false_args)

    def simplified_implies(
        self, antecedent: z3.BoolRef, consequent: z3.BoolRef
    ) -> z3.BoolRef:
        if not self.enabled:
            return z3.Implies(antecedent, consequent)

        if z3.is_false(antecedent) or z3.is_true(consequent):
            return self.true
        elif z3.is_true(antecedent):
            return consequent
        elif z3.is_false(consequent):
            return z3.Not(antecedent)
        else:
            return z3.Implies(antecedent, consequent)

    def simplified_iff(self, left: z3.BoolRef, right: z3.BoolRef) -> z3.BoolRef:
        if not self.enabled:
            return left == right
        if _is_const_bool(left) and _is_const_bool(right):
            if z3.is_true(left) == z3.is_true(right):
                return self.true
            else:
                return self.false
        elif _is_const_bool(left):
            if z3.is_true(left):
                return right
            else:
                return z3.Not(right)
        elif _is_const_bool(right):
            if z3.is_true(right):
                return left
            else:
                return z3.Not(left)
        else:
            return left == right

    def simplified_xor(self, left: z3.BoolRef, right: z3.BoolRef) -> z3.BoolRef:
        if not self.enabled:
            return z3.Xor(left, right)

        if _is_const_bool(left) and _is_const_bool(right):
            if z3.is_true(left) != z3.is_true(right):
                return self.true
            else:
                return self.false
        elif _is_const_bool(left):
            if z3.is_true(left):
                return z3.Not(right)
            else:
                return right
        elif _is_const_bool(right):
            if z3.is_true(right):
                return z3.Not(left)
            else:
                return left
        else:
            return z3.Xor(left, right)


SIMPLIFIER = Simplifier()


def simplify(formula: z3.BoolRef) -> z3.BoolRef:
    if _is_const_bool(formula):
        return formula
    elif z3.is_quantifier(formula):
        simplified = transform_quantifier(formula, simplify)
        if _is_const_bool(simplified.body()):
            return simplified.body()
        else:
            return simplified
    elif formula.decl().kind() == z3.Z3_OP_UNINTERPRETED:
        return formula
    elif z3.is_eq(formula):
        return _simplify_equality(formula)
    elif z3.is_distinct(formula):
        return simplified_not(_simplify_equality(formula))
    elif z3.is_not(formula):
        return _simplify_not(formula)
    elif z3.is_and(formula):
        return _simplify_and(formula)
    elif z3.is_or(formula):
        return _simplify_or(formula)
    elif z3.is_implies(formula):
        return _simplify_implies(formula)
    elif z3.is_app(formula) and formula.decl().name() == "xor":
        return _simplify_xor(formula)
    else:
        return formula


def _is_const_bool(formula: z3.BoolRef) -> bool:
    return z3.is_true(formula) or z3.is_false(formula)


def _simplify_equality(formula: z3.BoolRef) -> z3.BoolRef:
    assert len(formula.children()) == 2, "Non-binary equality unsupported"

    left, right = formula.children()

    if z3.is_bool(left):  # iff construction
        assert z3.is_bool(right)

        simplified_left = simplify(left)
        simplified_right = simplify(right)
        return simplified_iff(simplified_left, simplified_right)
    else:
        if z3.is_int_value(left) and z3.is_int_value(right):
            return z3.BoolVal(left.as_long() == right.as_long())
        return left == right


def _simplify_not(formula: z3.BoolRef) -> z3.BoolRef:
    return simplified_not(simplify(cast(z3.BoolRef, formula.children()[0])))


def _simplify_and(formula: z3.BoolRef) -> z3.BoolRef:
    return simplified_and(
        *(
            simplified_child
            for child in formula.children()
            if not z3.is_true(simplified_child := simplify(cast(z3.BoolRef, child)))
        )
    )


def _simplify_or(formula: z3.BoolRef) -> z3.BoolRef:
    return simplified_or(
        *(
            simplified_child
            for child in formula.children()
            if not z3.is_false(simplified_child := simplify(cast(z3.BoolRef, child)))
        )
    )


def _simplify_implies(formula: z3.BoolRef) -> z3.BoolRef:
    return simplified_implies(
        *(simplify(cast(z3.BoolRef, child)) for child in formula.children())
    )


def _simplify_xor(formula: z3.BoolRef) -> z3.BoolRef:
    return simplified_xor(
        *(simplify(cast(z3.BoolRef, child)) for child in formula.children())
    )


simplified_not = SIMPLIFIER.simplified_not
simplified_and = SIMPLIFIER.simplified_and
simplified_or = SIMPLIFIER.simplified_or
simplified_implies = SIMPLIFIER.simplified_implies
simplified_iff = SIMPLIFIER.simplified_iff
simplified_xor = SIMPLIFIER.simplified_xor
