// SPDX-License-Identifier: Apache-2.0
//
// Copyright The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//go:build windows

package pdh

import (
	"fmt"
	"log/slog"
	"math"
	"reflect"
	"runtime"
	"slices"
	"strings"
	"testing"
	"unsafe"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/stretchr/testify/require"
	"golang.org/x/sys/windows"
)

type rowSetValues struct {
	Name     string
	A        float64 `perfdata:"A"`
	B        float64 `perfdata:"B"`
	BSecond  float64 `perfdata:"B,secondvalue"`
	Optional float64 `perfdata:"Optional"`
	// Occurrence is last, so that the field indices of the counters stay the same.
	Occurrence int
}

// newRowSetTestCollector returns a collector with the counters A and B, which
// have instances, and Optional, which was skipped at Build and has none.
func newRowSetTestCollector(partialRows bool) *Collector[rowSetValues] {
	var row rowSetValues

	return &Collector[rowSetValues]{
		partialRows: partialRows,
		logger:      slog.New(slog.DiscardHandler),
		rows:        structRowAccessor[rowSetValues](int(unsafe.Offsetof(row.Name)), int(unsafe.Offsetof(row.Occurrence))),
		counters: []Counter{
			{Name: "A", Instances: map[string]pdhCounterHandle{"*": 1}, FieldIndexValue: int(unsafe.Offsetof(row.A)), FieldIndexSecondValue: -1},
			{Name: "B", Instances: map[string]pdhCounterHandle{"*": 2}, FieldIndexValue: int(unsafe.Offsetof(row.B)), FieldIndexSecondValue: int(unsafe.Offsetof(row.BSecond))},
			{Name: "Optional", Instances: map[string]pdhCounterHandle{}, FieldIndexValue: int(unsafe.Offsetof(row.Optional)), FieldIndexSecondValue: -1},
		},
		counterIndex: map[string]int{"A": 0, "B": 1, "Optional": 2},
	}
}

// fillRows adds the instance "complete" with valid values for A and B, and the
// instance "partial" with a valid value for A only, as if its item for B had an
// invalid status.
func fillRows(t *testing.T, c *Collector[rowSetValues]) []rowSetValues {
	t.Helper()

	var dst []rowSetValues

	rows := newRowSet(c, &dst, &collectState{})

	for _, name := range []string{"complete", "partial"} {
		row, ok := rows.row(&c.counters[0], windows.StringToUTF16Ptr(name), CstatusValidData)
		require.True(t, ok)
		require.True(t, c.setRawValue(&dst[row], &c.counters[0], RawCounter{FirstValue: 1}))
		rows.setValid(row, 0)
	}

	clear(rows.seen)

	_, ok := rows.row(&c.counters[1], windows.StringToUTF16Ptr("partial"), CstatusInvalidData)
	require.False(t, ok)

	row, ok := rows.row(&c.counters[1], windows.StringToUTF16Ptr("complete"), CstatusNewData)
	require.True(t, ok)
	require.True(t, c.setRawValue(&dst[row], &c.counters[1], RawCounter{FirstValue: 2, SecondValue: 3}))
	rows.setValid(row, 1)

	rows.finish()

	return dst
}

func TestRowSetOmitsIncompleteInstances(t *testing.T) {
	t.Parallel()

	rows := fillRows(t, newRowSetTestCollector(false))

	// "partial" has no valid value for B. Without partial rows, it must be left
	// out instead of reporting B as zero. Optional has no instances, so it is
	// not required.
	require.Equal(t, []rowSetValues{{Name: "complete", A: 1, B: 2, BSecond: 3}}, rows)
}

func TestRowSetPartialRows(t *testing.T) {
	t.Parallel()

	rows := fillRows(t, newRowSetTestCollector(true))
	require.Len(t, rows, 2)

	require.Equal(t, rowSetValues{Name: "complete", A: 1, B: 2, BSecond: 3}, rows[0])

	require.Equal(t, "partial", rows[1].Name)
	require.InDelta(t, 1, rows[1].A, 0)
	require.True(t, math.IsNaN(rows[1].B), "B: %v", rows[1].B)
	require.True(t, math.IsNaN(rows[1].BSecond), "BSecond: %v", rows[1].BSecond)
	require.Zero(t, rows[1].Optional)
}

func TestRowSetDuplicateInstances(t *testing.T) {
	t.Parallel()

	nan := math.NaN()

	// wantRow identifies a row by name and occurrence. suffixed is the name
	// that SuffixDuplicates gives the row.
	type wantRow struct {
		name       string
		occurrence int
		suffixed   string
		a, b       float64
	}

	for _, tc := range []struct {
		name        string
		names       []string
		secondNames []string
		a           []float64
		b           []float64
		repeatArray bool
		want        []wantRow
	}{
		{
			name:  "duplicates",
			names: []string{"A", "A", "A"},
			a:     []float64{1, 2, 3},
			b:     []float64{10, 20, 30},
			want: []wantRow{
				{name: "A", suffixed: "A", a: 1, b: 10},
				{name: "A", occurrence: 1, suffixed: "A#1", a: 2, b: 20},
				{name: "A", occurrence: 2, suffixed: "A#2", a: 3, b: 30},
			},
		},
		{
			name:        "interleaved_names",
			names:       []string{"A", "B", "A"},
			secondNames: []string{"B", "A", "A"},
			a:           []float64{1, 2, 3},
			b:           []float64{20, 10, 30},
			want: []wantRow{
				{name: "A", suffixed: "A", a: 1, b: 10},
				{name: "B", suffixed: "B", a: 2, b: 20},
				{name: "A", occurrence: 1, suffixed: "A#1", a: 3, b: 30},
			},
		},
		{
			// A name with a suffix reported by Windows is a different instance.
			name:        "literal_suffix",
			names:       []string{"A", "A", "A#1"},
			secondNames: []string{"A#1", "A", "A"},
			a:           []float64{1, 2, 3},
			b:           []float64{30, 10, 20},
			want: []wantRow{
				{name: "A", suffixed: "A", a: 1, b: 10},
				{name: "A", occurrence: 1, suffixed: "A#2", a: 2, b: 20},
				{name: "A#1", suffixed: "A#1", a: 3, b: 30},
			},
		},
		{
			name:  "invalid_middle_item",
			names: []string{"A", "A", "A"},
			a:     []float64{1, nan, 3},
			b:     []float64{10, 20, 30},
			want: []wantRow{
				{name: "A", suffixed: "A", a: 1, b: 10},
				{name: "A", occurrence: 1, suffixed: "A#1", a: nan, b: 20},
				{name: "A", occurrence: 2, suffixed: "A#2", a: 3, b: 30},
			},
		},
		{
			name:  "invalid_first_item",
			names: []string{"A", "A", "A"},
			a:     []float64{1, 2, 3},
			b:     []float64{nan, 20, 30},
			want: []wantRow{
				{name: "A", suffixed: "A", a: 1, b: nan},
				{name: "A", occurrence: 1, suffixed: "A#1", a: 2, b: 20},
				{name: "A", occurrence: 2, suffixed: "A#2", a: 3, b: 30},
			},
		},
		{
			// The invalid "A#1" has no row, so its suffix is free.
			name:  "invalid_literal_suffix",
			names: []string{"A", "A", "A#1"},
			a:     []float64{1, 2, nan},
			b:     []float64{10, 20, nan},
			want: []wantRow{
				{name: "A", suffixed: "A", a: 1, b: 10},
				{name: "A", occurrence: 1, suffixed: "A#1", a: 2, b: 20},
			},
		},
		{
			name:        "overlapping_counter_arrays",
			names:       []string{"A", "A"},
			a:           []float64{1, 2},
			b:           []float64{10, 20},
			repeatArray: true,
			want: []wantRow{
				{name: "A", suffixed: "A", a: 1, b: 10},
				{name: "A", occurrence: 1, suffixed: "A#1", a: 2, b: 20},
			},
		},
	} {
		for _, resultType := range []CounterType{CounterTypeRaw, CounterTypeFormatted} {
			for _, partialRows := range []bool{false, true} {
				for _, keepsOccurrences := range []bool{true, false} {
					t.Run(fmt.Sprintf("%s/%s/partial=%t/occurrences=%t", tc.name, resultType, partialRows, keepsOccurrences), func(t *testing.T) {
						t.Parallel()

						c := newRowSetTestCollector(partialRows)
						c.counters[1].FieldIndexSecondValue = -1
						c.rows.keepsOccurrences = keepsOccurrences

						var dst []rowSetValues

						rows := newRowSet(c, &dst, &collectState{})

						addDuplicateTestItems(&rows, resultType, 0, tc.names, tc.a)

						if tc.repeatArray {
							addDuplicateTestItems(&rows, resultType, 0, tc.names, tc.a)
						}

						secondNames := tc.secondNames
						if secondNames == nil {
							secondNames = tc.names
						}

						addDuplicateTestItems(&rows, resultType, 1, secondNames, tc.b)
						rows.finish()

						want := slices.DeleteFunc(slices.Clone(tc.want), func(row wantRow) bool {
							return !partialRows && (math.IsNaN(row.a) || math.IsNaN(row.b))
						})

						if !keepsOccurrences {
							// Without an Occurrence field, only the first remaining instance of a name is collected.
							seen := map[string]bool{}
							want = slices.DeleteFunc(want, func(row wantRow) bool {
								duplicate := seen[row.name]
								seen[row.name] = true

								return duplicate
							})
						}

						require.Len(t, dst, len(want))

						for _, expected := range want {
							i := slices.IndexFunc(dst, func(row rowSetValues) bool {
								return row.Name == expected.name && row.Occurrence == expected.occurrence
							})
							require.NotEqual(t, -1, i, "missing instance %s occurrence %d", expected.name, expected.occurrence)

							for field, value := range map[string]float64{"A": expected.a, "B": expected.b} {
								got := reflect.ValueOf(dst[i]).FieldByName(field).Float()
								if math.IsNaN(value) {
									require.True(t, math.IsNaN(got), "%s/%s: %v", expected.name, field, got)

									continue
								}

								require.InDelta(t, value, got, 0, "%s/%s", expected.name, field)
							}
						}

						if !keepsOccurrences {
							return
						}

						// SuffixDuplicates gives every row a unique name.
						SuffixDuplicates(dst,
							func(row *rowSetValues) *string { return &row.Name },
							func(row *rowSetValues) int { return row.Occurrence },
						)

						for _, expected := range want {
							i := slices.IndexFunc(dst, func(row rowSetValues) bool { return row.Name == expected.suffixed })
							require.NotEqual(t, -1, i, "missing suffixed instance %s", expected.suffixed)
							require.Equal(t, expected.occurrence, dst[i].Occurrence, expected.suffixed)
						}
					})
				}
			}
		}
	}
}

// addDuplicateTestItems supplies counter arrays without querying PDH. NaN marks an invalid item.
func addDuplicateTestItems(
	rows *rowSet[rowSetValues],
	resultType CounterType,
	counterIndex int,
	names []string,
	values []float64,
) {
	if resultType == CounterTypeRaw {
		items := make([]RawCounterItem, len(names))
		for i, name := range names {
			status := CstatusValidData
			if math.IsNaN(values[i]) {
				status = CstatusInvalidData
			}

			items[i] = RawCounterItem{
				SzName: windows.StringToUTF16Ptr(name),
				RawValue: RawCounter{
					CStatus:    status,
					FirstValue: int64(values[i]),
				},
			}
		}

		itemSize := int(unsafe.Sizeof(RawCounterItem{}))
		buf := unsafe.Slice((*byte)(unsafe.Pointer(unsafe.SliceData(items))), len(items)*itemSize)
		rows.addRawItems(counterIndex, buf, uint32(len(items)))
		runtime.KeepAlive(items)

		return
	}

	items := make([]FmtCounterValueItemDouble, len(names))
	for i, name := range names {
		status := CstatusNewData
		if math.IsNaN(values[i]) {
			status = CstatusInvalidData
		}

		items[i] = FmtCounterValueItemDouble{
			SzName: windows.StringToUTF16Ptr(name),
			FmtValue: FmtCounterValueDouble{
				CStatus:     status,
				DoubleValue: values[i],
			},
		}
	}

	itemSize := int(unsafe.Sizeof(FmtCounterValueItemDouble{}))
	buf := unsafe.Slice((*byte)(unsafe.Pointer(unsafe.SliceData(items))), len(items)*itemSize)
	rows.addFormattedItems(counterIndex, buf, uint32(len(items)))
	runtime.KeepAlive(items)
}

func TestSetRawValueElapsedTime(t *testing.T) {
	t.Parallel()

	c := newRowSetTestCollector(false)

	var values rowSetValues

	counter := &Counter{Type: PERF_ELAPSED_TIME, Frequency: 10_000_000, FieldIndexValue: int(unsafe.Offsetof(values.A)), FieldIndexSecondValue: -1}

	// 2.5 s at 10 MHz. Integer division truncated this to 2.
	require.True(t, c.setRawValue(&values, counter, RawCounter{FirstValue: 100_000_000, SecondValue: 125_000_000}))
	require.InDelta(t, 2.5, values.A, 0)

	// A zero frequency panicked with an integer divide by zero.
	counter.Frequency = 0
	require.False(t, c.setRawValue(&values, counter, RawCounter{FirstValue: 1, SecondValue: 2}))

	// A counter with only a second value field must not set field -1.
	counter = &Counter{Type: PERF_ELAPSED_TIME, Frequency: 1, FieldIndexValue: -1, FieldIndexSecondValue: int(unsafe.Offsetof(values.BSecond))}
	require.True(t, c.setRawValue(&values, counter, RawCounter{FirstValue: 1, SecondValue: 2}))
}

// TestCollectLoadsMissingTimeBase checks that an elapsed time counter whose time
// base was not available when the collector was created gets it on a later
// sample, instead of omitting every instance until a restart (#2313).
func TestCollectLoadsMissingTimeBase(t *testing.T) {
	t.Parallel()

	type systemUpTime struct {
		SystemUpTime float64 `perfdata:"System Up Time"`
	}

	collector, err := NewCollector[systemUpTime](slog.New(slog.DiscardHandler), CounterTypeRaw, "System", nil)
	require.NoError(t, err)

	t.Cleanup(collector.Close)

	require.Len(t, collector.counters, 1)

	counter := &collector.counters[0]
	require.Equal(t, uint32(PERF_ELAPSED_TIME), counter.Type)

	frequency := counter.Frequency
	require.Positive(t, frequency)

	// As if GetCounterTimeBase reported no data at startup.
	counter.Frequency = 0

	var dst []systemUpTime

	require.NoError(t, collector.Collect(&dst))
	require.Equal(t, frequency, counter.Frequency)
	require.Len(t, dst, 1)
	require.Positive(t, dst[0].SystemUpTime)
}

// TestFrequencyOfAverageTimer checks that the time base of a raw average timer
// is available, also when it was missing at startup, and that other counters
// have none.
func TestFrequencyOfAverageTimer(t *testing.T) {
	t.Parallel()

	type physicalDisk struct {
		AvgDiskSecPerRead      float64 `perfdata:"Avg. Disk sec/Read"`
		AvgDiskSecPerReadBase  float64 `perfdata:"Avg. Disk sec/Read,secondvalue"`
		DiskReadsPerSec        float64 `perfdata:"Disk Reads/sec"`
		CurrentDiskQueueLength float64 `perfdata:"Current Disk Queue Length"`
	}

	collector, err := NewCollector[physicalDisk](slog.New(slog.DiscardHandler), CounterTypeRaw, "PhysicalDisk", []string{"_Total"})
	require.NoError(t, err)

	t.Cleanup(collector.Close)

	frequency, ok := collector.Frequency("Avg. Disk sec/Read")
	require.True(t, ok)
	require.Positive(t, frequency)

	for _, name := range []string{"Disk Reads/sec", "Current Disk Queue Length", "unknown"} {
		_, ok = collector.Frequency(name)
		require.False(t, ok, name)
	}

	// As if GetCounterTimeBase failed at startup.
	idx := slices.IndexFunc(collector.counters, func(counter Counter) bool { return counter.Name == "Avg. Disk sec/Read" })
	require.NotEqual(t, -1, idx)

	collector.counters[idx].Frequency = 0

	_, ok = collector.Frequency("Avg. Disk sec/Read")
	require.False(t, ok)

	var dst []physicalDisk

	require.NoError(t, collector.Collect(&dst))

	reloaded, ok := collector.Frequency("Avg. Disk sec/Read")
	require.True(t, ok)
	require.Equal(t, frequency, reloaded)

	var nilCollector *Collector[physicalDisk]

	_, ok = nilCollector.Frequency("Avg. Disk sec/Read")
	require.False(t, ok)
}

func TestMetricTypeForCounterType(t *testing.T) {
	t.Parallel()

	require.Equal(t, prometheus.GaugeValue, metricType(CounterTypeRaw, PERF_COUNTER_RAWCOUNT))
	require.Equal(t, prometheus.CounterValue, metricType(CounterTypeRaw, PERF_COUNTER_BULK_COUNT))
	require.Equal(t, prometheus.GaugeValue, metricType(CounterTypeRaw, 0xFFFFFFFF))
	require.Equal(t, prometheus.GaugeValue, metricType(CounterTypeFormatted, PERF_COUNTER_BULK_COUNT))
}

type processThreads struct {
	Name        string
	ThreadCount float64 `perfdata:"Thread Count"`
}

// TestCollectRecoversPanic checks that a panic while collecting is returned as an error.
func TestCollectRecoversPanic(t *testing.T) {
	t.Parallel()

	c, err := NewCollector[processThreads](slog.New(slog.DiscardHandler), CounterTypeRaw, "Process", InstancesAll)
	require.NoError(t, err)

	t.Cleanup(c.Close)

	var dst []processThreads

	require.NoError(t, c.Collect(&dst))

	setValue := setPanicOnSetValue(c)

	err = c.Collect(&dst)
	require.ErrorContains(t, err, "panic while collecting performance counters of Process")

	// Collection still works after the panic.
	c.mu.Lock()
	c.rows.setValue = setValue
	c.mu.Unlock()

	require.NoError(t, c.Collect(&dst))
	require.NotEmpty(t, dst)
}

// TestCollectKeepsInstanceNames checks that Collect reuses decoded names between samples and
// removes the cached names that a failed sample did not reach.
func TestCollectKeepsInstanceNames(t *testing.T) {
	t.Parallel()

	c, err := NewCollector[processThreads](slog.New(slog.DiscardHandler), CounterTypeRaw, "Process", InstancesAll)
	require.NoError(t, err)

	t.Cleanup(c.Close)

	systemName := func(rows []processThreads) string {
		i := slices.IndexFunc(rows, func(row processThreads) bool { return row.Name == "System" })
		require.NotEqual(t, -1, i, "missing instance System")

		return rows[i].Name
	}

	var first, second []processThreads

	require.NoError(t, c.Collect(&first))
	require.NoError(t, c.Collect(&second))
	require.Same(t, unsafe.StringData(systemName(first)), unsafe.StringData(systemName(second)), "name decoded again")

	// Assertions run without holding the lock, because a failed assertion would leave it locked for Close.
	setValue := setPanicOnSetValue(c)

	cachedNames := func() int {
		c.mu.Lock()
		defer c.mu.Unlock()

		return len(c.state.names.names)
	}

	// Let the first item panic.
	require.Greater(t, cachedNames(), 1)
	require.Error(t, c.Collect(&first))
	require.Equal(t, 1, cachedNames(), "names not reached by the failed sample kept")

	c.mu.Lock()
	c.rows.setValue = setValue
	c.mu.Unlock()

	require.NoError(t, c.Collect(&first))
	require.Equal(t, systemName(second), systemName(first))
}

func TestInstanceNameCache(t *testing.T) {
	t.Parallel()

	var cache instanceNameCache

	// PDH writes the names of every counter array into the same buffer.
	buf := make([]uint16, 8)
	write := func(name string) *uint16 {
		copy(buf, windows.StringToUTF16(name))

		return &buf[0]
	}

	cache.startSample()

	first := cache.decode(write("a"))
	require.Equal(t, "a", first)
	require.Equal(t, "bb", cache.decode(write("bb")))
	require.Empty(t, cache.decode(nil))

	cache.removeUnseen()
	require.Len(t, cache.names, 2)

	cache.startSample()

	second := cache.decode(write("a"))
	require.Equal(t, "a", second)
	require.Same(t, unsafe.StringData(first), unsafe.StringData(second), "name decoded again")

	cache.removeUnseen()
	require.Len(t, cache.names, 1, "name of a vanished instance kept")
	require.Contains(t, cache.names, "a\x00", "keyed by the UTF-16LE content")
}

// TestRowSetReusesState checks that a sample does not see the instances of an
// earlier sample that shared its state, even if that sample did not finish.
// TestRowSetWarnsOnceAboutDroppedDuplicates checks that dropping duplicate
// instances of a row type without an Occurrence field is logged once.
func TestRowSetWarnsOnceAboutDroppedDuplicates(t *testing.T) {
	t.Parallel()

	var logs strings.Builder

	c := newRowSetTestCollector(true)
	c.logger = slog.New(slog.NewTextHandler(&logs, &slog.HandlerOptions{Level: slog.LevelWarn}))
	c.rows.keepsOccurrences = false
	state := &collectState{}

	for range 3 {
		var dst []rowSetValues

		rows := newRowSet(c, &dst, state)
		addDuplicateTestItems(&rows, CounterTypeRaw, 0, []string{"a", "a", "b"}, []float64{1, 2, 3})
		addDuplicateTestItems(&rows, CounterTypeRaw, 1, []string{"a", "a", "b"}, []float64{4, 5, 6})
		rows.finish()

		state.keys = rows.keys

		require.Len(t, dst, 2)
		require.InDelta(t, 1, dst[0].A, 0, "the first instance of a is kept")
	}

	require.Equal(t, 1, strings.Count(logs.String(), "level=WARN"), logs.String())
}

func TestRowSetReusesState(t *testing.T) {
	t.Parallel()

	c := newRowSetTestCollector(false)
	c.counters[1].FieldIndexSecondValue = -1
	state := &collectState{}

	var dst []rowSetValues

	// A sample that fails after its first counter array.
	rows := newRowSet(c, &dst, state)
	addDuplicateTestItems(&rows, CounterTypeRaw, 0, []string{"a", "a", "b"}, []float64{1, 2, 3})

	for range 2 {
		dst = nil
		rows = newRowSet(c, &dst, state)
		addDuplicateTestItems(&rows, CounterTypeRaw, 0, []string{"b", "c"}, []float64{4, 5})
		addDuplicateTestItems(&rows, CounterTypeRaw, 1, []string{"b", "c"}, []float64{6, 7})
		rows.finish()

		state.valid = rows.valid

		require.Equal(t, []rowSetValues{
			{Name: "b", A: 4, B: 6},
			{Name: "c", A: 5, B: 7},
		}, dst)
	}
}

// setPanicOnSetValue lets every value that c sets panic and returns the
// previous setter.
func setPanicOnSetValue[T any](c *Collector[T]) func(row *T, field int, value float64) {
	c.mu.Lock()
	defer c.mu.Unlock()

	setValue := c.rows.setValue
	c.rows.setValue = func(*T, int, float64) { panic("test panic") }

	return setValue
}