package glm

import (
	"isle/daemon/daecommon"
	"isle/garage"
)

func allocsByRPCPort(
	allocs []daecommon.ConfigStorageAllocation,
) map[int]daecommon.ConfigStorageAllocation {
	m := map[int]daecommon.ConfigStorageAllocation{}
	for _, alloc := range allocs {
		m[alloc.RPCPort] = alloc
	}
	return m
}

// AllocationWithKnownNode pairs an active storage allocation with its
// corresponding status information from garage.
type AllocationWithKnownNode struct {
	garage.KnownNode
	daecommon.ConfigStorageAllocation
}

func allocsWithKnownNode(
	allocs []daecommon.ConfigStorageAllocation,
	knownNodes []garage.KnownNode,
) (
	res []AllocationWithKnownNode,
) {
outer:
	for _, alloc := range allocs {
		for _, knownNode := range knownNodes {
			if alloc.RPCPort == int(knownNode.Addr.Port()) {
				res = append(res, AllocationWithKnownNode{knownNode, alloc})
				continue outer
			}
		}

		res = append(res, AllocationWithKnownNode{ConfigStorageAllocation: alloc})
	}

	return
}

// StateTransition describes the allocation changes which should be made as a
// result of the current state of the cluster and the target allocation state.
type StateTransition struct {
	// AddModifyAllocations should be added to the garage cluster layout, if
	// they are not already.
	AddModifyAllocations []daecommon.ConfigStorageAllocation

	// DrainAllocations should be "removed" from the garage cluster layout, if
	// the are not already. In reality "removing" from the garage cluster
	// layout only removes the node's role, and puts it into the draining state.
	// The node itself remains active.
	DrainAllocations []AllocationWithKnownNode
}

// DrainAllocationIDs returns the IDs of all allocations in the DrainAllocations
// field.
func (st StateTransition) DrainAllocationIDs() []string {
	ids := make([]string, len(st.DrainAllocations))
	for i, alloc := range st.DrainAllocations {
		ids[i] = alloc.ID
	}
	return ids
}

// ActiveAllocations returns all allocations which should be active, ie should
// have corresponding garage instances.
func (st StateTransition) ActiveAllocations() []daecommon.ConfigStorageAllocation {
	allocs := make(
		[]daecommon.ConfigStorageAllocation,
		0,
		len(st.AddModifyAllocations)+len(st.DrainAllocations),
	)

	for _, alloc := range st.AddModifyAllocations {
		allocs = append(allocs, alloc)
	}

	for _, alloc := range st.DrainAllocations {
		allocs = append(allocs, alloc.ConfigStorageAllocation)
	}

	return allocs
}

// calcStateTransition calculates the StateTransition which should be made given
// the current state of the garage layout.
//
// knownNodes is the set of KnownNode values returned from a
// [garage.ClusterStatus], filtered to only include those pertinent to this
// host.
//
// It is assumed that the active/targetAllocs pair has already been
// validated using validateTargetAllocs.
func calcStateTransition(
	activeAllocs []daecommon.ConfigStorageAllocation,
	knownNodes []garage.KnownNode,
	targetAllocs []daecommon.ConfigStorageAllocation,
) (
	res StateTransition,
) {
	var (
		activeAllocsWithKnownNodes = allocsWithKnownNode(
			activeAllocs, knownNodes,
		)
		targetAllocsM = allocsByRPCPort(targetAllocs)
	)

	// Deal with additions/modifications to the cluster. These are easy, since
	// all target allocations belong in the cluster no matter what.
	for _, targetAlloc := range targetAllocs {
		res.AddModifyAllocations = append(res.AddModifyAllocations, targetAlloc)
	}

	// Deal with removals from the cluster. These are harder.
	for _, activeAlloc := range activeAllocsWithKnownNodes {

		// If there's a corresponding targetAlloc then this alloc was added to
		// AddModifyAllocations already. Even if it was previously draining,
		// once it's added/modified it won't be anymore.
		if _, ok := targetAllocsM[activeAlloc.RPCPort]; ok {
			continue
		}

		// If it's draining then let it continue draining.
		//
		// If there's an associated role then this alloc is still part of the
		// cluster, but it's not in the target set so it needs to be drained
		// and removed.
		if activeAlloc.Draining || activeAlloc.Role != nil {
			res.DrainAllocations = append(res.DrainAllocations, activeAlloc)
			continue
		}

		// The alloc is not in the target alloc set, it's not draining, and it
		// has no role in the cluster.
	}

	return
}