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

package io

import (
	"encoding/binary"
	"math"
	"slices"
	"strings"
	"testing"
	"time"
)

// Regression pins: continuation links the reader sliced past the
// end of, a local heap header sized by the wrong field, the diamond of
// hard links that multiplied the walk, the read budget that only
// counted value bytes, the B-tree sizes pinned to eight-byte addresses,
// negative FITS column counts, the NAXIS prefix that swallowed user
// keywords, and the version 2 header features no test exercised.

// h5SizedHostileFile returns an n-byte HDF5 file with the signature and
// a version 0 superblock of the given address and length sizes whose
// root object header sits at rootAt. h5HostileFile is the 8/8 case;
// the hostile files here need the others.
func h5SizedHostileFile(n, offSize, lenSize int, rootAt uint64) []byte {
	f := make([]byte, n)
	copy(f, hdf5Magic)
	f[8] = 0 // superblock version 0
	f[13] = byte(offSize)
	f[14] = byte(lenSize)
	// Four addresses of offSize bytes from 24: base, free space, end of
	// file, driver information.
	putAddr := func(at int, v uint64) {
		switch offSize {
		case 4:
			binary.LittleEndian.PutUint32(f[at:], uint32(v))
		case 8:
			binary.LittleEndian.PutUint64(f[at:], v)
		}
	}
	putAddr(24, 0)
	putAddr(28, math.MaxUint64)
	putAddr(32, uint64(n))
	putAddr(36, math.MaxUint64)
	// The root symbol table entry: a link name offset of the length
	// size, then the object header address of the address size.
	entry := 24 + 4*offSize
	putAddr(entry+lenSize, rootAt)
	return f
}

// h5HardLink renders a version 1 hard-link message body: a one-byte
// name length, the name and the object header address.
func h5HardLink(name string, addr uint64) []byte {
	b := append([]byte{1, 0, byte(len(name))}, name...)
	return binary.LittleEndian.AppendUint64(b, addr)
}

// h5V3Superblock returns an n-byte file with a version 3 superblock
// (eight-byte addresses and lengths, the whole block under a lookup3
// checksum) whose root object header sits at rootAt.
func h5V3Superblock(n int, rootAt uint64) []byte {
	f := make([]byte, n)
	copy(f, hdf5Magic)
	f[8] = 3
	f[9] = 8                                              // address size
	f[10] = 8                                             // length size
	f[11] = 0                                             // consistency flags
	binary.LittleEndian.PutUint64(f[12:], 0)              // base address
	binary.LittleEndian.PutUint64(f[20:], math.MaxUint64) // no extension
	binary.LittleEndian.PutUint64(f[28:], uint64(n))      // end of file
	binary.LittleEndian.PutUint64(f[36:], rootAt)
	binary.LittleEndian.PutUint32(f[44:], hdf5Lookup3(f[:44]))
	return f
}

// h5WriteOHDRv2 writes a version 2 object header at at whose first
// message region is chunk, under the given flag byte, computing and
// storing the lookup3 checksum, and returns the offset just past the
// header. Flag-selected prefix fields are written as zeros.
func h5WriteOHDRv2(f []byte, at int, flags byte, chunk []byte) int {
	copy(f[at:], hdf5ObjHdr2)
	f[at+4] = 2
	f[at+5] = flags
	p := at + 6
	if flags&0x20 != 0 {
		p += 16 // access, modification, change and birth times
	}
	if flags&0x10 != 0 {
		p += 4 // max compact and min dense attribute counts
	}
	width := 1 << (flags & 0x03)
	switch width {
	case 1:
		f[p] = byte(len(chunk))
	case 2:
		binary.LittleEndian.PutUint16(f[p:], uint16(len(chunk)))
	case 4:
		binary.LittleEndian.PutUint32(f[p:], uint32(len(chunk)))
	case 8:
		binary.LittleEndian.PutUint64(f[p:], uint64(len(chunk)))
	}
	p += width
	copy(f[p:], chunk)
	p += len(chunk)
	binary.LittleEndian.PutUint32(f[p:], hdf5Lookup3(f[at:p]))
	return p + 4
}

// h5V2Msg renders one version 2 object header message: a type byte, a
// two-byte size and a flag byte, widened by a two-byte creation order
// when the header tracks it, then the body.
func h5V2Msg(typ byte, order uint16, body []byte, ordered bool) []byte {
	head := 4
	if ordered {
		head = 6
	}
	m := make([]byte, head+len(body))
	m[0] = typ
	binary.LittleEndian.PutUint16(m[1:], uint16(len(body)))
	if ordered {
		binary.LittleEndian.PutUint16(m[4:], order)
	}
	copy(m[head:], body)
	return m
}

// h5V2Dataspace renders a version 2 dataspace message body: rank one,
// the given extent in eight bytes.
func h5V2Dataspace(dim uint64) []byte {
	b := make([]byte, 12)
	b[0] = 2 // version
	b[1] = 1 // rank
	binary.LittleEndian.PutUint64(b[4:], dim)
	return b
}

// TestLoadHDF5ShortV1Continuation pins a minimal hostile
// file: a continuation block that ends right after its eight-byte
// message header carries a zero-size continuation message, and the
// reader used to slice the link's offset and length out of bytes past
// the block (a [16:8] panic out of LoadHDF5). It must be a named
// error.
func TestLoadHDF5ShortV1Continuation(t *testing.T) {
	f := h5HostileFile(512)
	h5ObjectHeader(f, 96,
		h5Msg{hdf5MsgContinuation, h5Link(136, 8)},
	)
	// The block: exactly eight bytes, one continuation header, no body.
	binary.LittleEndian.PutUint16(f[136:], hdf5MsgContinuation)
	binary.LittleEndian.PutUint16(f[138:], 0)
	_, err := LoadHDF5(writeHostile(t, "v1short.h5", f))
	if err == nil {
		t.Fatal("LoadHDF5 accepted a continuation message with no link body")
	}
	if !strings.Contains(err.Error(), "shorter than an offset and a length") {
		t.Fatalf("error = %v, want the short-link refusal", err)
	}
}

// TestLoadHDF5ShortV2Continuation pins the version 2 twin: a header
// whose whole message region is one continuation message of size zero
// panicked the same way ([12:4]), because the stream read the link's
// offset and length past the region.
func TestLoadHDF5ShortV2Continuation(t *testing.T) {
	f := h5V3Superblock(512, 96)
	chunk := []byte{hdf5MsgContinuation, 0, 0, 0} // type 16, size 0, flags 0
	h5WriteOHDRv2(f, 96, 0, chunk)
	_, err := LoadHDF5(writeHostile(t, "v2short.h5", f))
	if err == nil {
		t.Fatal("LoadHDF5 accepted a version 2 continuation message with no link body")
	}
	if !strings.Contains(err.Error(), "shorter than an offset and a length") {
		t.Fatalf("error = %v, want the short-link refusal", err)
	}
}

// TestLoadHDF5LocalHeapMixedSizes pins the local heap header size: the
// header holds two length-size fields and then one address-size field,
// but the reader sized it with three address-size fields, so a 4/8 file
// (twenty-byte header) was sliced at [24:] and panicked. A heap whose
// header lies about nothing must still be refused cleanly when what it
// points at is absent.
func TestLoadHDF5LocalHeapMixedSizes(t *testing.T) {
	f := h5SizedHostileFile(512, 4, 8, 96)
	body := make([]byte, 2*4)
	binary.LittleEndian.PutUint32(body, math.MaxUint32) // B-tree: undefined
	binary.LittleEndian.PutUint32(body[4:], 256)        // local heap at 256
	h5ObjectHeader(f, 96, h5Msg{hdf5MsgSymbolTable, body})
	copy(f[256:], hdf5LocalHeap)
	f[260] = 1 // version
	_, err := LoadHDF5(writeHostile(t, "heap48.h5", f))
	if err != nil && strings.Contains(err.Error(), "runtime error") {
		t.Fatalf("the 4/8 local heap panicked: %v", err)
	}
	_ = err // any named refusal is fine; the panic is the defect
}

// TestWalkDiamondReadsOnce pins the diamond: two names on one group
// used to walk the group twice, and a diamond of depth d walked it 2^d
// times, which stopped the reader for hours on a file of a kilobyte.
// The object is read once, under the first path the traversal reaches,
// and a link that closes a cycle along the current path is still an
// error.
func TestWalkDiamondReadsOnce(t *testing.T) {
	guard := time.AfterFunc(20*time.Second, func() { panic("diamond walk did not return") })
	defer guard.Stop()

	const childAt, datasetAt, dataAt = 256, 384, 512
	f := h5HostileFile(576)
	h5ObjectHeader(f, 96,
		h5Msg{hdf5MsgLink, h5HardLink("a", childAt)},
		h5Msg{hdf5MsgLink, h5HardLink("b", childAt)},
	)
	h5ObjectHeader(f, childAt,
		h5Msg{hdf5MsgLink, h5HardLink("d", datasetAt)},
	)
	h5ObjectHeader(f, datasetAt,
		h5Msg{hdf5MsgDataspace, h5Dataspace(1)},
		h5Msg{hdf5MsgDatatype, h5FloatType(8)},
		h5Msg{hdf5MsgDataLayout, h5ContiguousLayout(dataAt, 8)},
	)
	binary.LittleEndian.PutUint64(f[dataAt:], math.Float64bits(2.5))
	sets, err := LoadHDF5(writeHostile(t, "diamond.h5", f))
	if err != nil {
		t.Fatalf("LoadHDF5 on a diamond of hard links: %v", err)
	}
	if len(sets) != 1 {
		t.Fatalf("datasets = %d, want the single dataset once", len(sets))
	}
	// Deterministic: the links of a group are visited in file order, so
	// the first path wins and the listing is stable.
	if sets[0].Path != "/a/d" {
		t.Fatalf("path = %q, want /a/d, the first path to the object", sets[0].Path)
	}
	if got := sets[0].Values.FloatAt(0); got != 2.5 {
		t.Fatalf("value = %v, want 2.5", got)
	}
	t.Run("cycle", func(t *testing.T) {
		// The skip must never swallow a cycle: an object that closes a
		// loop along the current path is refused, not skipped.
		if _, err := LoadHDF5(writeHostile(t, "diamond-cycle.h5", selfLink())); err == nil || !strings.Contains(err.Error(), "hard-link cycle") {
			t.Fatalf("a hard-link cycle: err = %v, want the cycle refusal", err)
		}
	})
}

// TestHDF5BudgetChargesDatasetEnvelope pins the aggregate budget: it
// used to deduct only the value bytes, so millions of empty datasets
// would allocate their result structures, paths and attribute maps
// without ever touching the budget. The walk charges a fixed envelope
// plus the dataset's path beside the values, so a budget below one
// envelope refuses even a file of empty datasets.
func TestHDF5BudgetChargesDatasetEnvelope(t *testing.T) {
	const dataAt, n = 448, 512
	raw := h5HostileFile(n)
	h5ObjectHeader(raw, 96,
		h5Msg{hdf5MsgDataspace, h5Dataspace(1)},
		h5Msg{hdf5MsgDatatype, h5FloatType(8)},
		h5Msg{hdf5MsgDataLayout, h5ContiguousLayout(dataAt, 8)},
	)
	binary.LittleEndian.PutUint64(raw[dataAt:], math.Float64bits(2.5))
	// The file itself is legal: at the default budget it loads.
	sets, err := LoadHDF5(writeHostile(t, "envelope.h5", raw))
	if err != nil || len(sets) != 1 {
		t.Fatalf("LoadHDF5 at the default budget: %d datasets, err = %v", len(sets), err)
	}
	// Below one envelope the same file refuses: the envelope plus the
	// path is charged before any allocation happens.
	f, err := newHDF5File(raw)
	if err != nil {
		t.Fatalf("newHDF5File: %v", err)
	}
	st := newHDF5WalkState()
	st.budget = hdf5DatasetEnvelope / 2
	var out []HDF5Dataset
	if err := f.walk(f.rootAddress, st, &out); err == nil || !strings.Contains(err.Error(), "budget") {
		t.Fatalf("walk with a budget below one envelope: err = %v, want the budget refusal", err)
	}
}

// TestLoadHDF5SymbolTable4of4 pins the four-byte-address layout of the
// group structures: the version 1 B-tree header is 8+2*offSize wide
// (not 24), a symbol node entry is lenSize+offSize+24 (not 40) and the
// local heap header is 8+2*lenSize+offSize. A legal 4/4 file with a
// symbol-table group of two links used to die with "no symbol table
// node at 0"; the 8/8 fixtures are untouched by the arithmetic.
func TestLoadHDF5SymbolTable4of4(t *testing.T) {
	const (
		treeAt = 200
		snodAt = 240
		heapAt = 320
		segAt  = 352
		objA   = 384
		objB   = 512
		dataA  = 640
		dataB  = 648
	)
	f := h5SizedHostileFile(1024, 4, 4, 96)
	body := make([]byte, 8)
	binary.LittleEndian.PutUint32(body, treeAt)
	binary.LittleEndian.PutUint32(body[4:], heapAt)
	h5ObjectHeader(f, 96, h5Msg{hdf5MsgSymbolTable, body})

	// The group B-tree: a sixteen-byte header (signature, type, level,
	// one entry, two sibling addresses), key0 of four bytes, the child
	// address, one trailing key.
	copy(f[treeAt:], hdf5Tree)
	f[treeAt+4] = 0 // group
	f[treeAt+5] = 0 // leaf
	binary.LittleEndian.PutUint16(f[treeAt+6:], 1)
	binary.LittleEndian.PutUint32(f[treeAt+8:], math.MaxUint32)
	binary.LittleEndian.PutUint32(f[treeAt+12:], math.MaxUint32)
	binary.LittleEndian.PutUint32(f[treeAt+16:], 0)      // key0
	binary.LittleEndian.PutUint32(f[treeAt+20:], snodAt) // child
	binary.LittleEndian.PutUint32(f[treeAt+24:], 4)      // trailing key

	// The symbol node: two entries of 32 bytes (heap offset, object
	// address, cache type, reserved, scratch pad).
	copy(f[snodAt:], hdf5SymbolNode)
	f[snodAt+4] = 1
	binary.LittleEndian.PutUint16(f[snodAt+6:], 2)
	binary.LittleEndian.PutUint32(f[snodAt+8:], 0) // "a" at heap offset 0
	binary.LittleEndian.PutUint32(f[snodAt+12:], objA)
	binary.LittleEndian.PutUint32(f[snodAt+40:], 2) // "b" at heap offset 2
	binary.LittleEndian.PutUint32(f[snodAt+44:], objB)

	// The local heap: a twenty-byte header (signature, version, data
	// segment size, free-list head, data segment address).
	copy(f[heapAt:], hdf5LocalHeap)
	f[heapAt+4] = 1
	binary.LittleEndian.PutUint32(f[heapAt+8:], 8)
	binary.LittleEndian.PutUint32(f[heapAt+12:], math.MaxUint32)
	binary.LittleEndian.PutUint32(f[heapAt+16:], segAt)
	copy(f[segAt:], "a\x00b\x00\x00\x00\x00\x00")

	h5ObjectHeader(f, objA,
		h5Msg{hdf5MsgDataspace, h5Dataspace(1)},
		h5Msg{hdf5MsgDatatype, h5FloatType(8)},
		h5Msg{hdf5MsgDataLayout, h5ContiguousLayout(dataA, 8)},
	)
	h5ObjectHeader(f, objB,
		h5Msg{hdf5MsgDataspace, h5Dataspace(1)},
		h5Msg{hdf5MsgDatatype, h5FloatType(8)},
		h5Msg{hdf5MsgDataLayout, h5ContiguousLayout(dataB, 8)},
	)
	binary.LittleEndian.PutUint64(f[dataA:], math.Float64bits(4.5))
	binary.LittleEndian.PutUint64(f[dataB:], math.Float64bits(-2.5))

	sets, err := LoadHDF5(writeHostile(t, "snod44.h5", f))
	if err != nil {
		t.Fatalf("LoadHDF5 refused a legal 4/4 symbol-table file: %v", err)
	}
	if len(sets) != 2 {
		t.Fatalf("datasets = %d, want 2", len(sets))
	}
	if sets[0].Path != "/a" || sets[1].Path != "/b" {
		t.Fatalf("paths = %q, %q, want /a and /b", sets[0].Path, sets[1].Path)
	}
	if got := sets[0].Values.FloatAt(0); got != 4.5 {
		t.Fatalf("/a = %v, want 4.5", got)
	}
	if got := sets[1].Values.FloatAt(0); got != -2.5 {
		t.Fatalf("/b = %v, want -2.5", got)
	}
}

// TestLoadHDF5ContinuationChainDepth pins the version 1 depth guard: a
// chain of continuation blocks one longer than the cap must be refused
// like the version 2 walk refuses its own, not walked to the end.
func TestLoadHDF5ContinuationChainDepth(t *testing.T) {
	const first = 224
	blocks := hdf5MaxHeaderBlocks + 2
	f := h5HostileFile(first + blocks*24)
	h5ObjectHeader(f, 96,
		h5Msg{hdf5MsgContinuation, h5Link(first, 24)},
	)
	for k := range blocks {
		at := first + k*24
		binary.LittleEndian.PutUint16(f[at:], hdf5MsgContinuation)
		if k == blocks-1 {
			// The tail block carries nothing: the walk must be refused
			// on reaching it, not on reading it.
			binary.LittleEndian.PutUint16(f[at+2:], 0)
			continue
		}
		binary.LittleEndian.PutUint16(f[at+2:], 16)
		binary.LittleEndian.PutUint64(f[at+8:], uint64(at+24))
		binary.LittleEndian.PutUint64(f[at+16:], 24)
	}
	_, err := LoadHDF5(writeHostile(t, "chain.h5", f))
	if err == nil {
		t.Fatalf("LoadHDF5 walked a chain of %d continuation blocks", blocks)
	}
	if !strings.Contains(err.Error(), "continuation blocks") {
		t.Fatalf("error = %v, want the chain-depth refusal", err)
	}
}

// TestLoadHDF5V2HeaderFlags covers the version 2 header features the
// reference fixture does not carry: the four times (0x20), the attribute
// counts (0x10) and creation order tracking (0x04, which widens every
// message by its order field). A header with all three set must read
// like any other.
func TestLoadHDF5V2HeaderFlags(t *testing.T) {
	const dataAt = 320
	f := h5V3Superblock(512, 96)
	msgs := slices.Concat(
		h5V2Msg(hdf5MsgDataspace, 1, h5V2Dataspace(2), true),
		h5V2Msg(hdf5MsgDatatype, 2, h5FloatType(8), true),
		h5V2Msg(hdf5MsgDataLayout, 3, h5ContiguousLayout(dataAt, 16), true),
	)
	h5WriteOHDRv2(f, 96, 0x34, msgs)
	binary.LittleEndian.PutUint64(f[dataAt:], math.Float64bits(1.5))
	binary.LittleEndian.PutUint64(f[dataAt+8:], math.Float64bits(-2.5))

	sets, err := LoadHDF5(writeHostile(t, "v2flags.h5", f))
	if err != nil {
		t.Fatalf("LoadHDF5 refused a flagged version 2 header: %v", err)
	}
	if len(sets) != 1 || sets[0].Path != "/" {
		t.Fatalf("datasets = %v, want one dataset at /", sets)
	}
	if got := sets[0].Values.FloatAt(1); got != -2.5 {
		t.Fatalf("value = %v, want -2.5", got)
	}
}

// TestLoadHDF5V2ContinuationChecksum covers the happy path of a version
// 2 continuation: the dataset's messages live in an OCHK block whose
// lookup3 checksum covers the signature and the messages alike, and a
// correct checksum must be accepted (the hostile files above only ever
// see a broken one).
func TestLoadHDF5V2ContinuationChecksum(t *testing.T) {
	const blockAt, dataAt = 224, 320
	f := h5V3Superblock(512, 96)
	msgs := slices.Concat(
		h5V2Msg(hdf5MsgDataspace, 0, h5V2Dataspace(2), false),
		h5V2Msg(hdf5MsgDatatype, 0, h5FloatType(8), false),
		h5V2Msg(hdf5MsgDataLayout, 0, h5ContiguousLayout(dataAt, 16), false),
	)
	block := append([]byte{}, hdf5Chunk2...)
	block = append(block, msgs...)
	block = binary.LittleEndian.AppendUint32(block, hdf5Lookup3(block))
	copy(f[blockAt:], block)

	link := binary.LittleEndian.AppendUint64(
		binary.LittleEndian.AppendUint64([]byte{}, blockAt), uint64(len(block)))
	chunk := make([]byte, 4+len(link))
	chunk[0] = hdf5MsgContinuation
	binary.LittleEndian.PutUint16(chunk[1:], uint16(len(link)))
	copy(chunk[4:], link)
	h5WriteOHDRv2(f, 96, 0, chunk)

	binary.LittleEndian.PutUint64(f[dataAt:], math.Float64bits(1.5))
	binary.LittleEndian.PutUint64(f[dataAt+8:], math.Float64bits(-2.5))

	sets, err := LoadHDF5(writeHostile(t, "v2cont-ok.h5", f))
	if err != nil {
		t.Fatalf("LoadHDF5 refused a checksummed version 2 continuation: %v", err)
	}
	if len(sets) != 1 || sets[0].Path != "/" {
		t.Fatalf("datasets = %v, want one dataset at /", sets)
	}
	if got := sets[0].Values.FloatAt(0); got != 1.5 {
		t.Fatalf("value = %v, want 1.5", got)
	}
}

// TestLoadFITSTableNegativeTFIELDS pins the column count: a negative
// TFIELDS used to size the per-column slices with a negative length and
// panicked, in the binary and the ASCII branch alike. It is refused by
// name in both.
func TestLoadFITSTableNegativeTFIELDS(t *testing.T) {
	for _, kind := range []string{"BINTABLE", "TABLE"} {
		t.Run(kind, func(t *testing.T) {
			hdr := cardBlock(
				card("XTENSION= '"+kind+"'"),
				card("BITPIX  =                    8"),
				card("NAXIS   =                    2"),
				card("NAXIS1  =                    8"),
				card("NAXIS2  =                    1"),
				card("PCOUNT  =                    0"),
				card("GCOUNT  =                    1"),
				card("TFIELDS =                   -1"),
				card("END"),
			)
			path := writeHostile(t, "tfields.fits", append(hdr, make([]byte, 2880)...))
			_, err := LoadFITSTable(path)
			if err == nil {
				t.Fatal("LoadFITSTable accepted TFIELDS = -1")
			}
			if !strings.Contains(err.Error(), "negative") {
				t.Fatalf("error = %v, want the negative-count refusal", err)
			}
		})
	}
}

// TestLoadFITSNaxisRefKeyword pins the NAXIS prefix: any keyword that
// started with NAXIS used to count as an axis, so a user keyword
// NAXISREF = 7 answered "NAXIS = 1 with 2 NAXISn cards" and refused a
// legal file. Only a number after the prefix is an axis, the rule the
// writer applies; NAXIS1 keeps counting.
func TestLoadFITSNaxisRefKeyword(t *testing.T) {
	hdr := cardBlock(
		card("SIMPLE  =                    T"),
		card("BITPIX  =                  -64"),
		card("NAXIS   =                    1"),
		card("NAXIS1  =                    2"),
		card("NAXISREF=                    7"),
		card("END"),
	)
	payload := binary.BigEndian.AppendUint64(nil, math.Float64bits(1.5))
	payload = binary.BigEndian.AppendUint64(payload, math.Float64bits(-0.5))
	a, headers, err := LoadFITS(writeHostile(t, "naxisref.fits", append(hdr, payload...)))
	if err != nil {
		t.Fatalf("LoadFITS refused a file with a NAXISREF keyword: %v", err)
	}
	if a.Len() != 2 || a.FloatAt(0) != 1.5 || a.FloatAt(1) != -0.5 {
		t.Fatalf("values = %v, want 1.5 and -0.5", a)
	}
	if got := headers["NAXISREF"]; got != "7" {
		t.Fatalf("NAXISREF = %q, want it reported as a user keyword", got)
	}
}