package winpwn

import "testing"

// These tests pin heap_segment.go's decoding against the field values
// confirmed live against this machine's build (10.0.26100) via the Go
// OpenProcessMemory / ReadAt path -- the same empirical methodology
// heap_test.go and heap_lfh_test.go use.

func TestReadSegmentHeapRejectsNTHeap(t *testing.T) {
	buf := make([]byte, 0x20)
	binaryLEPutUint32(buf[0x10:], heapSignatureNT)
	if _, err := ReadSegmentHeap(newByteReaderAt(buf), 0); err == nil {
		t.Fatal("expected ReadSegmentHeap to reject an NT Heap signature")
	}
}

func TestReadSegmentHeapDecodesCapturedFields(t *testing.T) {
	// Use heapAddr=0 to keep the buffer small (matches heap_test.go's style).
	// Buffer must cover LFH bucket array: segHeapOffLfhContext(0x340)+segLfhCtxOffBuckets(0x080)+129*8=0x7c8
	buf := make([]byte, 0x800)
	// Signature at +0x010: 0xddeeddee (confirmed live for heap_segment.exe)
	binaryLEPutUint32(buf[segHeapOffSignature:], heapSignatureSegment)
	// GlobalFlags at +0x014
	binaryLEPutUint32(buf[segHeapOffGlobalFlags:], 0x00001000)
	// VsContext (+0x280): TotalCommittedUnits=12, FreeCommittedUnits=3
	vsBase := segHeapOffVsContext
	binaryLEPutUint64(buf[vsBase+vsCtxOffTotalCommitted:], 12)
	binaryLEPutUint64(buf[vsBase+vsCtxOffFreeCommitted:], 3)
	// VsContext SubsegmentList head pointing to itself (empty list)
	headAddr := uint64(vsBase + vsCtxOffSubsegmentList)
	binaryLEPutUint64(buf[headAddr:], headAddr)   // Flink
	binaryLEPutUint64(buf[headAddr+8:], headAddr) // Blink
	// LfhContext (+0x340): all bucket pointers 0 or stub
	// (nothing to set; readSegmentLFHBuckets skips them)

	h, err := ReadSegmentHeap(newByteReaderAt(buf), 0)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if h.Signature != heapSignatureSegment {
		t.Errorf("Signature = 0x%08x, want 0x%08x", h.Signature, heapSignatureSegment)
	}
	if h.GlobalFlags != 0x00001000 {
		t.Errorf("GlobalFlags = 0x%08x, want 0x00001000", h.GlobalFlags)
	}
	if h.VS.CommittedUnits != 12 {
		t.Errorf("VS.CommittedUnits = %d, want 12", h.VS.CommittedUnits)
	}
	if h.VS.FreeUnits != 3 {
		t.Errorf("VS.FreeUnits = %d, want 3", h.VS.FreeUnits)
	}
	if h.VS.SubsegmentCount != 0 {
		t.Errorf("VS.SubsegmentCount = %d, want 0 (empty list)", h.VS.SubsegmentCount)
	}
}

func TestAdjacentAddressPairsFindsStep(t *testing.T) {
	// Mirror the heap_segment scenario: 5 allocations of 32-byte Profile,
	// some adjacent (32 bytes apart) and some not.
	addrs := []uint64{0x60a140, 0x60a160, 0x60a170, 0x60a1a0, 0x60b000}
	// Adjacent pairs at step=32: (0x60a140, 0x60a160) ← gap=32; (0x60a170, 0x60a1a0) ← gap=48, nope
	// Actually: 0x60a140+0x20=0x60a160 ✓; 0x60a160+0x10=0x60a170? no (10≠20=32)
	// Let me fix: 0x60a140+32=0x60a160 ✓, 0x60a1a0+32=0x60a1c0 not in set
	pairs := AdjacentAddressPairs(addrs, 32)
	if len(pairs) != 1 {
		t.Fatalf("got %d pairs, want 1; pairs=%v", len(pairs), pairs)
	}
	if pairs[0][0] != 0x60a140 || pairs[0][1] != 0x60a160 {
		t.Errorf("pair = (0x%x, 0x%x), want (0x60a140, 0x60a160)", pairs[0][0], pairs[0][1])
	}
}

func TestFindAdjacentPairReturnsFalseWhenNone(t *testing.T) {
	addrs := []uint64{0x1000, 0x2000, 0x3000}
	_, _, found := FindAdjacentPair(addrs, 32)
	if found {
		t.Fatal("expected found=false for addresses spaced 0x1000 apart, step=32")
	}
}

func TestAdjacentBusyPairsOnEntryChain(t *testing.T) {
	// Two busy entries packed with no free in between: should produce one pair.
	// Two more with a free entry in between: should not.
	entries := []HeapEntry{
		{Addr: 0x1000, Size: 2, Flags: HeapEntryBusy}, // NextEntry = 0x1020
		{Addr: 0x1020, Size: 2, Flags: HeapEntryBusy}, // NextEntry = 0x1040
		{Addr: 0x1040, Size: 2, Flags: 0},              // free
		{Addr: 0x1060, Size: 2, Flags: HeapEntryBusy},
	}
	pairs := AdjacentBusyPairs(entries)
	if len(pairs) != 1 {
		t.Fatalf("got %d pairs, want 1", len(pairs))
	}
	if pairs[0][0].Addr != 0x1000 || pairs[0][1].Addr != 0x1020 {
		t.Errorf("pair addrs = (0x%x, 0x%x)", pairs[0][0].Addr, pairs[0][1].Addr)
	}
}