from collections.abc import Iterable

import z3

from common import config
from common.relations import Relation
from common.z3_utils.expr_utils import uninterpreted_relations


@config(order_relations=config.fest.order_relations)
def find_order_relation(
    source: z3.BoolRef | Iterable[Relation], *, order_relations: list[str] = []
) -> Relation | None:
    if isinstance(source, z3.BoolRef):
        source = uninterpreted_relations(source)

    names = frozenset(order_relations)
    for relation in source:
        if (
            relation.name().lower() in names
            and relation.arity() == 2
            and relation.domain(0) == relation.domain(1)
            and relation.domain(0).kind() == z3.Z3_UNINTERPRETED_SORT
        ):
            return relation
    return None


def well_definedness_formula(
    consts: Iterable[z3.Const], funs: Iterable[z3.FuncDeclRef]
) -> z3.BoolRef:
    return z3.And(well_defined_constants(consts), well_defined_functions(funs))


def well_defined_constants(consts: Iterable[z3.Const]) -> z3.BoolRef:
    return z3.And(
        *(
            z3.Exists(x, x == const)
            for const in consts
            if z3.is_const(x := z3.Const(f"x_{const}", const.sort()))
        )
    )


def well_defined_functions(funs: Iterable[z3.FuncDeclRef]) -> z3.BoolRef:
    return z3.And(
        *(
            z3.ForAll(args, z3.Exists(target, fun(*args) == target))
            for fun in funs
            if (
                args := [
                    z3.Const(f"x_{fun}_{i}", fun.domain(i)) for i in range(fun.arity())
                ]
            )
            if z3.is_const(target := z3.Const(f"x_{fun}", fun.range()))
        )
    )
