package awsquota

import (
	"context"
	"sync"

	"github.com/go-logr/logr"
	"github.com/prometheus/client_golang/prometheus"
	"golang.org/x/exp/slices"

	cloudaws "github.com/tidbcloud/infra-provider/pkg/cloud/aws"
	"github.com/tidbcloud/infra-provider/pkg/util"
	"github.com/tidbcloud/infra-provider/quota-exporter/config"
	"github.com/tidbcloud/infra-provider/quota-exporter/exporter"
	"github.com/tidbcloud/infra-provider/quota-exporter/exporter/cache"
)

const (
	Name = "awsQuota"
)

// awsAccount is the internal representation of an AWS account read from
// the static config.
type awsAccount struct {
	AccountID string
}

type AWSQuotaExporter struct {
	cfg *config.AWSQuotaExporterConfig

	logger    logr.Logger
	clientset cloudaws.Client

	mu            sync.Mutex
	clients       map[string]QuotaClient
	clientFactory func(logr.Logger, cloudaws.Client, *cloudaws.AssumedRole, string, string) QuotaClient // for test
}

func NewExporter(cfg config.ExporterConfig, deps *exporter.Dependency) (exporter.Exporter, error) {
	myCfg := &config.AWSQuotaExporterConfig{}
	if cfg.AWSQuotaExporterConfig != nil {
		myCfg = cfg.AWSQuotaExporterConfig
	}
	err := myCfg.Complete()
	if err != nil {
		return nil, err
	}

	e := &AWSQuotaExporter{
		cfg:           myCfg,
		logger:        deps.Logger.WithName(Name),
		clientset:     deps.Clientset,
		mu:            sync.Mutex{},
		clients:       map[string]QuotaClient{},
		clientFactory: newQuotaClient,
	}

	if e.cfg.Cache != nil {
		return cache.NewAndRun(context.TODO(), e, e.cfg.Cache, e.logger)
	}
	return e, nil
}

func (e *AWSQuotaExporter) Name() string {
	return Name
}

func (e *AWSQuotaExporter) MetricPath() string {
	return e.cfg.MetricPath
}

func (e *AWSQuotaExporter) Describe(ch chan<- *prometheus.Desc) {
	ch <- descForQuotaUsage()
	ch <- descForQuotaLimit()
}

func (e *AWSQuotaExporter) Collect(ch chan<- prometheus.Metric) {
	accounts := make([]awsAccount, 0, len(e.cfg.AWSAccounts))
	for _, acct := range e.cfg.AWSAccounts {
		accounts = append(accounts, awsAccount{
			AccountID: acct.AccountID,
		})
	}

	e.logger.V(4).Info("collect quota for accounts", "accounts",
		util.MapSlice(accounts, func(a awsAccount) string { return a.AccountID }))

	e.syncClients(accounts)

	wg := sync.WaitGroup{}
	wg.Add(len(accounts))
	for i := range accounts {
		acct := accounts[i]
		go func() {
			e.collectOneAccount(ch, acct)
			wg.Done()
		}()
	}

	wg.Wait()
}

func (e *AWSQuotaExporter) collectOneAccount(ch chan<- prometheus.Metric, account awsAccount) {
	accountID := account.AccountID

	client, exist := e.getClient(accountID)
	if !exist {
		return
	}
	quotas, err := client.QueryQuotas(context.TODO())
	if err != nil {
		e.logger.Error(err, "failed to query quotas for an account", "account", accountID)
		return
	}

	for _, quota := range quotas {
		ch <- metricForQuotaUsage(quota, accountID)
		ch <- metricForQuotaLimit(quota, accountID)
	}
}

func (e *AWSQuotaExporter) syncClients(accounts []awsAccount) {
	roleName := e.cfg.AssumedRoleName
	externalID := e.cfg.ExternalID

	e.mu.Lock()
	defer e.mu.Unlock()

	for _, account := range accounts {
		accountID := account.AccountID
		if _, exist := e.clients[accountID]; exist {
			continue
		}
		assumedRole := &cloudaws.AssumedRole{
			RoleARN: util.GetRoleARN(accountID, roleName),
		}
		if externalID != nil {
			assumedRole.ExternalID = util.Pointer(*externalID)
		}

		e.clients[accountID] = e.clientFactory(
			e.logger.WithValues("account", accountID),
			e.clientset,
			assumedRole,
			accountID,
			e.cfg.GCPProjectID,
		)
		e.logger.Info("created a new quota client", "account", accountID)
	}

	for accountID := range e.clients {
		pred := func(acct awsAccount) bool {
			return acct.AccountID == accountID
		}

		if slices.IndexFunc(accounts, pred) < 0 {
			delete(e.clients, accountID)
			e.logger.Info("deleted a quota client", "account", accountID)
		}
	}
}

func (e *AWSQuotaExporter) getClient(accountID string) (QuotaClient, bool) {
	e.mu.Lock()
	defer e.mu.Unlock()

	client, exist := e.clients[accountID]
	return client, exist
}

func descForQuotaUsage() *prometheus.Desc {
	fqName := prometheus.BuildFQName("awsaccount", "quota", "usage")
	help := "Current used amount of the given resource."
	varLabels := []string{"aws_account", "region", "resource", "service"}
	constLabels := prometheus.Labels{}
	return prometheus.NewDesc(fqName, help, varLabels, constLabels)
}

func descForQuotaLimit() *prometheus.Desc {
	fqName := prometheus.BuildFQName("awsaccount", "quota", "limit")
	help := "Current limit of the given resource."
	varLabels := []string{"aws_account", "region", "resource", "service"}
	constLabels := prometheus.Labels{}
	return prometheus.NewDesc(fqName, help, varLabels, constLabels)
}

func metricForQuotaUsage(quota Quota, account string) prometheus.Metric {
	return prometheus.MustNewConstMetric(
		descForQuotaUsage(), prometheus.GaugeValue, quota.Usage,
		account, quota.Region, quota.Resource, quota.Service,
	)
}

func metricForQuotaLimit(quota Quota, account string) prometheus.Metric {
	return prometheus.MustNewConstMetric(
		descForQuotaLimit(), prometheus.GaugeValue, quota.Limit,
		account, quota.Region, quota.Resource, quota.Service,
	)
}