package schedule

import (
	"fmt"
	"math/rand"

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

	"github.com/tidbcloud/infra-provider/pkg/account/schedule/plugin"
	"github.com/tidbcloud/infra-provider/pkg/config"
)

type Scheduler interface {
	// Schedule a network to an account
	Schedule(sctx *plugin.SchedContext, network *cloudv1.Network, accts []*cloudv1.Account) (string, error)
}

func NewScheduler(cfg *config.AcctSchedConfig, registry *plugin.Registry) (Scheduler, error) {
	filterPlugins, err := buildPlugins[plugin.FilterPlugin](cfg.Plugins, registry.FilterPlugins)
	if err != nil {
		return nil, fmt.Errorf("filter plugins: %s", err)
	}

	scorePlugins, err := buildPlugins[plugin.ScorePlugin](cfg.Plugins, registry.ScorePlugins)
	if err != nil {
		return nil, fmt.Errorf("score plugins: %s", err)
	}

	return &scheduler{
		registry:      registry,
		filterPlugins: filterPlugins,
		scorePlugins:  scorePlugins,
	}, nil
}

func buildPlugins[T plugin.Plugin](cfg config.AcctSchedPluginConfig, registry map[string]plugin.Factory) ([]T, error) {
	result := make([]T, 0)
	for name, factory := range registry {
		pl, err := factory(cfg)
		if err != nil {
			return nil, fmt.Errorf("initialize plugin %s failed: %v", name, err)
		}
		switch fpl := pl.(type) {
		case T:
			result = append(result, fpl)
		default:
			return nil, fmt.Errorf("plugin type %s is wrong", name)
		}
	}
	return result, nil
}

type scheduler struct {
	registry *plugin.Registry

	filterPlugins []plugin.FilterPlugin
	scorePlugins  []plugin.ScorePlugin
}

func (s *scheduler) Schedule(sctx *plugin.SchedContext, network *cloudv1.Network, accts []*cloudv1.Account) (string, error) {
	feasibleAccts, err := s.filterAccts(sctx, network, accts)
	if err != nil {
		return "", fmt.Errorf("filter accounts failed: %w", err)
	}

	if len(feasibleAccts) == 0 {
		return "", fmt.Errorf("no account is schedulable")
	}
	if len(feasibleAccts) == 1 {
		return feasibleAccts[0].Name, nil
	}

	acctScores, err := s.scoreAccts(sctx, network, feasibleAccts)
	if err != nil {
		return "", fmt.Errorf("score accounts failed: %w", err)
	}

	selected := s.selectAcct(sctx, acctScores)
	return selected.AccountName, nil
}

// filterAccts filters out unschedulable accounts
func (s *scheduler) filterAccts(sctx *plugin.SchedContext, network *cloudv1.Network, accts []*cloudv1.Account) ([]*cloudv1.Account, error) {
	feasibleAccts := make([]*cloudv1.Account, 0, len(accts))

	for _, plugin := range s.filterPlugins {
		result := plugin.PreFilter(sctx, network, accts)
		if result.IsError() {
			return nil, result.AsError()
		}
	}

	for _, acct := range accts {
		schedulable := true

		for _, plugin := range s.filterPlugins {
			result := plugin.Filter(sctx, network, acct)
			if result.IsError() {
				return nil, result.AsError()
			}
			if result.IsUnschedulable() {
				schedulable = false
				sctx.Logger.Info("account is unschedulable by a plugin",
					"account", acct.Name, "plugin", plugin.Name(), "reason", result.Message())
				break
			}
		}

		if schedulable {
			feasibleAccts = append(feasibleAccts, acct)
		}
	}

	return feasibleAccts, nil
}

// scoreAccts scores all feasible accounts
func (s *scheduler) scoreAccts(sctx *plugin.SchedContext, network *cloudv1.Network, accts []*cloudv1.Account) ([]plugin.AcctScore, error) {
	total := make([]plugin.AcctScore, len(accts))
	for i, acct := range accts {
		total[i] = plugin.AcctScore{
			AccountName: acct.Name,
			Score:       0,
		}
	}

	if len(s.scorePlugins) == 0 {
		for i := range total {
			total[i].Score = 1
		}
		return total, nil
	}

	for _, plugin := range s.scorePlugins {
		scores, result := plugin.Score(sctx, network, accts)
		if !result.IsSuccess() {
			return nil, fmt.Errorf("plugin %s score failed: %w", plugin.Name(), result.AsError())
		}

		sctx.Logger.V(4).WithValues("plugin", plugin.Name()).Info("score result", "scores", scores)

		for i := range scores {
			total[i].Score += scores[i].Score
		}
	}
	return total, nil
}

// selectAcct returns the highest score account
func (s *scheduler) selectAcct(sctx *plugin.SchedContext, scores []plugin.AcctScore) plugin.AcctScore {
	highestIndex := 0
	cntOfMaxScore := 1

	for i := range scores {
		if scores[i].Score > scores[highestIndex].Score {
			highestIndex = i
			cntOfMaxScore = 1
			continue
		}

		if scores[i].Score == scores[highestIndex].Score {
			cntOfMaxScore++
			// if multi accounts have the same score, returns a random account not always the first one
			if rand.Intn(cntOfMaxScore) == 0 {
				highestIndex = i
			}
		}
	}
	return scores[highestIndex]
}