// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package io

import (
	"encoding/binary"
	"fmt"
	"math"
	"runtime"
	"strings"
	"testing"
	"time"
)

// Regression pins for the HDF5 group walk: deep chains that must stay
// linear, the depth cap, self-links and hard-link diamonds that must
// be refused as cycles.

// groupChain builds a well-formed HDF5 file holding a chain of
// depth groups, each carrying one named attribute and a hard link to the
// next. Every group is legal, the attributes are legal scalar attributes
// and the chain terminates, so the reader has no grounds to refuse it:
// the walk has to stay cheap on its own.
func groupChain(depth int) []byte {
	const (
		firstGroup = 128
		stride     = 88
	)
	n := firstGroup + depth*stride + 16
	f := make([]byte, n)
	copy(f, hdf5Magic)
	f[8] = 0 // superblock version 0
	f[13] = 8
	f[14] = 8
	binary.LittleEndian.PutUint64(f[32:], math.MaxUint64) // free space undefined
	binary.LittleEndian.PutUint64(f[40:], uint64(n))      // end of file
	binary.LittleEndian.PutUint64(f[48:], math.MaxUint64) // driver undefined
	binary.LittleEndian.PutUint64(f[64:], firstGroup)     // root object header

	for k := range depth {
		a := firstGroup + k*stride
		nmsg := 2
		if k == depth-1 {
			nmsg = 1 // the tail group only carries its attribute
		}
		f[a] = 1 // object header version 1
		binary.LittleEndian.PutUint16(f[a+2:], uint16(nmsg))
		binary.LittleEndian.PutUint32(f[a+4:], 1)  // reference count
		binary.LittleEndian.PutUint32(f[a+8:], 88) // message data size

		// Attribute message (type 12), size 40, body at a+24.
		binary.LittleEndian.PutUint16(f[a+16:], hdf5MsgAttribute)
		binary.LittleEndian.PutUint16(f[a+18:], 40)
		f[a+24] = 1                                // version
		binary.LittleEndian.PutUint16(f[a+26:], 8) // name length
		binary.LittleEndian.PutUint16(f[a+28:], 8) // datatype length
		binary.LittleEndian.PutUint16(f[a+30:], 8) // dataspace length
		copy(f[a+32:], fmt.Sprintf("a%07d", k))    // 8-byte attribute name
		f[a+40] = 0x10                             // fixed-point datatype
		binary.LittleEndian.PutUint32(f[a+44:], 1) // one byte
		f[a+48] = 1                                // scalar dataspace
		f[a+56] = byte(k)                          // one byte of value

		if k == depth-1 {
			continue
		}
		// Link message (type 6), size 12, body at a+72.
		binary.LittleEndian.PutUint16(f[a+64:], hdf5MsgLink)
		binary.LittleEndian.PutUint16(f[a+66:], 12)
		f[a+72] = 1 // version
		f[a+73] = 0 // flags: hard link, 1-byte name length
		f[a+74] = 1 // name length
		f[a+75] = 'a'
		binary.LittleEndian.PutUint64(f[a+76:], uint64(a+stride))
	}
	return f
}

// selfLink builds the smallest HDF5 file whose root group links to
// itself: 136 bytes, one object header, one link message.
func selfLink() []byte {
	const N = 136
	f := make([]byte, N)
	copy(f, hdf5Magic)
	f[8] = 0 // superblock version 0
	f[13] = 8
	f[14] = 8
	binary.LittleEndian.PutUint64(f[32:], math.MaxUint64)
	binary.LittleEndian.PutUint64(f[40:], N)
	binary.LittleEndian.PutUint64(f[48:], math.MaxUint64)
	binary.LittleEndian.PutUint64(f[64:], 96) // root object header

	f[96] = 1 // object header version 1
	binary.LittleEndian.PutUint16(f[98:], 1)
	binary.LittleEndian.PutUint32(f[100:], 1)
	binary.LittleEndian.PutUint32(f[104:], 24)

	binary.LittleEndian.PutUint16(f[112:], hdf5MsgLink)
	binary.LittleEndian.PutUint16(f[114:], 12)
	f[120] = 1 // link message version
	f[121] = 0 // hard link, 1-byte name length
	f[122] = 1 // name length
	f[123] = 'a'
	binary.LittleEndian.PutUint64(f[124:], 96) // the link target: itself
	return f
}

// TestWalkRefusesHardLinkCycle pins the crash that took the editor down:
// walk had no visited set, so a group linked into itself recursed for
// ever while the path string grew, and the heap grew with it at hundreds
// of megabytes per second until the host ran out of memory. The reader
// must refuse the file with an error.
func TestWalkRefusesHardLinkCycle(t *testing.T) {
	guard := time.AfterFunc(20*time.Second, func() { panic("LoadHDF5 on a cyclic file did not return") })
	defer guard.Stop()

	path := writeHostile(t, "selflink.h5", selfLink())
	stop := warnHeap(t, 256<<20)
	_, err := LoadHDF5(path)
	stop()
	if err == nil {
		t.Fatal("LoadHDF5 accepted a group that hard-links into itself")
	}
	if !strings.Contains(err.Error(), "hard-link cycle") {
		t.Fatalf("LoadHDF5 = %v, want the cycle refusal", err)
	}
}

// TestWalkDeepChainStaysLinear pins the other half of the same crash: a
// well-formed chain of groups used to cost memory quadratic in its
// depth, because every level copied the inherited attribute map and
// built a longer path string. The live memory must grow with the depth,
// not with its square, and the walk carries one shared path buffer and
// one shared attribute map to make that so.
func TestWalkDeepChainStaysLinear(t *testing.T) {
	guard := time.AfterFunc(60*time.Second, func() { panic("deep chain walk did not return") })
	defer guard.Stop()

	measure := func(depth int) (uint64, int) {
		data := groupChain(depth)
		path := writeHostile(t, fmt.Sprintf("chain%d.h5", depth), data)
		runtime.GC()
		var before, after runtime.MemStats
		runtime.ReadMemStats(&before)
		if _, err := LoadHDF5(path); err != nil {
			t.Fatalf("depth %d: LoadHDF5 refused a well-formed file: %v", depth, err)
		}
		runtime.ReadMemStats(&after)
		return after.TotalAlloc - before.TotalAlloc, len(data)
	}
	small, smallFile := measure(150)
	big, bigFile := measure(450)
	// Linear: three times the depth is about three times the bytes. The
	// per-level copies this replaced measured 16.5x for the same step.
	if ratio := float64(big) / float64(small); ratio > 6 {
		t.Errorf("allocation grows with the square of the depth: %.1fx for 3x depth (%d B for %d B, %d B for %d B)",
			ratio, big, bigFile, small, smallFile)
	}
}

// TestWalkRefusesUnboundedNesting pins the depth cap: past it the reader
// refuses the file loudly and cheaply instead of walking every level
// first.
func TestWalkRefusesUnboundedNesting(t *testing.T) {
	guard := time.AfterFunc(60*time.Second, func() { panic("deep chain walk did not return") })
	defer guard.Stop()

	deep := writeHostile(t, "deep.h5", groupChain(2*hdf5MaxGroupDepth))
	stop := warnHeap(t, 256<<20)
	_, err := LoadHDF5(deep)
	stop()
	if err == nil {
		t.Fatalf("LoadHDF5 accepted a chain %d groups deep", 2*hdf5MaxGroupDepth)
	}
	if !strings.Contains(err.Error(), "nest deeper than") {
		t.Fatalf("LoadHDF5 = %v, want the nesting refusal", err)
	}
}

// warnHeap fails the test as soon as the live heap passes limit, so a
// regression cannot consume the machine: the watchdog panics, which
// unwinds the offending walk instead of letting it allocate on.
func warnHeap(t *testing.T, limit uint64) func() {
	t.Helper()
	done := make(chan struct{})
	go func() {
		ticker := time.NewTicker(20 * time.Millisecond)
		defer ticker.Stop()
		for {
			select {
			case <-done:
				return
			case <-ticker.C:
				var m runtime.MemStats
				runtime.ReadMemStats(&m)
				if m.HeapAlloc > limit {
					panic("walk heap above its cap")
				}
			}
		}
	}()
	return func() { close(done) }
}