package main

import (
    "context"
    "github.com/goatx/goat"
)

type State struct {
    goat.State
    Name string
}

type MyStateMachine struct {
    goat.StateMachine
    Counter int
}

func main() {
    // Create state machine specification
    spec := goat.NewStateMachineSpec(&MyStateMachine{})
    stateA := &State{Name: "A"}
    stateB := &State{Name: "B"}

    // Define the states that the StateMachine can take and set initial state
    spec.
        DefineStates(stateA, stateB).
        SetInitialState(stateA)

    // Define behavior
    goat.OnEntry(spec, stateA, func(ctx context.Context, sm *MyStateMachine) {
        sm.Counter = 1
        goat.Goto(ctx, stateB)
    })

    goat.OnEntry(spec, stateB, func(ctx context.Context, sm *MyStateMachine) {
        sm.Counter = 2
    })

    // Create instance and run model checking
    sm, err := spec.NewInstance()
    if err != nil {
        log.Fatal(err)
    }
    cond := goat.NewCondition("counter<=2", sm, func(sm *MyStateMachine) bool {
        return sm.Counter <= 2
    })

    err = goat.Test(
        goat.WithStateMachines(sm),
        goat.WithRules(
            goat.Always(cond),
            goat.WheneverPEventuallyQ(cond, cond),
            goat.EventuallyAlways(cond),
            goat.AlwaysEventually(cond),
        ),
    )
    if err != nil {
        panic(err)
    }
}

cond := goat.NewCondition2("replication", primary, replica, func(p *Storage, r *Storage) bool {
    // Simple consistency: every key in primary exists in replica with the same value
    for key, pv := range p.Data {
        rv, ok := r.Data[key]
        if !ok || rv != pv {
            return false
        }
    }
    return true
})