package datakeycache

import (
	"errors"
	"math"
	"math/rand"
	"sync"
	"time"

	"code.justin.tv/amzn/TwitchOLE/ole/internal/keybytes"
	"code.justin.tv/amzn/TwitchOLE/ole/internal/stats"
	"github.com/hashicorp/golang-lru/simplelru"
)

// KeyFetcher is called to fetch data keys from kms
type KeyFetcher func() (plaintextDataKey, encryptedDataKey []byte, err error)

// CacheItem stores data keys and maybe tickets.
type CacheItem struct {
	conf configurator

	// All items below are guarded by `mu`
	mu               sync.Mutex
	expiration       int64
	tickets          int64
	plaintextDataKey []byte
	encryptedDataKey []byte
}

func (ci *CacheItem) withLock(f func()) {
	ci.mu.Lock()
	defer ci.mu.Unlock()
	f()
}

// clear the existing keys out of memory as best possible. This must be done holding a lock.
func (ci *CacheItem) clearKeys() {
	keybytes.Zero(ci.plaintextDataKey)
	keybytes.Zero(ci.encryptedDataKey)
	ci.plaintextDataKey, ci.encryptedDataKey = nil, nil
}

// GetOrFetchKeys fetches the keys inside this item.
// If the keys exist and are not expired, they are immediately returned.
// Otherwise, if the keys do not exist or the item is expired, then the passed 'fetcher' is called.
// Presuming 'fetcher' returns without error, the keys are then cached inside this item and also returned.
func (ci *CacheItem) GetOrFetchKeys(fetcher KeyFetcher) (plaintextDataKey, encryptedDataKey []byte, err error) {
	ci.mu.Lock()
	defer ci.mu.Unlock()
	if ci.plaintextDataKey != nil && (ci.expiration == 0 || time.Now().UnixNano() < ci.expiration) {
		if ci.tickets > 0 {
			ci.tickets--
			return keybytes.Copy(ci.plaintextDataKey), keybytes.Copy(ci.encryptedDataKey), nil
		}
	}

	ci.clearKeys()
	plaintextDataKey, encryptedDataKey, err = fetcher()
	if err != nil {
		return nil, nil, err
	}

	ci.plaintextDataKey = plaintextDataKey
	ci.encryptedDataKey = encryptedDataKey
	ci.expiration = ci.conf.calculateExpiry()
	ci.tickets = ci.conf.maxTickets() - 1
	return keybytes.Copy(plaintextDataKey), keybytes.Copy(encryptedDataKey), nil
}

// DataKeyCacher describes the data key cache interface
type DataKeyCacher interface {
	// retrieves key from cache
	Get(compositeKey CompositeKey) *CacheItem
}

type lruCacher interface {
	Add(key, value interface{}) (evicted bool)
	Get(key interface{}) (value interface{}, ok bool)
}

type cache struct {
	keyExpiration      time.Duration
	keyExpirationSplay time.Duration
	maxUses            int64

	// mutex guards 'cache'
	mutex    sync.Mutex
	cache    lruCacher
	reporter stats.SampleReporterAPI
	keyType  KeyType
}

// KeyType is encrypt / decrypt key
type KeyType string

// encrypt / decrypt key
const (
	KeyTypeEncryption KeyType = "Encrypt"
	KeyTypeDecryption KeyType = "Decrypt"
)

// CacheConfig holds configuration variables. KeyExpiration and MaxUses cannot
// be less than 0. Either KeyExpiration or MaxUses must be set to a value
// greater than 0. If KeyExpiration is set to 0, keys will never expire and
// will stay in memory until MaxUses is consumed. If MaxUses is set to 0, keys
// will be removed in memory after the KeyExpiration period has passed.
type CacheConfig struct {
	// how long a key is valid for in cache before being removed
	KeyExpiration time.Duration
	// to avoid many keys expiring at the same time, set a splay duration. a
	// random number between [0,n] where n is the splay value will be added to
	// each key's ttl.
	KeyExpirationSplay time.Duration
	// how many times get() can be called on a key before it is removed
	MaxUses int64
	// threshold in bytes of plaintext and encrypted data keys in cache where
	// new puts will cause the least recently used key to be removed from
	// memory. if zero, math.MaxInt64
	MaxSize int64

	// Telemetry stats reporter to report cache level metrics. Default: no metrics will be reported.
	Reporter stats.SampleReporterAPI

	// KeyType - cache is used to store KMS key used to encryption/decryption.
	// This will differentiate and report encryption decryption metrics.
	// Default: Cache KMS keys used for encryption.
	KeyType KeyType
}

func (cfg *CacheConfig) validate() error {
	if cfg.KeyExpiration == 0 && cfg.MaxUses == 0 {
		return errors.New("OLECache: either KeyExpiration or MaxUses must be set to a nonzero value")
	}

	if cfg.KeyExpiration < 0 {
		return errors.New("OLECache: KeyExpiration must be greater than or equal to 0")
	}

	if cfg.KeyExpirationSplay < 0 {
		return errors.New("OLECache: KeyExpirationSplay must be greater than or equal to 0")
	}

	if cfg.MaxUses < 0 {
		return errors.New("OLECache: MaxUses must be greater than or equal to 0")
	}

	if cfg.MaxSize < 0 {
		return errors.New("OLECache: MaxSize must be greater than or equal to 0")
	}

	if cfg.MaxSize == 0 {
		cfg.MaxSize = math.MaxInt64
	}

	if cfg.KeyType == "" {
		cfg.KeyType = KeyTypeEncryption
	}

	if cfg.Reporter == nil {
		cfg.Reporter = &stats.NoopReporter{}
	}
	return nil
}

// NewCache instantiates a data key cache. can be used for both encryption and
// decryption keys
func NewCache(cfg CacheConfig) (DataKeyCacher, error) {
	err := cfg.validate()
	if err != nil {
		return nil, err
	}

	lrucache, err := simplelru.NewLRU(int(cfg.MaxSize), onEvict)
	if err != nil {
		// https://github.com/hashicorp/golang-lru/blob/f379f4cbe61dcb4acf4634238e2bef27031daf27/simplelru/lru.go#L27-L29
		// this should never err since simplelru.NewLRU() only errors on
		// size <= 0 which is caught in cfg.validate() above
		return nil, err
	}

	cache := &cache{
		keyExpiration:      cfg.KeyExpiration,
		keyExpirationSplay: cfg.KeyExpirationSplay,
		maxUses:            cfg.MaxUses,
		cache:              lrucache,
		reporter:           cfg.Reporter,
		keyType:            cfg.KeyType,
	}
	return cache, nil
}

func onEvict(key, value interface{}) {
	if item, ok := value.(*CacheItem); ok && item != nil {
		item.withLock(item.clearKeys)
	}
}

// Get retrieves a cache item
func (c *cache) Get(key CompositeKey) *CacheItem {
	k := key.Key()

	c.mutex.Lock()
	defer c.mutex.Unlock()

	v, ok := c.cache.Get(k)
	if ok {
		if item, ok := v.(*CacheItem); ok {
			return item
		}
	}

	item := &CacheItem{
		conf: c,
	}
	c.cache.Add(k, item)

	return item
}

// unixnano with splay (if exists)
func (c *cache) calculateExpiry() int64 {
	if c.keyExpiration > 0 {
		t := time.Now().Add(c.keyExpiration)
		if c.keyExpirationSplay > 0 {
			t = t.Add(time.Duration(rand.Int63n(int64(c.keyExpirationSplay))))
		}
		return t.UnixNano()
	}
	return 0
}

func (c *cache) maxTickets() int64 {
	if maxUses := c.maxUses; maxUses > 0 {
		return c.maxUses
	}
	return math.MaxInt64
}

type configurator interface {
	maxTickets() int64
	calculateExpiry() int64
}
