package fsm

import (
	"sync"
	"time"
)

type condAndState struct {
	cond  *sync.Cond
	state *UserState
}

type WaitCallback struct {
	Waiter sync.Map
	mu     sync.Mutex
}

func NewWaitCallback() *WaitCallback {
	return &WaitCallback{}
}

func (w *WaitCallback) Done(us *UserState) {
	if v, ok := w.Waiter.Load(us.UserId); ok {
		w.mu.Lock()
		defer w.mu.Unlock()
		cas := v.(*condAndState)
		cas.state = us
		cas.cond.Broadcast()
	}
}

func (w *WaitCallback) Wait(key int64, timeout time.Duration) *UserState {
	cas := &condAndState{
		cond:  sync.NewCond(&w.mu),
		state: nil,
	}

	v, loaded := w.Waiter.LoadOrStore(key, cas)
	if loaded {
		cas = v.(*condAndState)
	}

	done := make(chan struct{})
	go func() {
		w.mu.Lock()
		defer w.mu.Unlock()
		cas.cond.Wait()
		close(done)
	}()

	select {
	case <-done:
		return cas.state
	case <-time.After(timeout):
		return nil
	}
}