package pipeline

import (
	"context"
	"fmt"
	"sync"
	"sync/atomic"
)

type memoryIdempotencyEntry struct {
	state IdempotencyState
	token string
}

// MemoryIdempotencyStore only protects concurrent deliveries in one process.
// Complete releases the lease so a later version with the same stable checkId can update.
type MemoryIdempotencyStore struct {
	mu      sync.Mutex
	entries map[string]memoryIdempotencyEntry
	next    atomic.Uint64
}

func NewMemoryIdempotencyStore() *MemoryIdempotencyStore {
	return &MemoryIdempotencyStore{entries: make(map[string]memoryIdempotencyEntry)}
}

func (s *MemoryIdempotencyStore) Acquire(ctx context.Context, key string) (IdempotencyLease, error) {
	if err := ctx.Err(); err != nil {
		return IdempotencyLease{}, err
	}
	if key == "" {
		return IdempotencyLease{}, fmt.Errorf("idempotency key is required")
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if entry, ok := s.entries[key]; ok {
		if entry.state == IdempotencyAcquired {
			return IdempotencyLease{State: IdempotencyInProgress}, nil
		}
		return IdempotencyLease{State: entry.state}, nil
	}
	token := fmt.Sprintf("lease-%d", s.next.Add(1))
	s.entries[key] = memoryIdempotencyEntry{state: IdempotencyAcquired, token: token}
	return IdempotencyLease{State: IdempotencyAcquired, Token: token}, nil
}

func (s *MemoryIdempotencyStore) Complete(ctx context.Context, key, token string) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	entry, ok := s.entries[key]
	if !ok || entry.state != IdempotencyAcquired || entry.token != token {
		return fmt.Errorf("idempotency lease is missing or stale")
	}
	delete(s.entries, key)
	return nil
}

func (s *MemoryIdempotencyStore) Release(ctx context.Context, key, token string) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	entry, ok := s.entries[key]
	if !ok || entry.state != IdempotencyAcquired || entry.token != token {
		return fmt.Errorf("idempotency lease is missing or stale")
	}
	delete(s.entries, key)
	return nil
}