// Copyright IBM Corp. 2015, 2025
// SPDX-License-Identifier: BUSL-1.1

package structs

import (
	"reflect"

	"github.com/hashicorp/nomad/client/lib/numalib"
	"github.com/hashicorp/nomad/client/lib/numalib/hw"
	"github.com/hashicorp/nomad/helper"
)

var (
	// extendedTypes is a mapping of extended types to their extension function
	// TODO: the duplicates could be simplified by looking up the base type in the case of a pointer type in ConvertExt
	extendedTypes = map[reflect.Type]extendFunc{
		reflect.TypeOf(Node{}):              nodeExt,
		reflect.TypeOf(&Node{}):             nodeExt,
		reflect.TypeOf(CSIVolume{}):         csiVolumeExt,
		reflect.TypeOf(&CSIVolume{}):        csiVolumeExt,
		reflect.TypeOf(&numalib.Topology{}): numaTopoExt,
	}
)

// numaTopoExt is used to JSON encode topology to correctly handle the private
// idset.Set fields and so that NUMA NodeIDs are encoded as []int because
// go-msgpack will further JSON encode []uint8 into a base64-encoded bytestring,
// rather than an array
func numaTopoExt(v interface{}) interface{} {
	topo := v.(*numalib.Topology)

	var nodes []int
	if topo.GetNodes() != nil {
		nodes = helper.ConvertSlice(
			topo.GetNodes().Slice(), func(n uint8) int { return int(n) })
	}

	return &struct {
		Nodes                  []int
		Distances              numalib.SLIT
		Cores                  []numalib.Core
		OverrideTotalCompute   hw.MHz
		OverrideWitholdCompute hw.MHz
	}{
		Nodes:                  nodes,
		Distances:              topo.Distances,
		Cores:                  topo.Cores,
		OverrideTotalCompute:   topo.OverrideTotalCompute,
		OverrideWitholdCompute: topo.OverrideWitholdCompute,
	}
}

// nodeExt ensures the node is sanitized and adds the legacy field .Drain back to encoded Node objects
func nodeExt(v interface{}) interface{} {
	node := v.(*Node).Sanitize()
	// transform to a struct with inlined Node fields plus the Drain field
	// - using defined type (not an alias!) EmbeddedNode gives us free conversion to a distinct type
	// - distinct type prevents this encoding extension from being called recursively/infinitely on the embedding
	// - pointers mean the conversion function doesn't have to make a copy during conversion
	type EmbeddedNode Node
	return &struct {
		*EmbeddedNode
		Drain bool
	}{
		EmbeddedNode: (*EmbeddedNode)(node),
		Drain:        node != nil && node.DrainStrategy != nil,
	}
}

func csiVolumeExt(v interface{}) interface{} {
	vol := v.(*CSIVolume).Sanitize()
	type EmbeddedCSIVolume CSIVolume

	allocCount := len(vol.ReadAllocs) + len(vol.WriteAllocs)

	apiVol := &struct {
		*EmbeddedCSIVolume
		Allocations []*AllocListStub
	}{
		EmbeddedCSIVolume: (*EmbeddedCSIVolume)(vol),
		Allocations:       make([]*AllocListStub, 0, allocCount),
	}

	// WriteAllocs and ReadAllocs will only ever contain the Allocation ID,
	// with a null value for the Allocation; these IDs are mapped to
	// allocation stubs in the Allocations field. This indirection is so the
	// API can support both the UI and CLI consumer in a safely backwards
	// compatible way
	for _, a := range vol.ReadAllocs {
		if a != nil {
			apiVol.ReadAllocs[a.ID] = nil
			apiVol.Allocations = append(apiVol.Allocations, a.Stub(nil))
		}
	}
	for _, a := range vol.WriteAllocs {
		if a != nil {
			apiVol.WriteAllocs[a.ID] = nil
			apiVol.Allocations = append(apiVol.Allocations, a.Stub(nil))
		}
	}

	return apiVol
}
