import z3

import problems
from problems.transition_systems import update_state
from problems.z3_utils import fresh_const, fresh_consts

from . import (
    Sort,
    SortType,
    dijkstra_self_stabilizing_system,
    first,
    is_privileged,
    is_process,
    key_post,
    key_pre,
    pred,
    succ,
)
from .sanity import sanity

acted_upon = z3.Const("acted_upon", Sort)


def first_acts() -> z3.BoolRef:
    return z3.And(
        acted_upon == first,
        is_privileged(first),
        update_state(key_pre, key_post, {(first,): succ(key_pre(first))}),
    )


def other_acts(s: SortType) -> z3.BoolRef:
    return z3.And(
        acted_upon == s,
        s != first,
        is_privileged(s),
        update_state(key_pre, key_post, {(s,): key_pre(pred(s))}),
    )


def transition() -> z3.BoolRef:
    s = fresh_const(Sort)
    return z3.Or(first_acts(), z3.Exists(s, other_acts(s)))


def safety(key: z3.FuncDeclRef = key_pre) -> z3.BoolRef:
    x, y = fresh_consts(Sort, Sort)
    return z3.Exists(
        x,
        z3.And(
            is_process(x),
            is_privileged(x, key),
            z3.ForAll(
                y, z3.Or(y == x, z3.Not(is_process(y)), z3.Not(is_privileged(y, key)))
            ),
        ),
    )


def safety_counterexample() -> z3.BoolRef:
    return z3.And(
        dijkstra_self_stabilizing_system(),
        safety(),
        sanity(),
        transition(),
        sanity(key_post),
        z3.Not(safety(key_post)),
    )


@problems.mark(
    model=True,
    slow="cvc5-conversion",
    bruteforce_axiom=True,
    guided_axiom=True,
    popl25=True,
)
def dijkstra_self_stabilizing_safety() -> z3.BoolRef:
    return safety_counterexample()
