package main

import (
	"context"
	"errors"
	"fmt"
	"math"
	"slices"
	"sort"
	"strconv"
	"time"

	internalInstance "github.com/lxc/incus/v7/internal/instance"
	"github.com/lxc/incus/v7/internal/server/cluster"
	"github.com/lxc/incus/v7/internal/server/db"
	dbCluster "github.com/lxc/incus/v7/internal/server/db/cluster"
	"github.com/lxc/incus/v7/internal/server/instance"
	"github.com/lxc/incus/v7/internal/server/instance/instancetype"
	"github.com/lxc/incus/v7/internal/server/project"
	"github.com/lxc/incus/v7/internal/server/scriptlet"
	"github.com/lxc/incus/v7/internal/server/state"
	"github.com/lxc/incus/v7/internal/server/task"
	"github.com/lxc/incus/v7/shared/api"
	apiScriptlet "github.com/lxc/incus/v7/shared/api/scriptlet"
	"github.com/lxc/incus/v7/shared/logger"
	"github.com/lxc/incus/v7/shared/osarch"
)

// ServerScore represents server score taken into account during load balancing.
type ServerScore struct {
	NodeInfo  db.NodeInfo
	Resources *api.Resources
	Score     uint8
}

// ServerUsage represents current server load.
type ServerUsage struct {
	MemoryUsage uint64
	MemoryTotal uint64
	CPUUsage    float64
	CPUTotal    uint64
}

// sortAndGroupByArch sorts servers by its score and groups them by cpu architecture.
func sortAndGroupByArch(servers []*ServerScore) map[string][]*ServerScore {
	sort.Slice(servers, func(i, j int) bool {
		return servers[i].Score > servers[j].Score
	})

	result := make(map[string][]*ServerScore)
	for _, s := range servers {
		arch := s.Resources.CPU.Architecture
		_, ok := result[arch]
		if !ok {
			result[arch] = []*ServerScore{}
		}

		result[arch] = append(result[arch], s)
	}

	return result
}

// calculateScore calculates score for single server.
func calculateScore(su *ServerUsage, au *ServerUsage) uint8 {
	memoryUsage := su.MemoryUsage
	memoryTotal := su.MemoryTotal
	cpuUsage := su.CPUUsage
	cpuTotal := su.CPUTotal

	if au != nil {
		memoryUsage += au.MemoryUsage
		memoryTotal += au.MemoryTotal
		cpuUsage += au.CPUUsage
		cpuTotal += au.CPUTotal
	}

	memoryScore := uint8(float64(memoryUsage) * 100 / float64(memoryTotal))
	cpuScore := uint8((cpuUsage * 100) / float64(cpuTotal))

	return (memoryScore + cpuScore) / 2
}

// calculateServersScore calculates score based on memory and CPU usage for servers in cluster.
func calculateServersScore(s *state.State, members []db.NodeInfo) (map[string][]*ServerScore, error) {
	scores := []*ServerScore{}
	for _, member := range members {
		clusterMember, err := cluster.Connect(member.Address, s.Endpoints.NetworkCert(), s.ServerCert(), nil, true)
		if err != nil {
			return nil, fmt.Errorf("Failed to connect to cluster member: %w", err)
		}

		res, err := clusterMember.GetServerResources()
		if err != nil {
			return nil, fmt.Errorf("Failed to get resources for cluster member: %w", err)
		}

		su := &ServerUsage{
			MemoryUsage: res.Memory.Used,
			MemoryTotal: res.Memory.Total,
			CPUUsage:    res.Load.Average1Min,
			CPUTotal:    res.CPU.Total,
		}

		serverScore := calculateScore(su, nil)
		scores = append(scores, &ServerScore{NodeInfo: member, Resources: res, Score: serverScore})
	}

	return sortAndGroupByArch(scores), nil
}

// clusterRebalanceServers is responsible for instances migration from the most busy server to less busy candidates.
func clusterRebalanceServers(ctx context.Context, s *state.State, srcServer *ServerScore, candidates []*ServerScore, leaderAddress string, maxToMigrate int64) (int64, error) {
	numOfMigrated := int64(0)

	// Restrict candidates to servers less loaded than the source.
	lessLoadedCandidates := make([]*ServerScore, 0, len(candidates))
	for _, c := range candidates {
		if c.Score >= srcServer.Score {
			continue
		}

		lessLoadedCandidates = append(lessLoadedCandidates, c)
	}

	if len(lessLoadedCandidates) == 0 {
		return numOfMigrated, nil
	}

	// The default target is the least-loaded candidate (last in the sorted list).
	dstServer := lessLoadedCandidates[len(lessLoadedCandidates)-1]

	// Get the list of instances on the source.
	var dbInstances []dbCluster.Instance
	err := s.DB.Cluster.Transaction(ctx, func(ctx context.Context, tx *db.ClusterTx) error {
		var err error

		instType := instancetype.VM
		dbInstances, err = dbCluster.GetInstances(ctx, tx.Tx(), dbCluster.InstanceFilter{Node: &srcServer.NodeInfo.Name, Type: &instType})
		if err != nil {
			return fmt.Errorf("Failed to get instances: %w", err)
		}

		return nil
	})
	if err != nil {
		return -1, fmt.Errorf("Failed to get instances: %w", err)
	}

	// Filter for instances that can be live migrated to a new target.
	var instances []instance.Instance
	for _, dbInst := range dbInstances {
		inst, err := instance.LoadByProjectAndName(s, dbInst.Project, dbInst.Name)
		if err != nil {
			return -1, fmt.Errorf("Failed to load instance: %w", err)
		}

		// Do not allow to migrate instance which doesn't support live migration.
		if inst.CanMigrate() != "live-migrate" {
			continue
		}

		// Check if instance is ready for next migration.
		lastMove := inst.LocalConfig()["volatile.rebalance.last_move"]
		cooldown := s.GlobalConfig.ClusterRebalanceCooldown()
		if lastMove != "" {
			v, err := strconv.ParseInt(lastMove, 10, 64)
			if err != nil {
				return -1, fmt.Errorf("Failed to parse last_move value: %w", err)
			}

			expiry, err := internalInstance.GetExpiry(time.Unix(v, 0), cooldown)
			if err != nil {
				return -1, fmt.Errorf("Failed to calculate expiration for cooldown time: %w", err)
			}

			if time.Now().Before(expiry) {
				continue
			}
		}

		instances = append(instances, inst)
	}

	// Map candidate name to its score data for quick lookup.
	candidateByName := make(map[string]*ServerScore, len(lessLoadedCandidates))
	for _, c := range lessLoadedCandidates {
		candidateByName[c.NodeInfo.Name] = c
	}

	// Track running usage and score per target so multiple instances heading to the same target accumulate correctly.
	runningUsage := make(map[string]*ServerUsage, len(lessLoadedCandidates))
	runningScore := make(map[string]uint8, len(lessLoadedCandidates))
	for _, c := range lessLoadedCandidates {
		runningUsage[c.NodeInfo.Name] = &ServerUsage{
			MemoryUsage: c.Resources.Memory.Used,
			MemoryTotal: c.Resources.Memory.Total,
			CPUUsage:    c.Resources.Load.Average1Min,
			CPUTotal:    c.Resources.CPU.Total,
		}

		runningScore[c.NodeInfo.Name] = c.Score
	}

	placementScriptletEnabled := s.GlobalConfig.InstancesPlacementScriptlet() != ""

	// Prepare the source API client.
	srcClient, err := cluster.Connect(srcServer.NodeInfo.Address, s.Endpoints.NetworkCert(), s.ServerCert(), nil, true)
	if err != nil {
		return -1, fmt.Errorf("Failed to connect to cluster member: %w", err)
	}

	for _, inst := range instances {
		if numOfMigrated >= maxToMigrate {
			// We're done moving instances for now.
			return numOfMigrated, nil
		}

		// Filter the candidate list for this instance using project restrictions.
		var instanceCandidates []*ServerScore
		err = s.DB.Cluster.Transaction(ctx, func(ctx context.Context, tx *db.ClusterTx) error {
			dbProject, err := dbCluster.GetProject(ctx, tx.Tx(), inst.Project().Name)
			if err != nil {
				return fmt.Errorf("Failed to get project: %w", err)
			}

			apiProject, err := dbProject.ToAPI(ctx, tx.Tx())
			if err != nil {
				return fmt.Errorf("Failed to load project: %w", err)
			}

			for _, c := range lessLoadedCandidates {
				_, _, err := project.CheckTarget(ctx, s.Authorizer, nil, tx, apiProject, c.NodeInfo.Name, []db.NodeInfo{c.NodeInfo})
				if err != nil {
					continue
				}

				instanceCandidates = append(instanceCandidates, c)
			}

			return nil
		})
		if err != nil {
			return -1, fmt.Errorf("Failed to filter candidates for instance %q in project %q: %w", inst.Name(), inst.Project().Name, err)
		}

		if len(instanceCandidates) == 0 {
			// No allowed targets for this instance.
			continue
		}

		// Default target is the least-loaded allowed candidate.
		instanceDstServer := instanceCandidates[len(instanceCandidates)-1]
		chosenTarget := &instanceDstServer.NodeInfo

		// Default fallback: if scriptlet not enabled and the global least-loaded candidate is allowed, prefer it.
		if !placementScriptletEnabled {
			for _, c := range instanceCandidates {
				if c.NodeInfo.Name == dstServer.NodeInfo.Name {
					chosenTarget = &dstServer.NodeInfo
					break
				}
			}
		}

		if placementScriptletEnabled {
			archName, err := osarch.ArchitectureName(inst.Architecture())
			if err != nil {
				return -1, fmt.Errorf("Failed getting architecture for instance %q in project %q: %w", inst.Name(), inst.Project().Name, err)
			}

			profileNames := make([]string, 0, len(inst.Profiles()))
			for _, p := range inst.Profiles() {
				profileNames = append(profileNames, p.Name)
			}

			placementReq := apiScriptlet.InstancePlacement{
				InstancesPost: api.InstancesPost{
					Name: inst.Name(),
					Type: api.InstanceType(inst.Type().String()),
					InstancePut: api.InstancePut{
						Architecture: archName,
						Config:       inst.ExpandedConfig(),
						Devices:      inst.ExpandedDevices().CloneNative(),
						Profiles:     profileNames,
					},
				},
				Project: inst.Project().Name,
				Reason:  apiScriptlet.InstancePlacementReasonRebalance,
			}

			// Build the scriptlet candidate list sorted from least to most loaded.
			sortedCandidates := make([]db.NodeInfo, 0, len(instanceCandidates))
			for _, instanceCandidate := range slices.Backward(instanceCandidates) {
				sortedCandidates = append(sortedCandidates, instanceCandidate.NodeInfo)
			}

			scriptCtx, cancel := context.WithTimeout(ctx, time.Second*5)
			scriptTarget, err := scriptlet.InstancePlacementRun(scriptCtx, logger.Log, s, &placementReq, sortedCandidates, leaderAddress)
			cancel()
			if err != nil {
				return -1, fmt.Errorf("Failed instance placement scriptlet for instance %q in project %q: %w", inst.Name(), inst.Project().Name, err)
			}

			if scriptTarget != nil {
				chosenTarget = scriptTarget
			}
		}

		// Skip if the chosen target is the source itself.
		if chosenTarget.Name == srcServer.NodeInfo.Name {
			continue
		}

		// Look up the score data for the chosen target.
		chosenScore, ok := candidateByName[chosenTarget.Name]
		if !ok {
			// Chosen target isn't in our candidate list (shouldn't happen).
			continue
		}

		// Per-target target score: midpoint between source and chosen target initial loads.
		targetScore := (srcServer.Score + chosenScore.Score) / 2
		if runningScore[chosenTarget.Name] >= targetScore {
			// We've already balanced load against this target.
			continue
		}

		// Calculate resource consumption.
		cpuUsage, memUsage, _, err := instance.ResourceUsage(inst.ExpandedConfig(), inst.ExpandedDevices().CloneNative(), api.InstanceType(inst.Type().String()))
		if err != nil {
			return -1, fmt.Errorf("Failed to establish instance resource usage: %w", err)
		}

		// Calculate impact of migration.
		additionalUsage := &ServerUsage{
			MemoryUsage: uint64(cpuUsage),
			CPUUsage:    float64(memUsage),
		}

		expectedScore := calculateScore(runningUsage[chosenTarget.Name], additionalUsage)
		if expectedScore >= targetScore {
			// Skip the instance as it would have too big an impact.
			continue
		}

		// Prepare for live migration.
		req := api.InstancePost{
			Migration: true,
			Live:      true,
		}

		targetClient := srcClient.UseTarget(chosenTarget.Name)

		migrationOp, err := targetClient.MigrateInstance(inst.Name(), req)
		if err != nil {
			return -1, fmt.Errorf("Migration API failure: %w", err)
		}

		err = migrationOp.Wait()
		if err != nil {
			return -1, fmt.Errorf("Failed to wait for migration to finish: %w", err)
		}

		// Record the migration in the instance volatile storage.
		err = inst.VolatileSet(map[string]string{"volatile.rebalance.last_move": strconv.FormatInt(time.Now().Unix(), 10)})
		if err != nil {
			return -1, err
		}

		// Update counters and per-target running state.
		numOfMigrated++
		runningScore[chosenTarget.Name] = expectedScore
		runningUsage[chosenTarget.Name].MemoryUsage += additionalUsage.MemoryUsage
		runningUsage[chosenTarget.Name].CPUUsage += additionalUsage.CPUUsage
	}

	return numOfMigrated, nil
}

// clusterRebalance performs cluster re-balancing.
func clusterRebalance(ctx context.Context, s *state.State, servers map[string][]*ServerScore, leaderAddress string) error {
	rebalanceThreshold := s.GlobalConfig.ClusterRebalanceThreshold()
	rebalanceBatch := s.GlobalConfig.ClusterRebalanceBatch()
	numOfMigrated := int64(0)

	for archName, v := range servers {
		if numOfMigrated >= rebalanceBatch {
			// Maximum number of instances already migrated in this run.
			continue
		}

		if len(v) < 2 {
			// Skip if there isn't at least 2 servers with specific arch.
			continue
		}

		if v[0].Score == 0 {
			// Don't migrate anything if most loaded isn't loaded.
			continue
		}

		leastBusyIndex := len(v) - 1
		percentageChange := int64(float64(v[0].Score-v[leastBusyIndex].Score) / float64(v[0].Score) * 100)
		logger.Debug("Automatic re-balancing", logger.Ctx{"Architecture": archName, "LeastBusy": v[leastBusyIndex].NodeInfo.Name, "LeastBusyScore": v[leastBusyIndex].Score, "MostBusy": v[0].NodeInfo.Name, "MostBusyScore": v[0].Score, "Difference": fmt.Sprintf("%d%%", percentageChange), "Threshold": fmt.Sprintf("%d%%", rebalanceThreshold)})

		if percentageChange < rebalanceThreshold {
			continue // Skip as threshold condition is not met.
		}

		n, err := clusterRebalanceServers(ctx, s, v[0], v[1:], leaderAddress, rebalanceBatch-numOfMigrated)
		if err != nil {
			return fmt.Errorf("Failed to rebalance cluster: %w", err)
		}

		numOfMigrated += n
	}

	return nil
}

func autoRebalanceCluster(ctx context.Context, d *Daemon) error {
	s := d.State()

	// Confirm we should run the rebalance.
	leader, err := s.Cluster.LeaderAddress()
	if err != nil {
		if errors.Is(err, cluster.ErrNodeIsNotClustered) {
			// Not clustered.
			return nil
		}

		return fmt.Errorf("Failed to get leader cluster member address: %w", err)
	}

	if s.LocalConfig.ClusterAddress() != leader {
		// Not the leader.
		return nil
	}

	// Get all online members
	var onlineMembers []db.NodeInfo
	err = s.DB.Cluster.Transaction(ctx, func(ctx context.Context, tx *db.ClusterTx) error {
		members, err := tx.GetNodes(ctx)
		if err != nil {
			return fmt.Errorf("Failed getting cluster members: %w", err)
		}

		onlineMembers, err = tx.GetCandidateMembers(ctx, members, nil, "", nil, s.GlobalConfig.OfflineThreshold())
		if err != nil {
			return fmt.Errorf("Failed getting online cluster members: %w", err)
		}

		return nil
	})
	if err != nil {
		return fmt.Errorf("Failed getting cluster members: %w", err)
	}

	servers, err := calculateServersScore(s, onlineMembers)
	if err != nil {
		return fmt.Errorf("Failed calculating servers score: %w", err)
	}

	err = clusterRebalance(ctx, s, servers, leader)
	if err != nil {
		return fmt.Errorf("Failed rebalancing cluster: %w", err)
	}

	return nil
}

func autoRebalanceClusterTask(d *Daemon) (task.Func, task.Schedule) {
	f := func(ctx context.Context) {
		s := d.State()

		// Check that we should run now.
		interval := s.GlobalConfig.ClusterRebalanceInterval()
		if interval <= 0 {
			// Re-balance is disabled.
			return
		}

		now := time.Now()
		elapsed := int64(math.Round(now.Sub(s.StartTime).Minutes()))
		if elapsed%interval != 0 {
			// It's not time for a re-balance.
			return
		}

		// Run the rebalance.
		err := autoRebalanceCluster(ctx, d)
		if err != nil {
			logger.Error("Failed during cluster auto rebalancing", logger.Ctx{"err": err})
		}
	}

	return f, task.Every(time.Minute, task.SkipFirst)
}
