package sts

import (
	"context"

	"code.justin.tv/eventbus/controlplane/internal/arn"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/client"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
	"github.com/aws/aws-sdk-go/service/sts"
	"github.com/aws/aws-sdk-go/service/sts/stsiface"
)

type Manager struct {
	baseConfig  client.ConfigProvider
	stsClient   stsiface.STSAPI
	environment string
}

func NewManager(c client.ConfigProvider, environment string) *Manager {
	return &Manager{
		baseConfig:  c,
		stsClient:   sts.New(c),
		environment: environment,
	}
}

func (m *Manager) AssumeRoleCredentials(accountID string) *credentials.Credentials {
	assumeRoleARN := arn.AssumeRoleARN(accountID, m.environment)
	return m.AssumeRoleCredentialsByARN(assumeRoleARN)
}

func (m *Manager) AssumeRoleCredentialsByARN(arn string) *credentials.Credentials {
	return stscreds.NewCredentials(m.baseConfig, arn)
}

func (m *Manager) GetBaseAWSAccountID(ctx context.Context) (string, error) {
	output, err := m.stsClient.GetCallerIdentityWithContext(ctx, &sts.GetCallerIdentityInput{})
	if err != nil {
		return "", err
	}
	return aws.StringValue(output.Account), nil
}
