package plugin

import (
	"context"
	"fmt"
	"testing"

	"github.com/go-logr/logr/testr"
	. "github.com/onsi/gomega"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
)

func TestPreFilter(t *testing.T) {
	cases := map[string]struct {
		acctCount int
		input     func(*SchedContext, *cloudv1.Network, []*cloudv1.Account)
		expect    func(*WithT, *SchedContext, *ActionResult)
	}{
		"no accounts": {
			acctCount: 0,
			input:     func(ctx *SchedContext, net *cloudv1.Network, accts []*cloudv1.Account) {},
			expect: func(g *WithT, sctx *SchedContext, result *ActionResult) {
				state, ok := sctx.State[dedicatedAccountPluginName].(*dedicatedAccountState)
				g.Expect(ok).Should(BeTrue(), "state should be dedicatedAccountState")
				g.Expect(state.acctForProjects).ShouldNot(BeNil())
				g.Expect(state.acctForTenants).ShouldNot(BeNil())
				g.Expect(state.acctForProjects).Should(BeEmpty())
				g.Expect(state.acctForTenants).Should(BeEmpty())
			},
		},
		"no dedicated accounts": {
			acctCount: 3,
			input:     func(ctx *SchedContext, net *cloudv1.Network, accts []*cloudv1.Account) {},
			expect: func(g *WithT, sctx *SchedContext, result *ActionResult) {
				state, ok := sctx.State[dedicatedAccountPluginName].(*dedicatedAccountState)
				g.Expect(ok).Should(BeTrue(), "state should be dedicatedAccountState")
				g.Expect(state.acctForProjects).ShouldNot(BeNil())
				g.Expect(state.acctForTenants).ShouldNot(BeNil())
				g.Expect(state.acctForProjects).Should(BeEmpty())
				g.Expect(state.acctForTenants).Should(BeEmpty())
			},
		},
		"one dedicated account without projects and tenants": {
			acctCount: 3,
			input: func(ctx *SchedContext, net *cloudv1.Network, accts []*cloudv1.Account) {
				accts[1].Annotations = map[string]string{}
				accts[1].Annotations[dedicatedAcctAnnoForProjects] = ""
				accts[1].Annotations[dedicatedAcctAnnoForTenants] = ""
			},
			expect: func(g *WithT, sctx *SchedContext, result *ActionResult) {
				state, ok := sctx.State[dedicatedAccountPluginName].(*dedicatedAccountState)
				g.Expect(ok).Should(BeTrue(), "state should be dedicatedAccountState")
				projects, exist := state.acctForProjects[accountName(1)]
				g.Expect(exist).Should(BeTrue())
				g.Expect(projects).Should(BeEmpty())
				tenant, exist := state.acctForTenants[accountName(1)]
				g.Expect(exist).Should(BeTrue())
				g.Expect(tenant).Should(BeEmpty())
			},
		},
		"dedicated accounts with projects": {
			acctCount: 3,
			input: func(ctx *SchedContext, net *cloudv1.Network, accts []*cloudv1.Account) {
				accts[1].Annotations = map[string]string{}
				accts[1].Annotations[dedicatedAcctAnnoForProjects] = "1215,6266,123415"
				accts[2].Annotations = map[string]string{}
				accts[2].Annotations[dedicatedAcctAnnoForProjects] = "ava1"
			},
			expect: func(g *WithT, sctx *SchedContext, result *ActionResult) {
				state, ok := sctx.State[dedicatedAccountPluginName].(*dedicatedAccountState)
				g.Expect(ok).Should(BeTrue(), "state should be dedicatedAccountState")
				projects1, exist := state.acctForProjects[accountName(1)]
				g.Expect(exist).Should(BeTrue())
				g.Expect(projects1).Should(Equal([]string{"1215", "6266", "123415"}))
				tenant1, exist := state.acctForTenants[accountName(1)]
				g.Expect(exist).Should(BeTrue())
				g.Expect(tenant1).Should(BeEmpty())
				projects2, exist := state.acctForProjects[accountName(2)]
				g.Expect(exist).Should(BeTrue())
				g.Expect(projects2).Should(Equal([]string{"ava1"}))
				tenant2, exist := state.acctForTenants[accountName(2)]
				g.Expect(exist).Should(BeTrue())
				g.Expect(tenant2).Should(BeEmpty())
			},
		},
		"dedicated accounts with tenants": {
			acctCount: 3,
			input: func(ctx *SchedContext, net *cloudv1.Network, accts []*cloudv1.Account) {
				accts[1].Annotations = map[string]string{}
				accts[1].Annotations[dedicatedAcctAnnoForTenants] = "1215,6266,123415"
				accts[2].Annotations = map[string]string{}
				accts[2].Annotations[dedicatedAcctAnnoForTenants] = "ava1"
			},
			expect: func(g *WithT, sctx *SchedContext, result *ActionResult) {
				state, ok := sctx.State[dedicatedAccountPluginName].(*dedicatedAccountState)
				g.Expect(ok).Should(BeTrue(), "state should be dedicatedAccountState")
				projects1, exist := state.acctForProjects[accountName(1)]
				g.Expect(exist).Should(BeTrue())
				g.Expect(projects1).Should(BeEmpty())
				tenant1, exist := state.acctForTenants[accountName(1)]
				g.Expect(exist).Should(BeTrue())
				g.Expect(tenant1).Should(Equal([]string{"1215", "6266", "123415"}))
				projects2, exist := state.acctForProjects[accountName(2)]
				g.Expect(exist).Should(BeTrue())
				g.Expect(projects2).Should(BeEmpty())
				tenant2, exist := state.acctForTenants[accountName(2)]
				g.Expect(exist).Should(BeTrue())
				g.Expect(tenant2).Should(Equal([]string{"ava1"}))
			},
		},
	}

	for name, c := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			network := &cloudv1.Network{}
			network.Name = "test-network"
			network.Namespace = "test-namespace"
			accts := accounts(c.acctCount)
			sctx := NewSchedContext(context.TODO(), testr.New(t))

			c.input(sctx, network, accts)

			p := &dedicatedAccount{}
			result := p.PreFilter(sctx, network, accts)

			c.expect(g, sctx, result)
		})
	}
}

//nolint:goconst
func TestFilter(t *testing.T) {
	cases := map[string]struct {
		input  func(*SchedContext, *cloudv1.Network, *cloudv1.Account)
		expect func(*WithT, *ActionResult)
	}{
		"success if no dedicated acount": {
			input: func(sc *SchedContext, n *cloudv1.Network, a *cloudv1.Account) {
				sc.State[dedicatedAccountPluginName] = &dedicatedAccountState{}
			},
			expect: func(g *WithT, result *ActionResult) {
				g.Expect(result).Should(Equal(NewActionResult(Success)))
			},
		},
		"success if the account is not dedicated and the network no need a dedicated account": {
			input: func(sc *SchedContext, network *cloudv1.Network, acct *cloudv1.Account) {
				sc.State[dedicatedAccountPluginName] = &dedicatedAccountState{
					acctForProjects: map[string][]string{
						"account-2": {"project-2"},
					},
					acctForTenants: map[string][]string{
						"account-2": {"tenant-2"},
					},
				}
				network.Labels[cloudv1.ProjectTagKey] = "project-1"
				network.Labels[cloudv1.TenantTagKey] = "tenant-1"
			},
			expect: func(g *WithT, result *ActionResult) {
				g.Expect(result).Should(Equal(NewActionResult(Success)))
			},
		},
		"success if the account is dedicated for the network by project id": {
			input: func(sc *SchedContext, network *cloudv1.Network, acct *cloudv1.Account) {
				sc.State[dedicatedAccountPluginName] = &dedicatedAccountState{
					acctForProjects: map[string][]string{
						"account-1": {"project-1"},
					},
					acctForTenants: map[string][]string{
						"account-2": {"tenant-2"},
					},
				}
				network.Labels[cloudv1.ProjectTagKey] = "project-1"
				network.Labels[cloudv1.TenantTagKey] = "tenant-1"
			},
			expect: func(g *WithT, result *ActionResult) {
				g.Expect(result).Should(Equal(NewActionResult(Success)))
			},
		},
		"success if the account is dedicated for the network by tenant id": {
			input: func(sc *SchedContext, network *cloudv1.Network, acct *cloudv1.Account) {
				sc.State[dedicatedAccountPluginName] = &dedicatedAccountState{
					acctForProjects: map[string][]string{
						"account-2": {"project-2"},
					},
					acctForTenants: map[string][]string{
						"account-1": {"tenant-1"},
					},
				}
				network.Labels[cloudv1.ProjectTagKey] = "project-1"
				network.Labels[cloudv1.TenantTagKey] = "tenant-1"
			},
			expect: func(g *WithT, result *ActionResult) {
				g.Expect(result).Should(Equal(NewActionResult(Success)))
			},
		},
		"unschedulable if the account is dedicated but not for the network": {
			input: func(sc *SchedContext, network *cloudv1.Network, acct *cloudv1.Account) {
				sc.State[dedicatedAccountPluginName] = &dedicatedAccountState{
					acctForProjects: map[string][]string{
						"account-1": {},
					},
				}
				network.Labels[cloudv1.ProjectTagKey] = "project-1"
				network.Labels[cloudv1.TenantTagKey] = "tenant-1"
			},
			expect: func(g *WithT, result *ActionResult) {
				g.Expect(result).Should(Equal(NewActionResult(Unschedulable, "can not use the dedicated account account-1")))
			},
		},
		"unschedulable if the account is not dedicated but the network needs a dedicated account by project id": {
			input: func(sc *SchedContext, network *cloudv1.Network, acct *cloudv1.Account) {
				sc.State[dedicatedAccountPluginName] = &dedicatedAccountState{
					acctForProjects: map[string][]string{
						"account-2": {"project-2"},
					},
				}
				network.Labels[cloudv1.ProjectTagKey] = "project-2"
				network.Labels[cloudv1.TenantTagKey] = "tenant-2"
			},
			expect: func(g *WithT, result *ActionResult) {
				g.Expect(result).Should(Equal(NewActionResult(Unschedulable, "must use account account-2")))
			},
		},
		"unschedulable if the account is not dedicated but the network needs a dedicated account by tenant id": {
			input: func(sc *SchedContext, network *cloudv1.Network, acct *cloudv1.Account) {
				sc.State[dedicatedAccountPluginName] = &dedicatedAccountState{
					acctForTenants: map[string][]string{
						"account-2": {"tenant-2"},
					},
				}
				network.Labels[cloudv1.ProjectTagKey] = "project-2"
				network.Labels[cloudv1.TenantTagKey] = "tenant-2"
			},
			expect: func(g *WithT, result *ActionResult) {
				g.Expect(result).Should(Equal(NewActionResult(Unschedulable, "must use account account-2")))
			},
		},
	}

	for name, c := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewWithT(t)

			network := &cloudv1.Network{}
			network.Name = "test-network"
			network.Namespace = "test-namespace"
			network.Labels = map[string]string{}
			account := &cloudv1.Account{}
			account.Name = "account-1"
			sctx := NewSchedContext(context.TODO(), testr.New(t))

			c.input(sctx, network, account)

			p := &dedicatedAccount{}
			result := p.Filter(sctx, network, account)

			c.expect(g, result)
		})
	}
}

func accounts(count int) []*cloudv1.Account {
	accts := make([]*cloudv1.Account, count)
	for i := 0; i < count; i++ {
		accts[i] = &cloudv1.Account{
			ObjectMeta: metav1.ObjectMeta{
				Name: accountName(i),
			},
		}
	}
	return accts
}

func accountName(index int) string {
	return fmt.Sprintf("account-%d", index)
}