package manager

import (
	"sync"

	"code.justin.tv/systems/sandstorm/resource"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/dynamodb"
	"github.com/aws/aws-sdk-go/service/dynamodb/dynamodbiface"
	"github.com/aws/aws-sdk-go/service/kms"
	"github.com/aws/aws-sdk-go/service/kms/kmsiface"
)

type clientLoader struct {
	Config     Config
	clientLock sync.Mutex
	dynamoDBs  map[string]dynamodbiface.DynamoDBAPI
	kmsClients map[string]kmsiface.KMSAPI
	sessions   map[string]*session.Session
}

func (c *clientLoader) awsConfigurer(region string) resource.AWSConfigurer {
	return c.Config.awsConfigurerForRegion(region)
}

func (c *clientLoader) session(region string) *session.Session {
	if c.sessions == nil {
		c.sessions = make(map[string]*session.Session)
	}

	sess, ok := c.sessions[region]
	if !ok {
		sess = c.awsConfigurer(region).STSSession(c.Config.awsCreds())
		c.sessions[region] = sess
	}
	return sess
}

func (c *clientLoader) KMS(region string) kmsiface.KMSAPI {
	c.clientLock.Lock()

	if c.kmsClients == nil {
		c.kmsClients = make(map[string]kmsiface.KMSAPI)
	}

	client, ok := c.kmsClients[region]
	if !ok {
		client = kms.New(c.session(region))
		c.kmsClients[region] = client
	}

	c.clientLock.Unlock()
	return client
}

func (c *clientLoader) DynamoDB(region string) dynamodbiface.DynamoDBAPI {
	c.clientLock.Lock()

	if c.dynamoDBs == nil {
		c.dynamoDBs = make(map[string]dynamodbiface.DynamoDBAPI)
	}

	ddb, ok := c.dynamoDBs[region]
	if !ok {
		ddb = dynamodb.New(c.session(region))
		c.dynamoDBs[region] = ddb
	}

	c.clientLock.Unlock()
	return ddb
}
