package awsquota

import (
	"context"
	"fmt"
	"sync"
	"testing"

	"github.com/go-logr/logr"
	"github.com/go-logr/logr/testr"
	. "github.com/onsi/gomega"
	"github.com/prometheus/client_golang/prometheus"

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

func TestExporterCollect(t *testing.T) {
	expectMetricCount := func(quotaCount, acctCount int) int {
		// quotaCount is the number of quotas fakeQuotaClient returns
		return quotaCount * acctCount * 2
	}

	cases := map[string]struct {
		acctCount int
		setup     func(*AWSQuotaExporter, *fakeQuotaClient, []config.AWSAccount)
		expect    func(*WithT, *AWSQuotaExporter, []prometheus.Metric)
	}{
		"normal": {
			acctCount: 2,
			setup: func(e *AWSQuotaExporter, cli *fakeQuotaClient, accounts []config.AWSAccount) {
				cli.quotas = []Quota{
					{ResourceID: "test"},
				}
				cli.err = nil
			},
			expect: func(g *WithT, e *AWSQuotaExporter, metrics []prometheus.Metric) {
				// check metrics
				g.Expect(metrics).To(HaveLen(expectMetricCount(1, 2)))
				// check clients
				g.Expect(e.clients).To(HaveLen(2))
			},
		},
		"remove clients": {
			acctCount: 2,
			setup: func(e *AWSQuotaExporter, cli *fakeQuotaClient, accounts []config.AWSAccount) {
				cli.quotas = []Quota{
					{ResourceID: "test"},
				}
				cli.err = nil
				e.clients = map[string]QuotaClient{
					"3": &fakeQuotaClient{}, // a client should be removed
					"4": &fakeQuotaClient{}, // a client should be removed
				}
			},
			expect: func(g *WithT, e *AWSQuotaExporter, metrics []prometheus.Metric) {
				// check metrics
				g.Expect(metrics).To(HaveLen(expectMetricCount(1, 2)))
				// check clients
				g.Expect(e.clients).To(HaveLen(2))
			},
		},
		"reuse existing clients": {
			acctCount: 2,
			setup: func(e *AWSQuotaExporter, cli *fakeQuotaClient, accounts []config.AWSAccount) {
				// new quota client return an error and should not be used.
				cli.quotas = nil
				cli.err = fmt.Errorf("should not use the new client")
				e.clients = map[string]QuotaClient{
					"0": &fakeQuotaClient{ // a existing client should be used
						quotas: []Quota{{ResourceID: "test"}},
						err:    nil,
					},
					"1": &fakeQuotaClient{ // a existing client should be used
						quotas: []Quota{{ResourceID: "test"}},
						err:    nil,
					},
				}
			},
			expect: func(g *WithT, e *AWSQuotaExporter, metrics []prometheus.Metric) {
				// check metrics
				g.Expect(metrics).To(HaveLen(expectMetricCount(1, 2)))
				// check clients
				g.Expect(e.clients).To(HaveLen(2))
			},
		},
	}

	for name, c := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			accounts := awsAccounts(c.acctCount)
			quotaClient := &fakeQuotaClient{}
			cfg := &config.AWSQuotaExporterConfig{
				AssumedRoleName: "test-role",
				AWSAccounts:     accounts,
			}
			_ = cfg.Complete()

			e := &AWSQuotaExporter{
				cfg: cfg,

				logger:    testr.New(t),
				clientset: cloudaws.NewClientSet(aws.Config{}),

				mu:      sync.Mutex{},
				clients: map[string]QuotaClient{},
				clientFactory: func(_ logr.Logger, _ cloudaws.Client, _ *cloudaws.AssumedRole, _, _ string) QuotaClient {
					return quotaClient
				},
			}

			c.setup(e, quotaClient, accounts)

			ch := make(chan prometheus.Metric, 100)
			e.Collect(ch)
			close(ch)
			metrics := []prometheus.Metric{}
			for metric := range ch {
				metrics = append(metrics, metric)
			}

			c.expect(g, e, metrics)
		})
	}
}

type fakeQuotaClient struct {
	err    error
	quotas []Quota
}

func (c *fakeQuotaClient) QueryQuotas(ctx context.Context) ([]Quota, error) {
	return c.quotas, c.err
}

func awsAccounts(count int) []config.AWSAccount {
	accounts := []config.AWSAccount{}
	for i := 0; i < count; i++ {
		accounts = append(accounts, config.AWSAccount{
			AccountID: fmt.Sprintf("%d", i),
		})
	}
	return accounts
}