import z3

import problems
from common.relations import Relation
from problems.transition_systems import (
    state_function,
    state_relation,
    unchanged,
    update_state,
)
from problems.z3_utils import TRUE, consts

from . import (
    Acceptor,
    AcceptorType,
    Ballot,
    BallotType,
    Quorum,
    QuorumType,
    Value,
    ValueType,
    _1,
    belongs_to,
    gt,
    gte,
    less_than,
    lte,
    order_axioms,
    pred,
    quorum_axioms,
)

none = z3.Const("none", Value)

# State
msg1a_pre, msg1a_post = state_relation("msg1a", Ballot)
"""
The global proposer says: let us consider this Ballot
"""

msg1b_pre, msg1b_post = state_relation("msg1b", Acceptor, Ballot, Ballot, Value)
"""
This acceptor promises never to vote on Ballot < B1, 
and says: I last voted at B2, and for Value
"""

msg2a_pre, msg2a_post = state_relation("msg2a", Ballot, Value)
"""
The global proposer says: let us vote for Value at Ballot
"""

msg2b_pre, msg2b_post = state_relation("msg2b", Acceptor, Ballot, Value)
"""
This acceptor votes for Value at Ballot
"""

max_bal_pre, max_bal_post = state_function("max_bal", Acceptor, Ballot)
"""
This acceptor promised to never vote on 'smaller' Ballots
"""

max_vbal_pre, max_vbal_post = state_function("max_vbal", Acceptor, Ballot)
"""
This acceptor last voted at Ballot
"""

max_val_pre, max_val_post = state_function("max_val", Acceptor, Value)
"""
This acceptor last voted for Value
"""

b, b_prime, v, v_prime, a = consts(
    ("b", Ballot),
    ("b_prime", Ballot),
    ("v", Value),
    ("v_prime", Value),
    ("a", Acceptor),
)

b1, v1, b2, v2 = consts(
    ("b1", Ballot),
    ("v1", Value),
    ("b2", Ballot),
    ("v2", Value),
)


def axiom1() -> z3.BoolRef:
    b, b_prime, v, b_star = consts(
        ("b", Ballot), ("b_prime", Ballot), ("v", Value), ("b_star", Ballot)
    )

    def predicate(ballot: BallotType, *, qid_prefix: str = "") -> z3.BoolRef:
        return z3.And(
            gt(ballot, b_prime), is_safe_at_paxos(ballot, v, qid_prefix=qid_prefix)
        )

    return z3.ForAll(
        [b_prime, v],
        z3.Implies(
            z3.Exists(b, predicate(b, qid_prefix="axiom1_q1_"), qid="axiom1_q1"),
            z3.Exists(
                b_star,
                z3.And(
                    predicate(b_star, qid_prefix="axiom1_b*_"),
                    z3.ForAll(
                        b,
                        z3.Implies(
                            gt(b_star, b), z3.Not(predicate(b, qid_prefix="axiom1_q2_"))
                        ),
                        qid="axiom1_q2",
                    ),
                ),
                qid="axiom1_q3",
            ),
        ),
        qid="axiom1_q4",
    )


def lemma2() -> z3.BoolRef:
    return z3.ForAll(
        [b, b_prime, v, v_prime],
        z3.Implies(
            lemma2_given(),
            lemma2_consequence(),
        ),
        qid="lemma2_q1",
    )


def lemma2_given() -> z3.BoolRef:
    return z3.And(
        gt(b, b_prime),
        gt(b_prime, _1),
        v != v_prime,
        is_safe_at_paxos(b, v, qid_prefix="lemma2_given_"),
        msg2a_pre(b_prime, v_prime),
    )


def lemma2_consequence() -> z3.BoolRef:
    b_plus, b_minus, q = consts(("b_plus", Ballot), ("b_minus", Ballot), ("q", Quorum))

    return z3.Exists(
        [b_plus, b_minus, q],
        lemma2_consequence_body(b_plus, b_minus, q),
        qid="lemma2_consequence",
    )


def lemma2_consequence_body(
    b_plus: BallotType, b_minus: BallotType, q: QuorumType
) -> z3.BoolRef:
    b_a, v_a = consts(("b_a", Ballot), ("v_a", Value))
    return z3.And(
        gt(b_plus, b_prime),
        gt(b_prime, b_minus),
        is_safe_at_paxos(b_plus, v, qid_prefix="lemma2_consequence_"),
        z3.ForAll(
            a,
            z3.Exists(
                [b_a, v_a],
                z3.Implies(
                    belongs_to(a, q),
                    z3.And(
                        lte(b_a, b_minus),
                        msg1b_pre(a, b_plus, b_a, v_a),
                    ),
                ),
                qid="lemma2_consequence_body_q1",
            ),
            qid="lemma2_consequence_body_q2",
        ),
    )


def chosen_at(
    b: BallotType, v: ValueType, msg2b: Relation, *, qid_prefix: str = ""
) -> z3.BoolRef:
    Q, A = consts(("Q", Quorum), ("A", Acceptor))
    return z3.Exists(
        Q,
        z3.ForAll(
            A,
            z3.Implies(belongs_to(A, Q), msg2b(A, b, v)),
            qid=f"{qid_prefix}chosen_at_q1",
        ),
        qid=f"{qid_prefix}chosen_at_q2",
    )


def chosen(v: ValueType, msg2b: Relation, *, qid_prefix: str = "") -> z3.BoolRef:
    (B,) = consts(("B", Ballot))

    return z3.Exists(
        B, chosen_at(B, v, msg2b, qid_prefix=qid_prefix), qid=f"{qid_prefix}chosen"
    )


def shows_safe_at_paxos_skolemized(
    q: QuorumType, b: BallotType, v: ValueType, msg1b: Relation, *, qid_prefix: str = ""
) -> z3.BoolRef:
    A, Mb, Mv, Mb2, Mv2 = consts(
        ("A", Acceptor), ("Mb", Ballot), ("Mv", Value), ("Mb2", Ballot), ("Mv2", Value)
    )
    return z3.And(
        z3.ForAll(
            A,
            z3.Implies(
                belongs_to(A, q),
                z3.Exists(
                    [Mb, Mv], msg1b(A, b, Mb, Mv), qid=f"{qid_prefix}shows_safe_at1"
                ),
            ),
            qid=f"{qid_prefix}shows_safe_at2",
        ),
        z3.Or(
            z3.ForAll(
                [A, Mb, Mv],
                z3.Not(z3.And(belongs_to(A, q), msg1b(A, b, Mb, Mv), Mb != _1)),
                qid=f"{qid_prefix}shows_safe_at3",
            ),
            z3.And(
                pred(b) != _1,
                less_than(pred(b), b),  # todo: could this be proved instead of assumed?
                z3.Exists(
                    A,
                    z3.And(belongs_to(A, q), msg1b(A, b, pred(b), v)),
                    qid=f"{qid_prefix}shows_safe_at4",
                ),
                z3.ForAll(
                    [A, Mb2, Mv2],
                    z3.Implies(
                        z3.And(belongs_to(A, q), msg1b(A, b, Mb2, Mv2), Mb2 != _1),
                        lte(Mb2, pred(b)),
                    ),
                    qid=f"{qid_prefix}shows_safe_at5",
                ),
            ),
        ),
    )


def is_safe_at_paxos(
    b: BallotType, v: ValueType, msg1b: Relation = msg1b_pre, *, qid_prefix: str = ""
) -> z3.BoolRef:
    q = z3.Const("q", Quorum)
    return z3.Exists(
        q,
        shows_safe_at_paxos_skolemized(q, b, v, msg1b, qid_prefix=qid_prefix),
        qid=f"{qid_prefix}is_safe_at_paxos",
    )


def phase1a(b: BallotType) -> z3.BoolRef:
    return z3.And(
        b != _1,
        update_state(msg1a_pre, msg1a_post, {(b,): TRUE}),
        unchanged(msg1b_pre, msg1b_post),
        unchanged(msg2a_pre, msg2a_post),
        unchanged(msg2b_pre, msg2b_post),
        unchanged(max_bal_pre, max_bal_post),
        unchanged(max_vbal_pre, max_vbal_post),
        unchanged(max_val_pre, max_val_post),
    )


def phase1b(a: AcceptorType, b: BallotType) -> z3.BoolRef:
    return z3.And(
        b != _1,
        msg1a_pre(b),
        gt(b, max_bal_pre(a)),
        update_state(max_bal_pre, max_bal_post, {(a,): b}),
        update_state(
            msg1b_pre,
            msg1b_post,
            {(a, b, max_vbal_pre(a), max_val_pre(a)): TRUE},
        ),
        unchanged(msg1a_pre, msg1a_post),
        unchanged(msg2a_pre, msg2a_post),
        unchanged(msg2b_pre, msg2b_post),
        unchanged(max_vbal_pre, max_vbal_post),
        unchanged(max_val_pre, max_val_post),
    )


def phase2a(b: BallotType, v: ValueType) -> z3.BoolRef:
    (V,) = consts(("V", Value))
    return z3.And(
        b != _1,
        v != none,
        z3.Not(z3.Exists(V, msg2a_pre(b, V))),
        is_safe_at_paxos(b, v, msg1b_pre),
        update_state(msg2a_pre, msg2a_post, {(b, v): TRUE}),
        unchanged(msg1a_pre, msg1a_post),
        unchanged(msg1b_pre, msg1b_post),
        unchanged(msg2b_pre, msg2b_post),
        unchanged(max_bal_pre, max_bal_post),
        unchanged(max_vbal_pre, max_vbal_post),
        unchanged(max_val_pre, max_val_post),
    )


def phase2b(a: AcceptorType, b: BallotType, v: ValueType) -> z3.BoolRef:
    return z3.And(
        b != _1,
        v != none,
        msg2a_pre(b, v),
        gte(b, max_bal_pre(a)),
        update_state(max_bal_pre, max_bal_post, {(a,): b}),
        update_state(max_vbal_pre, max_vbal_post, {(a,): b}),
        update_state(max_val_pre, max_val_post, {(a,): v}),
        update_state(msg2b_pre, msg2b_post, {(a, b, v): TRUE}),
        unchanged(msg1a_pre, msg1a_post),
        unchanged(msg1b_pre, msg1b_post),
        unchanged(msg2a_pre, msg2a_post),
    )


def transition() -> z3.BoolRef:
    A, B, V = consts(("A", Acceptor), ("B", Ballot), ("V", Value))
    return z3.Exists([A, B], phase2b(A, B, v2), qid="transition")


def simplified_transition() -> z3.BoolRef:
    A, B, V = consts(("A", Acceptor), ("B", Ballot), ("V", Value))
    return z3.Exists(
        [A, B, V],
        phase2b(A, B, V),
    )


def safety(msg2b: Relation, *, qid_prefix: str = "") -> z3.BoolRef:
    V1, V2 = consts(("V1", Value), ("V2", Value))

    return z3.ForAll(
        [V1, V2],
        z3.Implies(
            z3.And(
                chosen(V1, msg2b, qid_prefix=f"{qid_prefix}safety_v1_"),
                chosen(V2, msg2b, qid_prefix=f"{qid_prefix}v2_"),
            ),
            V1 == V2,
        ),
        qid=f"{qid_prefix}safety",
    )


# Invariants
def inv1(msg1b: Relation, msg2b: Relation) -> z3.BoolRef:
    A, B, V = consts(("A", Acceptor), ("B", Ballot), ("V", Value))
    return z3.ForAll(
        [A, B, V], z3.Implies(msg2b(A, B, V), is_safe_at_paxos(B, V, msg1b))
    )


def inv2(msg2b: Relation) -> z3.BoolRef:
    A, B, V1, V2 = consts(("A", Acceptor), ("B", Ballot), ("V1", Value), ("V2", Value))
    return z3.ForAll(
        [A, B, V1, V2],
        z3.Implies(z3.And(chosen_at(B, V1, msg2b), msg2b(A, B, V2)), V1 == V2),
    )


def inv3(msg2a: Relation = msg2a_pre, msg1b: Relation = msg1b_pre) -> z3.BoolRef:
    B, V = consts(("B", Ballot), ("V", Value))
    return z3.ForAll(
        [B, V], z3.Implies(msg2a(B, V), is_safe_at_paxos(B, V, msg1b)), qid="inv3"
    )


def inv4(msg2a: Relation = msg2a_pre) -> z3.BoolRef:
    B, V1, V2 = consts(("B", Ballot), ("V1", Value), ("V2", Value))
    return z3.ForAll(
        [B, V1, V2],
        z3.Implies(z3.And(msg2a(B, V1), msg2a(B, V2)), V1 == V2),
        qid="inv4",
    )


def inv5(msg2a: Relation = msg2a_pre, msg2b: Relation = msg2b_pre) -> z3.BoolRef:
    A, B, V = consts(("A", Acceptor), ("B", Ballot), ("V", Value))
    return z3.ForAll([A, B, V], z3.Implies(msg2b(A, B, V), msg2a(B, V)), qid="inv5")


def inv6(msg1b: Relation, max_bal: z3.FuncDeclRef) -> z3.BoolRef:
    A, B, Bmax, Vmax = consts(
        ("A", Acceptor), ("B", Ballot), ("Bmax", Ballot), ("Vmax", Value)
    )

    return z3.ForAll(
        [A, B, Bmax, Vmax], z3.Implies(msg1b(A, B, Bmax, Vmax), gte(max_bal(A), B))
    )


def inv7(msg1b: Relation = msg1b_pre, msg2b: Relation = msg2b_pre) -> z3.BoolRef:
    A, B, Bmax, Vmax = consts(
        ("A", Acceptor), ("B", Ballot), ("Bmax", Ballot), ("Vmax", Value)
    )
    return z3.ForAll(
        [A, B, Bmax, Vmax],
        z3.Implies(
            z3.And(
                B != _1,
                Bmax != _1,
                msg1b(A, B, Bmax, Vmax),
            ),
            msg2b(A, Bmax, Vmax),
        ),
        qid="inv7",
    )


def inv8(msg1b: Relation, msg2b: Relation) -> z3.BoolRef:
    A, B, Bmid, Bmax, V, Vmax = consts(
        ("A", Acceptor),
        ("B", Ballot),
        ("Bmid", Ballot),
        ("Bmax", Ballot),
        ("V", Value),
        ("Vmax", Value),
    )
    return z3.ForAll(
        [A, B, Bmid, Bmax, V, Vmax],
        z3.Implies(
            z3.And(
                gt(B, Bmid),
                gt(Bmid, Bmax),
                msg1b(A, B, Bmax, Vmax),
            ),
            z3.Not(msg2b(A, Bmid, V)),
        ),
    )


def inv9(max_vbal: z3.FuncDeclRef, max_bal: z3.FuncDeclRef) -> z3.BoolRef:
    (A,) = consts(("A", Acceptor))
    return z3.ForAll(A, lte(max_vbal(A), max_bal(A)))


def inv10(max_vbal: z3.FuncDeclRef, msg2b: Relation) -> z3.BoolRef:
    A, B, V = consts(("A", Acceptor), ("B", Ballot), ("V", Value))
    return z3.ForAll([A, B, V], z3.Implies(msg2b(A, B, V), gte(max_vbal(A), B)))


def inv11(
    max_vbal: z3.FuncDeclRef, max_val: z3.FuncDeclRef, msg2b: Relation
) -> z3.BoolRef:
    (A,) = consts(("A", Acceptor))
    return z3.ForAll(
        A, z3.Implies(max_vbal(A) != _1, msg2b(A, max_vbal(A), max_val(A)))
    )


def inv0(
    msg2b: Relation = msg2b_pre, msg1b: Relation = msg1b_pre, *, qid_prefix: str = ""
) -> z3.BoolRef:
    A, B, V = consts(("A", Acceptor), ("B", Ballot), ("V", Value))
    Mb, Mv = consts(("Mb", Ballot), ("Mv", Value))
    return z3.And(
        z3.ForAll(
            [A, B, V], z3.Implies(msg2b(A, B, V), B != _1), qid=f"{qid_prefix}inv0_1"
        ),
        z3.ForAll(
            [A, B, Mb, Mv],
            z3.Implies(z3.And(msg1b(A, B, Mb, Mv), Mb == _1), Mv == none),
            qid=f"{qid_prefix}inv0_2",
        ),
    )


def invariants() -> z3.BoolRef:
    return z3.And(
        inv1(msg1b_pre, msg2b_pre),
        inv2(msg2b_pre),
        inv3(msg2a_pre, msg1b_pre),
        inv4(msg2a_pre),
        inv5(msg2a_pre, msg2b_pre),
        inv6(msg1b_pre, max_bal_pre),
        inv7(msg1b_pre, msg2b_pre),
        inv8(msg1b_pre, msg2b_pre),
        inv9(max_vbal_pre, max_bal_pre),
        inv10(max_vbal_pre, msg2b_pre),
        inv11(max_vbal_pre, max_val_pre, msg2b_pre),
        inv0(msg2b_pre),
    )


def post_invariants() -> z3.BoolRef:
    return z3.And(
        inv1(msg1b_post, msg2b_post),
        inv2(msg2b_post),
        inv3(msg2a_post, msg1b_post),
        inv4(msg2a_post),
        inv5(msg2a_post, msg2b_post),
        inv6(msg1b_post, max_bal_post),
        inv7(msg1b_post, msg2b_post),
        inv8(msg1b_post, msg2b_post),
        inv9(max_vbal_post, max_bal_post),
        inv10(max_vbal_post, msg2b_post),
        inv11(max_vbal_post, max_val_post, msg2b_post),
        inv0(msg2b_post),
    )


def two_chosen() -> z3.BoolRef:
    return z3.And(
        chosen(v1, msg2b_post, qid_prefix="v1_"),
        chosen(v2, msg2b_post, qid_prefix="v2_"),
    )


def cex_given() -> z3.BoolRef:
    return z3.And(
        quorum_axioms(),
        inv0(msg2b_post, msg1b_post, qid_prefix="post_"),
        transition(),
        safety(msg2b_pre),
        two_chosen(),
    )


def counter_example() -> z3.BoolRef:
    return z3.And(
        quorum_axioms(),
        order_axioms(),
        safety(msg2b_pre),
        invariants(),
        transition(),
        two_chosen(),
        v1 != v2,
    )


@problems.mark(
    model=True,
    slow=True,
    popl24=True,
)
def skolemized_paxos() -> z3.BoolRef:
    return counter_example()
