package sqs

import (
	"context"

	"code.justin.tv/eventbus/controlplane/internal/sqsutil"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/pkg/errors"
)

type CredentialFetcher interface {
	AssumeRoleCredentials(accountID string) *credentials.Credentials
}

type Manager struct {
	baseConfig  *aws.Config // assume role sessions will be spawned off of this base configuration
	roleAssumer CredentialFetcher
}

func NewManager(baseSession *session.Session, assumer CredentialFetcher) *Manager {
	// Technically, only the underlying config from the session is needed to create
	// an sqs.Manager, but the constructor uses the wrapping session struct to be
	// consistent with other "Manager" constructors
	var cfg *aws.Config
	if baseSession == nil || baseSession.Config == nil {
		cfg = aws.NewConfig()
	} else {
		cfg = baseSession.Config
	}

	return &Manager{
		baseConfig:  cfg,
		roleAssumer: assumer,
	}
}

func (m *Manager) GetQueueAttributes(ctx context.Context, url string) (map[string]string, error) {
	accountID, err := sqsutil.AccountIDFromQueueURL(url)
	if err != nil {
		return nil, errors.Wrap(err, "could not get aws account id from sqs queue url")
	}

	sess, err := m.getSession(accountID)
	if err != nil {
		return nil, errors.Wrap(err, "could not initialize session")
	}

	sqsClient := sqs.New(sess)
	output, err := sqsClient.GetQueueAttributesWithContext(ctx, &sqs.GetQueueAttributesInput{
		QueueUrl:       aws.String(url),
		AttributeNames: aws.StringSlice([]string{sqsutil.KeyAll}),
	})
	if err != nil {
		return nil, errors.Wrap(err, "could not fetch queue attributes")
	}

	return aws.StringValueMap(output.Attributes), nil
}

func (m *Manager) GetQueueURL(ctx context.Context, name, accountID string) (string, error) {
	sess, err := m.getSession(accountID)
	if err != nil {
		return "", errors.Wrap(err, "could not initialize session")
	}

	sqsClient := sqs.New(sess)
	output, err := sqsClient.GetQueueUrlWithContext(ctx, &sqs.GetQueueUrlInput{
		QueueName: aws.String(name),
	})
	if err != nil {
		return "", errors.Wrap(err, "could not get queue url")
	}

	return aws.StringValue(output.QueueUrl), nil
}

func (m *Manager) getSession(accountID string) (*session.Session, error) {
	assumeRoleCreds := m.roleAssumer.AssumeRoleCredentials(accountID)

	cfgCpy := m.baseConfig.Copy().WithCredentials(assumeRoleCreds)
	sess, err := session.NewSession(cfgCpy)
	if err != nil {
		return nil, errors.Wrap(err, "could not create assume role session")
	}

	return sess, nil
}
