from typing import Iterable, Mapping

import z3

from common.relations import Relation, relation
from common.z3_utils.expr_utils import domain, signature


def state_relation(name: str, *sig: z3.SortRef) -> tuple[Relation, Relation]:
    return relation(f"{name}_pre", *sig), relation(f"{name}_post", *sig)


def state_function(
    name: str, *sig: z3.SortRef
) -> tuple[z3.FuncDeclRef, z3.FuncDeclRef]:
    return z3.Function(f"{name}_pre", *sig), z3.Function(f"{name}_post", *sig)


def update_state(
    pre: z3.FuncDeclRef,
    post: z3.FuncDeclRef,
    conditions: Mapping[Iterable[z3.ExprRef], z3.ExprRef],
) -> z3.BoolRef:
    assert signature(pre) == signature(post), f"Bad pre/post pair `{pre}` / `{post}`"

    counter = 0

    def _next_const(sort: z3.SortRef) -> z3.Const:
        nonlocal counter
        counter = counter + 1
        return z3.Const(f"{str(sort)[0]}{counter}", sort)

    variables = [_next_const(sort) for sort in domain(pre)]

    result = post(*variables) == pre(*variables)

    for values, then in conditions.items():
        condition = z3.And(*(var == value for var, value in zip(variables, values)))
        result = z3.If(condition, post(*variables) == then, result)
    return z3.ForAll(variables, result, qid=f"update_{pre}")


def unchanged(pre: z3.FuncDeclRef, post: z3.FuncDeclRef) -> z3.BoolRef:
    return update_state(pre, post, {})
