package memory

import (
	"context"
	"github.com/abelith/gomax/v1/fsm/state"
	"github.com/abelith/gomax/v1/isolation"
	"github.com/google/uuid"
	"sync"
	"time"
)

type FSMStorage struct {
	m sync.Map
}

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

type fsmKey struct {
	userID int64
	chatID int64
}

type eventIsolationKey struct {
	userID int
	chatID int
}

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

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

func (f *FSMStorage) Get(_ context.Context, scope isolation.Scope) (state.Snapshot, error) {
	st, ok := f.m.Load(fsmKey{
		userID: scope.UserID,
		chatID: scope.ChatID,
	})
	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(fsmKey{
		userID: snapshot.Scope.UserID,
		chatID: snapshot.Scope.ChatID,
	}, snapshot)
	return nil
}

func (f *FSMStorage) Delete(_ context.Context, scope isolation.Scope) error {
	f.m.Delete(fsmKey{
		userID: scope.UserID,
		chatID: scope.ChatID,
	})
	return nil
}

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

func (ei *EventIsolation) Lock(ctx context.Context, key isolation.Key) error {
	scope := key.Scope()

	for {
		ei.mu.Lock()
		entry, busy := ei.locks[scope]
		if !busy {
			ei.acquireLocked(scope, key.UpdateID)
			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, key isolation.Key) error {
	scope := key.Scope()

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

	entry, ok := ei.locks[scope]
	if !ok || entry.owner != key.UpdateID {
		return isolation.ErrNotOwner
	}
	ei.releaseLocked(scope, entry)
	return nil
}

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

func (ei *EventIsolation) releaseLocked(scope isolation.Scope, entry *lockEntry) {
	entry.timer.Stop()
	delete(ei.locks, scope)
	close(entry.done)
}