package ole

import (
	"errors"
	"io"
	"time"

	"code.justin.tv/amzn/TwitchOLE/ole/internal/datakeycache"
	"code.justin.tv/amzn/TwitchOLE/ole/internal/stats"
)

// NewKMSOleClient initializes an ole client backed by kms keys
func NewKMSOleClient(kms KMSAPI, cfg KMSOleClientConfig) (*KMSClient, error) {
	err := cfg.validate()
	if err != nil {
		return nil, err
	}

	client := &KMSClient{
		kms:      kms,
		cfg:      cfg,
		reporter: cfg.Reporter,
	}

	client.encryptionKeyCache, err = datakeycache.NewCache(datakeycache.CacheConfig{
		KeyExpiration:      cfg.DataKeyTTL,
		KeyExpirationSplay: cfg.DataKeyTTLSplay,
		MaxUses:            cfg.EncryptionKeyMaxUsage,
		MaxSize:            cfg.MaxCacheSize,
		Reporter:           cfg.Reporter,
		KeyType:            datakeycache.KeyTypeEncryption,
	})
	if err != nil {
		return nil, err
	}

	client.decryptionKeyCache, err = datakeycache.NewCache(datakeycache.CacheConfig{
		KeyExpiration:      cfg.DataKeyTTL,
		KeyExpirationSplay: cfg.DataKeyTTLSplay,
		MaxSize:            cfg.MaxCacheSize,
		Reporter:           cfg.Reporter,
		KeyType:            datakeycache.KeyTypeDecryption,
	})
	if err != nil {
		return nil, err
	}

	return client, nil
}

// KMSOleClientConfig contains config options for kms ole client
type KMSOleClientConfig struct {
	// ARN of the KMS CMK
	CMKArn string

	// duration that a plaintext data key should be cached in memory
	DataKeyTTL 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.
	DataKeyTTLSplay time.Duration

	// number of times a plaintext data key can be used for encryption before
	// eviction
	EncryptionKeyMaxUsage int64

	// maximum size limit for each of the data key caches (encryption & decryption)
	// past this size limit, the least recently used key will be removed from
	// memory. If zero, math.MaxInt64.
	MaxCacheSize int64

	// this will always be AES256GCM
	Algorithm Algorithm

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

func (cfg *KMSOleClientConfig) validate() error {
	if cfg.CMKArn == "" {
		return errors.New("ole: cmk arn must not be empty")
	}

	if cfg.Algorithm == nil {
		cfg.Algorithm = AES256GCM{}
	}

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

// KMSClient implements an OLE client, backed by KMS
type KMSClient struct {
	kms KMSAPI
	cfg KMSOleClientConfig

	encryptionKeyCache datakeycache.DataKeyCacher

	decryptionKeyCache datakeycache.DataKeyCacher
	//reporter
	reporter stats.SampleReporterAPI
}

// NewEncryptor returns an io.Writer that encrypts data and writes it to the
// underlying io.Writer. Each object should be written with its own encryptor.
func (k *KMSClient) NewEncryptor(encryptionContext map[string]string, w io.Writer) io.Writer {
	return &encryptionWriter{
		k:                 k,
		encryptionContext: encryptionContext,
		w:                 w,
	}
}

// NewDecryptor reads from the provided io.Reader and decrypts before returning
// bytes.
func (k *KMSClient) NewDecryptor(r io.Reader) io.Reader {
	return &decryptionReader{
		k: k,
		r: r,
	}
}
