package api

import (
	"context"
	"strings"
	"sync"
	"time"

	"golang.org/x/sync/singleflight"

	"github.com/tidbcloud/AgentTeam/internal/store/workspace"
)

// DefaultWorkspaceResolverCacheTTL is the short process-local TTL used to
// collapse page-init and realtime reconnect bursts without hiding membership
// changes for long.
const DefaultWorkspaceResolverCacheTTL = 5 * time.Second

// CachedWorkspaceResolver shares successful workspace summary lookups across
// HTTP middleware and realtime handshakes.
type CachedWorkspaceResolver struct {
	next WorkspaceResolver
	ttl  time.Duration
	now  func() time.Time

	mu      sync.Mutex
	entries map[workspaceResolverCacheKey]workspaceResolverCacheEntry
	flights singleflight.Group
}

type workspaceResolverCacheKey struct {
	userID      string
	workspaceID string
}

func newWorkspaceResolverCacheKey(userID string, workspaceID string) workspaceResolverCacheKey {
	return workspaceResolverCacheKey{
		userID:      strings.TrimSpace(userID),
		workspaceID: strings.TrimSpace(workspaceID),
	}
}

func (key workspaceResolverCacheKey) valid() bool {
	return key.userID != "" && key.workspaceID != ""
}

func (key workspaceResolverCacheKey) flightKey() string {
	return key.userID + "\x00" + key.workspaceID
}

type workspaceResolverCacheEntry struct {
	summary   workspace.Summary
	expiresAt time.Time
}

// NewCachedWorkspaceResolver decorates a WorkspaceResolver with a short TTL
// cache and singleflight. Only successful summaries are cached.
func NewCachedWorkspaceResolver(next WorkspaceResolver, ttl time.Duration) *CachedWorkspaceResolver {
	if ttl <= 0 {
		ttl = DefaultWorkspaceResolverCacheTTL
	}
	return newCachedWorkspaceResolver(next, ttl, time.Now)
}

func newCachedWorkspaceResolver(next WorkspaceResolver, ttl time.Duration, now func() time.Time) *CachedWorkspaceResolver {
	if now == nil {
		now = time.Now
	}
	return &CachedWorkspaceResolver{
		next:    next,
		ttl:     ttl,
		now:     now,
		entries: make(map[workspaceResolverCacheKey]workspaceResolverCacheEntry),
	}
}

func (resolver *CachedWorkspaceResolver) GetWorkspaceForUser(ctx context.Context, userID string, workspaceID string) (workspace.Summary, error) {
	if resolver == nil || resolver.next == nil {
		return workspace.Summary{}, ErrForbidden
	}
	key := newWorkspaceResolverCacheKey(userID, workspaceID)
	if !key.valid() {
		return resolver.next.GetWorkspaceForUser(ctx, key.userID, key.workspaceID)
	}
	if summary, ok := resolver.lookup(key, resolver.now()); ok {
		return summary, nil
	}

	value, err, _ := resolver.flights.Do(key.flightKey(), func() (any, error) {
		if summary, ok := resolver.lookup(key, resolver.now()); ok {
			return summary, nil
		}
		summary, err := resolver.next.GetWorkspaceForUser(ctx, key.userID, key.workspaceID)
		if err != nil {
			return workspace.Summary{}, err
		}
		resolver.store(key, summary, resolver.now().Add(resolver.ttl))
		return summary, nil
	})
	if err != nil {
		return workspace.Summary{}, err
	}
	summary, ok := value.(workspace.Summary)
	if !ok {
		return workspace.Summary{}, ErrForbidden
	}
	return summary, nil
}

func (resolver *CachedWorkspaceResolver) lookup(key workspaceResolverCacheKey, now time.Time) (workspace.Summary, bool) {
	resolver.mu.Lock()
	defer resolver.mu.Unlock()

	entry, ok := resolver.entries[key]
	if !ok {
		return workspace.Summary{}, false
	}
	if !now.Before(entry.expiresAt) {
		delete(resolver.entries, key)
		return workspace.Summary{}, false
	}
	return entry.summary, true
}

func (resolver *CachedWorkspaceResolver) store(key workspaceResolverCacheKey, summary workspace.Summary, expiresAt time.Time) {
	resolver.mu.Lock()
	defer resolver.mu.Unlock()

	resolver.entries[key] = workspaceResolverCacheEntry{
		summary:   summary,
		expiresAt: expiresAt,
	}
}

func (resolver *CachedWorkspaceResolver) invalidateWhere(matches func(workspaceResolverCacheKey, workspaceResolverCacheEntry) bool) {
	if resolver == nil || matches == nil {
		return
	}
	var flightKeys []string
	resolver.mu.Lock()
	for key, entry := range resolver.entries {
		if !matches(key, entry) {
			continue
		}
		delete(resolver.entries, key)
		flightKeys = append(flightKeys, key.flightKey())
	}
	resolver.mu.Unlock()
	for _, flightKey := range flightKeys {
		resolver.flights.Forget(flightKey)
	}
}

// Invalidate removes one cached workspace summary for a user/workspace pair.
func (resolver *CachedWorkspaceResolver) Invalidate(userID string, workspaceID string) {
	if resolver == nil {
		return
	}
	key := newWorkspaceResolverCacheKey(userID, workspaceID)
	if !key.valid() {
		return
	}

	resolver.invalidateWhere(func(candidate workspaceResolverCacheKey, _ workspaceResolverCacheEntry) bool {
		return candidate == key
	})
}

// InvalidateWorkspace removes cached summaries for a workspace across users.
func (resolver *CachedWorkspaceResolver) InvalidateWorkspace(workspaceID string) {
	if resolver == nil {
		return
	}
	trimmedWorkspaceID := strings.TrimSpace(workspaceID)
	if trimmedWorkspaceID == "" {
		return
	}

	resolver.invalidateWhere(func(key workspaceResolverCacheKey, _ workspaceResolverCacheEntry) bool {
		return key.workspaceID == trimmedWorkspaceID
	})
}

// InvalidateUser removes cached summaries for one user across workspaces.
func (resolver *CachedWorkspaceResolver) InvalidateUser(userID string) {
	if resolver == nil {
		return
	}
	trimmedUserID := strings.TrimSpace(userID)
	if trimmedUserID == "" {
		return
	}

	resolver.invalidateWhere(func(key workspaceResolverCacheKey, _ workspaceResolverCacheEntry) bool {
		return key.userID == trimmedUserID
	})
}

// InvalidateOrganization removes cached summaries for workspaces in one organization.
func (resolver *CachedWorkspaceResolver) InvalidateOrganization(organizationID string) {
	if resolver == nil {
		return
	}
	trimmedOrganizationID := strings.TrimSpace(organizationID)
	if trimmedOrganizationID == "" {
		return
	}

	resolver.invalidateWhere(func(_ workspaceResolverCacheKey, entry workspaceResolverCacheEntry) bool {
		return entry.summary.OrganizationID == trimmedOrganizationID
	})
}

// Clear removes every cached workspace summary.
func (resolver *CachedWorkspaceResolver) Clear() {
	if resolver == nil {
		return
	}
	var flightKeys []string
	resolver.mu.Lock()
	for key := range resolver.entries {
		flightKeys = append(flightKeys, key.flightKey())
	}
	resolver.entries = make(map[workspaceResolverCacheKey]workspaceResolverCacheEntry)
	resolver.mu.Unlock()
	for _, flightKey := range flightKeys {
		resolver.flights.Forget(flightKey)
	}
}