package privatelink

import (
	"context"
	"sync"
	"time"

	"golang.org/x/sync/singleflight"

	regionalapiv1 "github.com/tidbcloud/cluster-service-ng/api/tidbcloud/dataplane-manager/v1"
)

type endpointServiceZonesCacheEntry struct {
	zones     []*regionalapiv1.EndpointServiceAvailabilityZone
	expiresAt time.Time
}

type endpointServiceZonesCache struct {
	ttl time.Duration
	now func() time.Time

	mu      sync.RWMutex
	entries map[string]endpointServiceZonesCacheEntry
	group   singleflight.Group
}

func newEndpointServiceZonesCache(ttl time.Duration, now func() time.Time) *endpointServiceZonesCache {
	if now == nil {
		now = time.Now
	}
	return &endpointServiceZonesCache{
		ttl:     ttl,
		now:     now,
		entries: map[string]endpointServiceZonesCacheEntry{},
	}
}

func (c *endpointServiceZonesCache) resolve(
	ctx context.Context,
	key string,
	query func() ([]*regionalapiv1.EndpointServiceAvailabilityZone, error),
) ([]*regionalapiv1.EndpointServiceAvailabilityZone, error) {
	if zones, ok := c.get(key); ok {
		return zones, nil
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}

	resultCh := c.group.DoChan(key, func() (any, error) {
		if zones, ok := c.get(key); ok {
			return zones, nil
		}
		zones, err := query()
		if err != nil {
			return nil, err
		}
		// Only successful lookups are cached so transient cloud or service state changes recover immediately.
		c.set(key, zones)
		return copyEndpointServiceAvailabilityZones(zones), nil
	})

	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case result := <-resultCh:
		if result.Err != nil {
			return nil, result.Err
		}
		return copyEndpointServiceAvailabilityZones(result.Val.([]*regionalapiv1.EndpointServiceAvailabilityZone)), nil
	}
}

func (c *endpointServiceZonesCache) get(key string) ([]*regionalapiv1.EndpointServiceAvailabilityZone, bool) {
	now := c.now()
	c.mu.RLock()
	entry, ok := c.entries[key]
	c.mu.RUnlock()
	if !ok {
		return nil, false
	}
	if !now.Before(entry.expiresAt) {
		c.mu.Lock()
		if current, ok := c.entries[key]; ok && !now.Before(current.expiresAt) {
			delete(c.entries, key)
		}
		c.mu.Unlock()
		return nil, false
	}
	return copyEndpointServiceAvailabilityZones(entry.zones), true
}

func (c *endpointServiceZonesCache) set(key string, zones []*regionalapiv1.EndpointServiceAvailabilityZone) {
	now := c.now()
	c.mu.Lock()
	defer c.mu.Unlock()
	// Prune on writes so one-off cache keys expire even when they are never read again.
	for cacheKey, entry := range c.entries {
		if !now.Before(entry.expiresAt) {
			delete(c.entries, cacheKey)
		}
	}
	c.entries[key] = endpointServiceZonesCacheEntry{
		zones:     copyEndpointServiceAvailabilityZones(zones),
		expiresAt: now.Add(c.ttl),
	}
}