package nixcache

import (
	"context"
	"log/slog"
	"maps"
	"slices"
	"sync"
	"time"
)

const (
	useFlushInterval = 10 * time.Second
	useMaxPending    = 4096
	useFlushTimeout  = 10 * time.Second
)

type UseRecordFunc func(ctx context.Context, cache string, storeHashes []string) error

type UseBatcher struct {
	cache  string
	record UseRecordFunc
	logger *slog.Logger
	max    int

	mu      sync.Mutex
	pending map[string]struct{}
	dropped int

	stop chan struct{}
	done chan struct{}
	once sync.Once
}

func NewUseBatcher(cache string, record UseRecordFunc, logger *slog.Logger) *UseBatcher {
	return newUseBatcher(cache, record, logger, useFlushInterval, useMaxPending)
}

func newUseBatcher(cache string, record UseRecordFunc, logger *slog.Logger, interval time.Duration, maxPending int) *UseBatcher {
	if logger == nil {
		logger = slog.Default()
	}
	b := &UseBatcher{
		cache:   cache,
		record:  record,
		logger:  logger.With("where", "nix_use", "cache", cache),
		max:     maxPending,
		pending: make(map[string]struct{}),
		stop:    make(chan struct{}),
		done:    make(chan struct{}),
	}
	go b.loop(interval)
	return b
}

// Add drops the hash past the pending cap, and only counts it for the log.
func (b *UseBatcher) Add(storeHash string) {
	if b == nil {
		return
	}
	b.mu.Lock()
	defer b.mu.Unlock()
	if _, ok := b.pending[storeHash]; ok {
		return
	}
	if len(b.pending) >= b.max {
		b.dropped++
		return
	}
	b.pending[storeHash] = struct{}{}
}

func (b *UseBatcher) Close() error {
	if b == nil {
		return nil
	}
	b.once.Do(func() { close(b.stop) })
	<-b.done
	return nil
}

func (b *UseBatcher) loop(interval time.Duration) {
	defer close(b.done)
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			b.flush()
		case <-b.stop:
			b.flush()
			return
		}
	}
}

func (b *UseBatcher) flush() {
	b.mu.Lock()
	hashes := slices.Collect(maps.Keys(b.pending))
	clear(b.pending)
	dropped := b.dropped
	b.dropped = 0
	b.mu.Unlock()

	if dropped > 0 {
		b.logger.Debug("dropped nix use reports past the pending cap", "dropped", dropped)
	}
	if len(hashes) == 0 {
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), useFlushTimeout)
	defer cancel()
	if err := b.record(ctx, b.cache, hashes); err != nil {
		b.logger.Debug("report nix use failed", "hashes", len(hashes), "error", err)
	}
}