from dataclasses import dataclass, field
from typing import cast

import z3

from common.z3_utils.expr_utils import transform_quantifier, unpack_quantifier


@dataclass
class Skolemized:
    formula: z3.BoolRef
    skolem_functions: dict[str, z3.FuncDeclRef] = field(default_factory=dict)


def skolemize(
    formula: z3.BoolRef, negated: bool = False, closure: tuple[z3.Const, ...] = ()
) -> Skolemized:
    if z3.is_quantifier(formula):
        if formula.is_forall() == negated:
            qid = formula.qid()
            assert qid != "<null>", f"Quantifier without id: {formula.sexpr()}"

            variables = [
                (formula.var_name(i), formula.var_sort(i))
                for i in range(formula.num_vars())
            ]
            parameters = [variable.sort() for variable in closure]

            skolem_functions = {
                f"{qid}:{variable}": z3.Function(f"{qid}:{variable}", *parameters, sort)
                for variable, sort in variables
            }

            applications = [
                skolem_functions[f"{qid}:{variable}"](*closure)
                for variable, _sort in variables
            ]

            substitutions = reversed(applications)
            skolemized_once = z3.substitute_vars(formula.body(), *substitutions)
            skolemized_fully = skolemize(skolemized_once, negated)
            skolemized_fully.skolem_functions.update(skolem_functions)
            return skolemized_fully
        else:
            skolem_functions = {}

            quantifier_variables, _body = unpack_quantifier(formula)

            new_closure = closure + tuple(quantifier_variables)

            def transform_body(body: z3.BoolRef) -> z3.BoolRef:
                skolemized_body = skolemize(body, negated, new_closure)
                skolem_functions.update(skolemized_body.skolem_functions)
                return skolemized_body.formula

            skolemized: z3.BoolRef = transform_quantifier(formula, transform_body)
            return Skolemized(skolemized, skolem_functions)
    elif (
        z3.is_app(formula)
        and not z3.is_const(formula)
        and z3.is_bool(formula.children()[0])
        and formula.decl().kind() != z3.Z3_OP_UNINTERPRETED
    ):

        if z3.is_not(formula):
            skolemized_children = [
                skolemize(cast(z3.BoolRef, formula.children()[0]), not negated, closure)
            ]
        elif z3.is_implies(formula):
            antecedent, consequent = formula.children()
            skolemized_children = [
                skolemize(cast(z3.BoolRef, antecedent), not negated, closure),
                skolemize(cast(z3.BoolRef, consequent), negated, closure),
            ]
        else:
            skolemized_children = [
                skolemize(cast(z3.BoolRef, child), negated, closure)
                for child in formula.children()
            ]

        skolemized_formulas = []
        skolem_functions = {}
        for child in skolemized_children:
            skolem_functions.update(child.skolem_functions)
            skolemized_formulas.append(child.formula)

        if z3.is_eq(formula):
            left, right = skolemized_formulas
            skolemized = left == right
        elif z3.is_distinct(formula):
            left, right = skolemized_formulas
            skolemized = left != right
        else:
            skolemized = cast(
                z3.BoolRef, getattr(z3, str(formula.decl()))(*skolemized_formulas)
            )

        return Skolemized(skolemized, skolem_functions)
    else:
        return Skolemized(formula)
