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 to_lia_variable(variable: z3.Const) -> z3.Int:
    if z3.is_int(variable):
        return z3.Int(str(variable))
    else:
        return z3.Int(f"i_{variable}")
