package memory

import (
	"context"
	"github.com/abelith/gomax/state"
	"sync"
	"time"
)

type FSMStorage struct {
	m sync.Map
}

type EventIsolation struct {
	mu    sync.Mutex
	locks map[eventIsolationKey]*lockEntry
	ttl   time.Duration
}

type eventIsolationKey struct {
	userID int
	chatID int
}

type lockEntry struct {
	timer *time.Timer
	done  chan struct{}
}

func NewFSMStorage() *FSMStorage {
	return &FSMStorage{}
}

func (f *FSMStorage) Get(_ context.Context, userID int) (state.Snapshot, error) {
	st, ok := f.m.Load(userID)
	if !ok {
		return state.Snapshot{}, state.ErrStateNotFound
	}

	return st.(state.Snapshot), nil
}

func (f *FSMStorage) Put(_ context.Context, snapshot state.Snapshot) error {
	f.m.Store(snapshot.ID, snapshot)
	return nil
}

func (f *FSMStorage) Delete(_ context.Context, userID int) error {
	f.m.Delete(userID)
	return nil
}

func NewEventIsolation(ttl time.Duration) *EventIsolation {
	return &EventIsolation{
		locks: make(map[eventIsolationKey]*lockEntry),
		ttl:   ttl,
	}
}

func (ei *EventIsolation) Lock(ctx context.Context, userID, chatID int) error {
	key := eventIsolationKey{userID: userID, chatID: chatID}

	for {
		ei.mu.Lock()
		entry, busy := ei.locks[key]
		if !busy {
			ei.acquireLocked(key)
			ei.mu.Unlock()
			return nil
		}
		done := entry.done
		ei.mu.Unlock()

		select {
		case <-done:
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

func (ei *EventIsolation) Unlock(_ context.Context, userID, chatID int) error {
	key := eventIsolationKey{userID: userID, chatID: chatID}

	ei.mu.Lock()
	defer ei.mu.Unlock()

	if entry, ok := ei.locks[key]; ok {
		ei.releaseLocked(key, entry)
	}
	return nil
}

func (ei *EventIsolation) acquireLocked(key eventIsolationKey) {
	entry := &lockEntry{done: make(chan struct{})}
	entry.timer = time.AfterFunc(ei.ttl, func() {
		ei.mu.Lock()
		defer ei.mu.Unlock()
		if cur, ok := ei.locks[key]; ok && cur == entry {
			ei.releaseLocked(key, entry)
		}
	})
	ei.locks[key] = entry
}

func (ei *EventIsolation) releaseLocked(key eventIsolationKey, entry *lockEntry) {
	entry.timer.Stop()
	delete(ei.locks, key)
	close(entry.done)
}