import z3

from common.z3_utils.expr_utils import constants

from . import Assignment, SymStruct
from .model_checker import resolve_assignment


def find_assignment(formula: z3.BoolRef, struct: SymStruct) -> Assignment | None:
    variables = constants(formula)
    for assignment in struct.assignments(variables):  # todo: parallelize?
        result = resolve_assignment(formula, struct, assignment)
        if result == "unsat" or result == "unknown":
            continue

        return Assignment(
            {
                key: (
                    element.with_index(result.eval_expr(key))
                    if element.is_summary()
                    else element
                )
                for key, element in assignment.items()
            }
        )

    return None
