package cpu

import "golang.org/x/sys/cpu"

// Features captures the SIMD capabilities we care about for dispatching
// CPU kernels. Values are populated once at init using the runtime cpuid
// information from x/sys/cpu.
type Features struct {
	FMA      bool
	AVX2     bool
	AVX512F  bool
	AVX512DQ bool
	AVX512BW bool
	AVX512VL bool

	AVX512VNNI bool

	AMXTile bool
	AMXInt8 bool
	AMXBF16 bool
}

var detected Features

func init() {
	detected = Features{
		FMA:      cpu.X86.HasFMA,
		AVX2:     cpu.X86.HasAVX2,
		AVX512F:  cpu.X86.HasAVX512F,
		AVX512DQ: cpu.X86.HasAVX512DQ,
		AVX512BW: cpu.X86.HasAVX512BW,
		AVX512VL: cpu.X86.HasAVX512VL,

		AVX512VNNI: cpu.X86.HasAVX512VNNI,

		AMXTile: cpu.X86.HasAMXTile,
		AMXInt8: cpu.X86.HasAMXInt8,
		AMXBF16: cpu.X86.HasAMXBF16,
	}
}

// CPUFeatures returns the detected SIMD feature set.
func CPUFeatures() Features {
	return detected
}

// SupportsAVX512 reports whether the CPU supports the AVX-512 subset we use.
// We require F, DQ, BW and VL to match ggml-style kernels.
func SupportsAVX512() bool {
	return detected.FMA && detected.AVX512F && detected.AVX512DQ && detected.AVX512BW && detected.AVX512VL
}

// SupportsAVX2 reports whether AVX2 is available.
func SupportsAVX2() bool {
	return detected.FMA && detected.AVX2
}

// SupportsAVX512VNNI reports whether the CPU supports AVX-512 VNNI instructions.
func SupportsAVX512VNNI() bool {
	return SupportsAVX512() && detected.AVX512VNNI
}

// SupportsAMXInt8 reports whether the CPU supports AMX INT8 tile ops.
func SupportsAMXInt8() bool {
	return detected.AMXTile && detected.AMXInt8
}

// SupportsAMXBF16 reports whether the CPU supports AMX BF16 tile ops.
func SupportsAMXBF16() bool {
	return detected.AMXTile && detected.AMXBF16
}

// SIMDLevel represents the best available SIMD tier.
type SIMDLevel int

const (
	SIMDNone SIMDLevel = iota
	SIMDAVX2
	SIMDAVX512
)

// BestSIMD returns the highest SIMD level supported (AVX-512 over AVX2).
func BestSIMD() SIMDLevel {
	switch {
	case SupportsAVX512():
		return SIMDAVX512
	case SupportsAVX2():
		return SIMDAVX2
	default:
		return SIMDNone
	}
}