package main

import (
	"context"
	"fmt"
	"log"
	"time"

	"github.com/patrickmn/go-cache"
)

// Define the function signature for the underlying function we want to cache.
// The key is removed, as per the user request.  Added context.Context.
type DataFetcher func(ctx context.Context, app, principal string) (map[string]struct{}, error)

// cacheKey generates a unique cache key based on the app and principal.
// The key is NOT included here.
func cacheKey(app, principal string) string {
	return fmt.Sprintf("%s:%s", app, principal)
}

// NewPullThroughCache creates a new pull-through cache using go-cache.
func NewPullThroughCache(fetcher DataFetcher) *PullThroughCache {
	// Create a new go-cache with a default expiration time of 5 minutes and
	// a purge interval of 10 minutes.  The purge interval doesn't drastically
	// affect the sliding expiration, but it's good practice to set it.
	c := cache.New(5*time.Minute, 10*time.Minute)
	return &PullThroughCache{
		cache:   c,
		fetcher: fetcher,
	}
}

// PullThroughCache is a struct that wraps the go-cache and the data fetching function.
type PullThroughCache struct {
	cache   *cache.Cache
	fetcher DataFetcher
}

// GetOrLoad retrieves data from the cache, or loads it using the provided
// fetcher function if it's not in the cache.
func (p *PullThroughCache) GetOrLoad(ctx context.Context, app, principal, key string) (map[string]struct{}, error) {
	cacheKey := cacheKey(app, principal) // key no longer part of cache key.

	// Check if the value is in the cache.
	cachedValue, found := p.cache.Get(cacheKey)
	if found {
		// Type assert the cached value to the correct type.  This is crucial.
		if value, ok := cachedValue.(map[string]struct{}); ok {
			return value, nil // Return the cached value.
		} else {
			// Should not happen, unless there's a bug in the cache or an external modification.
			log.Printf("error: unexpected type in cache for key %s: got %T, expected %T", cacheKey, cachedValue, map[string]struct{}{})
			p.cache.Delete(cacheKey) //remove the invalid entry
		}
	}

	// If the value is not in the cache, fetch it using the provided fetcher function.
	// The key is no longer passed to the fetcher.  Pass the context.
	value, err := p.fetcher(ctx, app, principal)
	if err != nil {
		return nil, err // Return the error from the fetcher function.
	}

	// Only cache the result if the key is present in the returned map.
	_, keyPresent := value[key]
	if keyPresent {
		// Set the value in the cache with a 5-minute expiration time.
		p.cache.Set(cacheKey, value, cache.DefaultExpiration)
	} else {
		// Invalidate the entire cache for this app and principal
		p.invalidateCache(app, principal)
	}

	return value, nil // Return the fetched value.
}

// invalidateCache invalidates all cache entries for a given app and principal.
// This is necessary when the underlying data changes.  The key is not part of the cache key.
func (p *PullThroughCache) invalidateCache(app, principal string) {
	p.cache.Delete(cacheKey(app, principal)) // Delete exact key
}

func main() {
	// Simulate a data fetching function.
	// The key parameter is removed from the function signature.
	// Added context parameter.
	fetchData := func(ctx context.Context, app, principal string) (map[string]struct{}, error) {
		fmt.Printf("Fetching data for app: %s, principal: %s from source with context: %v...\n", app, principal, ctx)
		// Simulate a database or API call with a delay.
		time.Sleep(500 * time.Millisecond) // Simulate latency

		// Simulate different results based on the key.
		//  The key logic remains here, to simulate different data sets
		//  being returned.
		if app == "myApp" && principal == "user123" {
			return map[string]struct{}{"key1": {}, "key2": {}}, nil
		} else if app == "myApp" && principal == "user456" {
			return map[string]struct{}{"keyA": {}, "keyB": {}}, nil
		}
		return map[string]struct{}{}, nil
	}

	// Create a new pull-through cache.
	cache := NewPullThroughCache(fetchData)

	// --- Example Usage ---
	app := "myApp"
	principal := "user123"
	ctx := context.Background() // Use a background context.

	// First call for key1: data will be fetched and cached.
	result1, err := cache.GetOrLoad(ctx, app, principal, "key1")
	if err != nil {
		log.Fatalf("Error getting data for key1: %v", err)
	}
	fmt.Printf("Result 1: %v\n", result1)

	// Second call for key1: data will be retrieved from the cache.
	result2, err := cache.GetOrLoad(ctx, app, principal, "key2")
	if err != nil {
		log.Fatalf("Error getting data for key2: %v", err)
	}
	fmt.Printf("Result 2: %v (from cache)\n", result2)

	// Call for key3: data will be fetched, and because key3 is not in the
	// result, the cache will be invalidated.
	result3, err := cache.GetOrLoad(ctx, app, principal, "key3")
	if err != nil {
		log.Fatalf("Error getting data for key3: %v", err)
	}
	fmt.Printf("Result 3: %v\n", result3)

	// Call for key1 again: data will be fetched from source because the cache
	// was invalidated in the previous step.
	result4, err := cache.GetOrLoad(ctx, app, principal, "key1")
	if err != nil {
		log.Fatalf("Error getting data for key1: %v", err)
	}
	fmt.Printf("Result 4: %v (after invalidation)\n", result4)

	// Wait for longer than the TTL.
	time.Sleep(6 * time.Minute)

	// Call for key1 again: Data will be fetched again after TTL expiry
	result5, err := cache.GetOrLoad(ctx, app, principal, "key1")
	if err != nil {
		log.Fatalf("Error getting data for key1: %v", err)
	}
	fmt.Printf("Result 5: %v (after TTL expiry)\n", result5)

	// Example with a different principal to show cache isolation.
	principal2 := "user456"
	result6, err := cache.GetOrLoad(ctx, app, principal2, "keyA") // Should not be affected by previous invalidation
	if err != nil {
		log.Fatalf("Error getting data for keyA: %v", err)
	}
	fmt.Printf("Result 6 (different principal): %v\n", result6)
}