import z3

from common.relations import relation
from problems.z3_utils import fresh_consts

ClientSort = z3.DeclareSort("ClientSort")
TimeSort = z3.DeclareSort("TimeSort")
type TimeType = z3.ExprRef

# sorts = [ClientSort, ServerSort, TimeSort]

C = z3.Const("C", ClientSort)
C_prime = z3.Const("C'", ClientSort)
C1 = z3.Const("C1", ClientSort)
C2 = z3.Const("C2", ClientSort)
# S = z3.Const("S", ServerSort)
# S1 = z3.Const("S1", ServerSort)
# S2 = z3.Const("S2", ServerSort)
T = z3.Const("T", TimeSort)
T1 = z3.Const("T1", TimeSort)
T2 = z3.Const("T2", TimeSort)
T3 = z3.Const("T3", TimeSort)


zero = z3.Const("zero", TimeSort)
lt = relation("lt", TimeSort, TimeSort)
prev = z3.Function("prev", TimeSort, TimeSort)
semaphore = relation("semaphore", TimeSort)
link = relation("link", TimeSort, ClientSort)


def leq(u: TimeType, v: TimeType) -> z3.BoolRef:
    return z3.Or(lt(u, v), u == v)


def succ_time(u: TimeType, v: TimeType) -> z3.BoolRef:
    return z3.And(lt(u, v), z3.ForAll(T, z3.Or(lt(v, T), lt(T, u), T == u, T == v)))


def no_init_time_specification() -> z3.BoolRef:
    return z3.ForAll(
        [T1, T2, T3],
        z3.And(
            z3.Implies(z3.And(lt(T1, T2), lt(T2, T3)), lt(T1, T3)),
            z3.Not(z3.And(lt(T1, T2), lt(T2, T1))),
            z3.Or(lt(T1, T2), lt(T2, T1), T1 == T2),
        ),
    )


def time_specification() -> z3.BoolRef:
    return z3.ForAll(
        [T1, T2, T3],
        z3.And(
            z3.Implies(z3.And(lt(T1, T2), lt(T2, T3)), lt(T1, T3)),
            z3.Not(z3.And(lt(T1, T2), lt(T2, T1))),
            z3.Or(lt(T1, T2), lt(T2, T1), T1 == T2),
            z3.ForAll(T1, z3.Or(lt(zero, T1), zero == T1)),
            prev(zero) == zero,
            z3.Implies(T1 != zero, succ_time(prev(T1), T1)),
            # succ_time(prev(T1), T1)
        ),
    )


def axiom(_t: TimeType) -> z3.BoolRef:
    return z3.BoolVal(True)


def globally_axiom() -> z3.BoolRef:
    return z3.ForAll(T, axiom(T))


def init(t: TimeType) -> z3.BoolRef:
    return z3.And(semaphore(t), z3.ForAll(C, z3.Not(link(t, C))))


def init_zero() -> z3.BoolRef:
    return init(zero)


def tr(t1: TimeType, t2: TimeType) -> z3.BoolRef:
    return z3.Or(
        # transition 'connect'
        z3.Exists(
            [C1],
            z3.And(
                semaphore(t1),
                z3.ForAll([C], link(t2, C) == z3.Or(link(t1, C), C == C1)),
                z3.Not(semaphore(t2)),
            ),
        ),
        # transition 'disconnect'
        z3.Exists(
            [C1],
            z3.And(
                link(t1, C1),
                # z3.Not(semaphore(t1))
                z3.ForAll([C], link(t2, C) == z3.And(link(t1, C), C != C1)),
                semaphore(t2),
            ),
        ),
    )


def no_init_globally_tr() -> z3.BoolRef:
    return z3.ForAll(T1, tr(prev(T1), T1))


def globally_tr() -> z3.BoolRef:
    return z3.ForAll(T1, z3.Implies(T1 != zero, tr(prev(T1), T1)))


def safety(t: TimeType) -> z3.BoolRef:
    return z3.ForAll([C1, C2], z3.Not(z3.And(C1 != C2, link(t, C1), link(t, C2))))


def globally_safe() -> z3.BoolRef:
    return z3.ForAll(T, safety(T))


def invariant(t: TimeType) -> z3.BoolRef:
    return z3.And(z3.ForAll([C], z3.Not(z3.And(link(t, C), semaphore(t)))), safety(t))


def invariant_without_safety(t: TimeType) -> z3.BoolRef:
    return z3.ForAll([C], z3.Not(z3.And(link(t, C), semaphore(t))))


def inductive_invariant() -> z3.BoolRef:
    return z3.Implies(
        z3.Exists(T1, z3.Not(invariant(T1))),
        z3.Exists(
            T2,
            z3.And(
                z3.Not(invariant(T2)),
                z3.ForAll(T, z3.Implies(z3.Not(invariant(T)), leq(T2, T))),
            ),
        ),
    )


def no_init_client_server() -> z3.BoolRef:
    return z3.And(
        no_init_time_specification(),
        globally_axiom(),
        # init_zero(),
        no_init_globally_tr(),
        z3.Not(globally_safe()),
        # z3.Exists([C, C_prime, T], z3.And(C != C_prime, link(T, C), link(T, C_prime)))
        # inductive_invariant(),
    )


def client_server() -> z3.BoolRef:
    return z3.And(
        time_specification(),
        globally_axiom(),
        init_zero(),
        globally_tr(),
        z3.Not(globally_safe()),
        # z3.Exists([C, C_prime, T], z3.And(C != C_prime, link(T, C), link(T, C_prime)))
        # inductive_invariant(),
    )


def liveness() -> z3.BoolRef:
    C, T, Tprime = fresh_consts(ClientSort, TimeSort, TimeSort)

    return z3.ForAll([C, T], z3.Exists(Tprime, z3.And(lt(T, Tprime), link(Tprime, C))))


def client_server_liveness() -> z3.BoolRef:
    return z3.And(
        time_specification(),
        globally_axiom(),
        init_zero(),
        globally_tr(),
        z3.Not(liveness()),
    )
