package schedule

import (
	"context"
	"fmt"
	"math/rand"
	"testing"
	"time"

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

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	pl "github.com/tidbcloud/infra-provider/pkg/account/schedule/plugin"
	"github.com/tidbcloud/infra-provider/pkg/config"
)

func TestFilterAccts(t *testing.T) {
	cases := map[string]struct {
		pluginCount int
		acctCount   int
		input       func(*scheduler, *pl.SchedContext, *cloudv1.Network, []*cloudv1.Account)
		expect      func(*WithT, []*cloudv1.Account, error)
	}{
		"no filter plugin": {
			pluginCount: 0,
			acctCount:   3,
			input:       func(s *scheduler, sctx *pl.SchedContext, network *cloudv1.Network, accts []*cloudv1.Account) {},
			expect: func(g *WithT, accts []*cloudv1.Account, err error) {
				g.Expect(err).To(Succeed())
				g.Expect(accts).To(HaveLen(3))
				expectAccts := accounts(3)
				g.Expect(cmp.Diff(accts, expectAccts)).To(BeEmpty(), "+got, -want")
			},
		},
		"one plugin return error when prefilter": {
			pluginCount: 3,
			acctCount:   3,
			input: func(s *scheduler, sctx *pl.SchedContext, network *cloudv1.Network, accts []*cloudv1.Account) {
				for i, p := range s.filterPlugins {
					p.(*pl.FakePlugin).SetPreFilterFuncAsSuccess()
					if i == 1 {
						p.(*pl.FakePlugin).PreFilterFunc = func(sctx *pl.SchedContext, network *cloudv1.Network, accts []*cloudv1.Account) *pl.ActionResult {
							return pl.NewActionResult(pl.Error, "prefilter failed")
						}
					}
				}
			},
			expect: func(g *WithT, accts []*cloudv1.Account, err error) {
				g.Expect(err).To(HaveOccurred())
				g.Expect(err.Error()).To(ContainSubstring("prefilter failed"))
			},
		},
		"all plugins return success": {
			pluginCount: 3,
			acctCount:   3,
			input: func(s *scheduler, sctx *pl.SchedContext, network *cloudv1.Network, accts []*cloudv1.Account) {
				for _, p := range s.filterPlugins {
					p.(*pl.FakePlugin).SetPreFilterFuncAsSuccess()
					p.(*pl.FakePlugin).SetFilterFuncAsSuccess()
				}
			},
			expect: func(g *WithT, accts []*cloudv1.Account, err error) {
				g.Expect(err).To(Succeed())
				g.Expect(accts).To(HaveLen(3))
				expectAccts := accounts(3)
				g.Expect(cmp.Diff(accts, expectAccts)).To(BeEmpty(), "+got, -want")
			},
		},
		"one plugin returns unschedulable": {
			pluginCount: 3,
			acctCount:   3,
			input: func(s *scheduler, sctx *pl.SchedContext, network *cloudv1.Network, accts []*cloudv1.Account) {
				// only pl 1 think account 1 is unschedulable
				for i, p := range s.filterPlugins {
					p.(*pl.FakePlugin).SetPreFilterFuncAsSuccess()
					if i != 1 {
						p.(*pl.FakePlugin).SetFilterFuncAsSuccess()
					} else {
						p.(*pl.FakePlugin).FilterFunc = func(sc *pl.SchedContext, n *cloudv1.Network, a *cloudv1.Account) *pl.ActionResult {
							if a.Name == "account-1" {
								return pl.NewActionResult(pl.Unschedulable, "unschedulable by fake plugin")
							}
							return pl.NewActionResult(pl.Success)
						}
					}
				}
			},
			expect: func(g *WithT, accts []*cloudv1.Account, err error) {
				g.Expect(err).To(Succeed())
				g.Expect(accts).To(HaveLen(2))
				allAccts := accounts(3)
				expectAccts := append(allAccts[:1], allAccts[2:]...)
				g.Expect(cmp.Diff(accts, expectAccts)).To(BeEmpty(), "+got, -want")
			},
		},
		"one plugin returns error": {
			pluginCount: 3,
			acctCount:   3,
			input: func(s *scheduler, sctx *pl.SchedContext, network *cloudv1.Network, accts []*cloudv1.Account) {
				// pl 1 return error
				for i, p := range s.filterPlugins {
					p.(*pl.FakePlugin).SetPreFilterFuncAsSuccess()
					if i != 1 {
						p.(*pl.FakePlugin).SetFilterFuncAsSuccess()
					} else {
						p.(*pl.FakePlugin).FilterFunc = func(sc *pl.SchedContext, n *cloudv1.Network, a *cloudv1.Account) *pl.ActionResult {
							if a.Name == "account-1" {
								return pl.NewActionResult(pl.Error, "unschedulable by fake plugin")
							}
							return pl.NewActionResult(pl.Success)
						}
					}
				}
			},
			expect: func(g *WithT, accts []*cloudv1.Account, err error) {
				g.Expect(err).To(HaveOccurred())
				g.Expect(err.Error()).To(ContainSubstring("unschedulable by fake pl"))
			},
		},
		"no account": {
			pluginCount: 3,
			acctCount:   0,
			input: func(s *scheduler, sc *pl.SchedContext, n *cloudv1.Network, a []*cloudv1.Account) {
				for _, p := range s.filterPlugins {
					p.(*pl.FakePlugin).SetPreFilterFuncAsSuccess()
				}
			},
			expect: func(g *WithT, a []*cloudv1.Account, err error) {
				g.Expect(err).To(Succeed())
				g.Expect(a).Should(HaveLen(0))
			},
		},
	}

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

			cfg := &config.AcctSchedConfig{}
			registry := fakePluginRegistry(c.pluginCount)
			sched, err := NewScheduler(cfg, registry)
			g.Expect(err).Should(Succeed())
			scheduler := sched.(*scheduler)

			ctx := pl.NewSchedContext(context.TODO(), testr.New(t))
			network := &cloudv1.Network{}
			accts := accounts(c.acctCount)
			c.input(scheduler, ctx, network, accts)

			result, err := scheduler.filterAccts(ctx, network, accts)

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

func TestScoreAccts(t *testing.T) {
	cases := map[string]struct {
		pluginCount int
		acctCount   int
		input       func(*scheduler, *pl.SchedContext, *cloudv1.Network, []*cloudv1.Account)
		expect      func(*WithT, []pl.AcctScore, error)
	}{
		"no score plugin": {
			pluginCount: 0,
			acctCount:   3,
			input:       func(s *scheduler, sc *pl.SchedContext, n *cloudv1.Network, a []*cloudv1.Account) {},
			expect: func(g *WithT, scores []pl.AcctScore, err error) {
				g.Expect(err).Should(Succeed())
				g.Expect(scores).Should(HaveLen(3))
				for i := range scores {
					g.Expect(scores[i].AccountName).Should(Equal(accountName(i)))
					g.Expect(scores[i].Score).Should(Equal(int64(1)))
				}
			},
		},
		"all plugins return a score": {
			pluginCount: 3,
			acctCount:   3,
			input: func(s *scheduler, sc *pl.SchedContext, n *cloudv1.Network, a []*cloudv1.Account) {
				for _, p := range s.scorePlugins {
					p.(*pl.FakePlugin).SetScoreFunc(2)
				}
			},
			expect: func(g *WithT, scores []pl.AcctScore, err error) {
				g.Expect(err).Should(Succeed())
				g.Expect(scores).Should(HaveLen(3))
				for i := range scores {
					g.Expect(scores[i].AccountName).Should(Equal(accountName(i)))
					g.Expect(scores[i].Score).Should(Equal(int64(6))) // 3 plugins, each return 2, so the final score is 6
				}
			},
		},
		"one plugin returns error": {
			pluginCount: 3,
			acctCount:   3,
			input: func(s *scheduler, sc *pl.SchedContext, n *cloudv1.Network, a []*cloudv1.Account) {
				// plugin 1 return error
				for i, p := range s.scorePlugins {
					if i != 1 {
						p.(*pl.FakePlugin).SetScoreFunc(2)
					} else {
						p.(*pl.FakePlugin).ScoreFunc = func(sc *pl.SchedContext, n *cloudv1.Network, accts []*cloudv1.Account) ([]pl.AcctScore, *pl.ActionResult) {
							return nil, pl.NewActionResult(pl.Error, "score error by fake plugin")
						}
					}
				}
			},
			expect: func(g *WithT, scores []pl.AcctScore, err error) {
				g.Expect(err).Should(HaveOccurred())
				g.Expect(err.Error()).Should(ContainSubstring("score error by fake plugin"))
			},
		},
	}

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

			cfg := &config.AcctSchedConfig{}
			registry := fakePluginRegistry(c.pluginCount)
			sched, err := NewScheduler(cfg, registry)
			g.Expect(err).Should(Succeed())
			scheduler := sched.(*scheduler)

			ctx := pl.NewSchedContext(context.TODO(), testr.New(t))
			network := &cloudv1.Network{}
			accts := accounts(c.acctCount)
			c.input(scheduler, ctx, network, accts)

			result, err := scheduler.scoreAccts(ctx, network, accts)

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

func TestSelectAccts(t *testing.T) {
	cases := map[string]struct {
		input  func(*scheduler) []pl.AcctScore
		expect func(*WithT, []pl.AcctScore, pl.AcctScore)
	}{
		"select highest score": {
			input: func(s *scheduler) []pl.AcctScore {
				rand.Seed(time.Now().UnixNano())
				count := 100
				scores := make([]pl.AcctScore, count)
				for i := range scores {
					scores[i].AccountName = accountName(i)
					scores[i].Score = int64(rand.Intn(100))
				}
				return scores
			},
			expect: func(g *WithT, scores []pl.AcctScore, result pl.AcctScore) {
				expectScore := scores[0].Score
				for _, as := range scores {
					if as.Score > expectScore {
						expectScore = as.Score
					}
				}
				g.Expect(result.Score).Should(Equal(expectScore))
			},
		},
		"select randomly when mutli highest score": {
			input: func(s *scheduler) []pl.AcctScore {
				rand.Seed(100)
				count := 10
				scores := make([]pl.AcctScore, count)
				for i := range scores {
					scores[i].AccountName = accountName(i)
					if i == 2 || i == 7 {
						scores[i].Score = 10
					} else {
						scores[i].Score = 5
					}
				}
				return scores
			},
			expect: func(g *WithT, scores []pl.AcctScore, result pl.AcctScore) {
				maxScore := int64(10)
				allowed := map[string]struct{}{
					accountName(2): {},
					accountName(7): {},
				}
				g.Expect(result.Score).Should(Equal(maxScore))
				_, ok := allowed[result.AccountName]
				g.Expect(ok).Should(BeTrue())
			},
		},
	}

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

			ctx := pl.NewSchedContext(context.TODO(), testr.New(t))
			scores := c.input(scheduler)

			result := scheduler.selectAcct(ctx, scores)
			c.expect(g, scores, 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 fakePluginRegistry(count int) *pl.Registry {
	registry := &pl.Registry{
		FilterPlugins: map[string]pl.Factory{},
		ScorePlugins:  map[string]pl.Factory{},
	}

	for i := 0; i < count; i++ {
		name := fmt.Sprintf("fake-plugin-%d", i)
		registry.FilterPlugins[name] = pl.NewFakePluginFactory(name)
		registry.ScorePlugins[name] = pl.NewFakePluginFactory(name)
	}
	return registry
}

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