from collections.abc import Sequence

import z3


def set_pretty_print() -> None:
    z3.set_param("pp.max_depth", 100)
    z3.set_param("pp.max_indent", 60)
    z3.set_param("pp.flat_assoc", False)
    z3.set_param("pp.min_alias_size", 500_000)


def pretty_print(expr: z3.ExprRef) -> None:
    set_pretty_print()
    print(expr.sexpr())
    z3.reset_params()


def check_with_timeout(formula: z3.BoolRef, timeout_in_ms: int) -> z3.CheckSatResult:
    solver = z3.Solver()
    solver.set("timeout", timeout_in_ms)
    return solver.check(formula)


def validate(
    consequent: z3.BoolRef,
    *antecedents: z3.BoolRef,
) -> None:
    disprove(z3.Not(consequent), *antecedents)


def disprove(*assumptions: z3.BoolRef, log_instances: bool = False) -> None:
    solver = z3.Solver()

    def _try_decl(expr: z3.ExprRef) -> str:
        try:
            return expr.decl().name()
        except Exception as e:
            return f"?<{e}>"

    def log_instance(
        inference: z3.ExprRef, deps: list[z3.AstRef], clause: Sequence[z3.BoolRef]
    ) -> None:
        print(f"log clause: {inference.__class__}, {clause.__class__}")
        if len(deps) > 0:
            print(deps)
        if inference.decl().name() != "inst":
            return
        source = inference.arg(0)
        if not z3.is_quantifier(source):
            return
        qid = source.qid()
        if qid == "<null>":
            return

        for child in inference.children():
            if _try_decl(child) == "bind":
                print(f"instantiation for {qid}: {child.children()}")

    solver.add(*assumptions)

    if log_instances:
        z3.OnClause(solver, log_instance)

    assert solver.check() == z3.unsat
