package winpwn

import (
	"fmt"
	"io"
)

// This file is the Segment Heap layer on top of heap.go's HeapKind
// detection -- the ROADMAP.md Phase 9 step 4 item: structural parsing of
// _SEGMENT_HEAP and its two sub-backends.
//
// Offsets confirmed via `dt ntdll!_SEGMENT_HEAP`, `dt ntdll!_HEAP_VS_CONTEXT`,
// and `dt ntdll!_HEAP_LFH_CONTEXT` against this machine's build
// (10.0.26100) using the same cdb methodology as heap.go and heap_lfh.go.
//
// What IS and IS NOT implemented in this pass:
//
//   - SegmentHeap outer struct (Signature, GlobalFlags, per-backend
//     summary fields): fully validated against a live heap_segment.exe
//     process reading PEB.ProcessHeaps → DetectHeapKind → ReadSegmentHeap.
//
//   - VS context subsegment enumeration (SubsegmentList walk) and summary
//     stats (TotalCommittedUnits, FreeCommittedUnits): validated.
//
//   - Segment Heap LFH context bucket enumeration (which buckets are
//     active, TotalBlockCount per bucket): validated.
//
//   - Individual VS chunk headers and Segment Heap LFH subsegment
//     BlockOffsets: NOT decoded yet. Both are XOR-encoded against a
//     per-subsegment/per-page key (confirmed empirically: direct reads of
//     _HEAP_VS_CHUNK_HEADER.Sizes and
//     _HEAP_LFH_SUBSEGMENT.BlockOffsets.EncodedData produced implausible
//     field values -- same class of problem as NT Heap LFH's EncodedOffsets,
//     which took its own empirical investigation pass to fix). Decoding them
//     requires recovering the per-page segment offset key, which is its own
//     future validation pass. The AdjacentAddressPairs / FindAdjacentPair
//     helpers in heap.go fill the practical gap for the most common CTF
//     need (finding adjacent same-size allocations from a set of leaked
//     pointers) without needing chunk-level decode at all.

// Offsets confirmed via `dt ntdll!_SEGMENT_HEAP` on build 10.0.26100:
const (
	segHeapOffSignature   = 0x010 // Uint4B -- 0xddeeddee
	segHeapOffGlobalFlags = 0x014 // Uint4B
	segHeapOffVsContext   = 0x280 // inline _HEAP_VS_CONTEXT
	segHeapOffLfhContext  = 0x340 // inline _HEAP_LFH_CONTEXT
)

// VS context sub-offsets (from `dt ntdll!_HEAP_VS_CONTEXT`):
const (
	vsCtxOffFreeChunkTree  = 0x010 // _RTL_RB_TREE (16 bytes, free chunk rb-tree)
	vsCtxOffSubsegmentList = 0x020 // _LIST_ENTRY (head of all VS subsegments)
	vsCtxOffTotalCommitted = 0x030 // Uint8B: committed units
	vsCtxOffFreeCommitted  = 0x038 // Uint8B: free committed units
)

// LFH context sub-offsets (from `dt ntdll!_HEAP_LFH_CONTEXT`):
const (
	segLfhCtxOffBuckets   = 0x080 // [129]Ptr64 _HEAP_LFH_BUCKET
	segLfhBucketCount     = 129
	segLfhBucketOffTotalBlocks = 0x038 // _HEAP_LFH_BUCKET.TotalBlockCount (Uint8B)
)

// VS subsegment (from `dt ntdll!_HEAP_VS_SUBSEGMENT`):
const (
	vsSubsegOffListEntry = 0x000 // _LIST_ENTRY, links into VsContext.SubsegmentList
	vsSubsegOffSize      = 0x020 // Uint2B: size in page-granularity units
	vsSubsegOffSigBits   = 0x022 // bitfield: bits 0-14 = signature, bit 15 = FullCommit
)

// SegmentHeap is the decoded outer shell of ntdll's _SEGMENT_HEAP -- the
// handle passed to HeapAlloc/HeapFree when a process opts into Segment Heap
// (most commonly via an embedded manifest <heapType>segmentHeap</heapType>).
// Use ReadSegmentHeap after DetectHeapKind confirms HeapKindSegment.
type SegmentHeap struct {
	Addr        uint64
	Signature   uint32 // 0xddeeddee -- distinct from NT Heap's 0xeeffeeff and segment-signature 0xffeeffee
	GlobalFlags uint32
	VS          SegmentVSContext  // variable-size backend summary
	LFH         SegmentLFHContext // segment-heap-native LFH summary
}

// SegmentVSContext summarises the VS (variable-size) backend inside a
// Segment Heap -- where allocations outside the LFH's fixed-size buckets
// land. CommittedUnits and FreeUnits are in internal granularity units
// (not bytes); SubsegmentCount is the length of the subsegment list.
type SegmentVSContext struct {
	Addr            uint64 // address of _HEAP_VS_CONTEXT inside the _SEGMENT_HEAP
	CommittedUnits  uint64
	FreeUnits       uint64
	SubsegmentCount int
	Subsegments     []uint64 // address of each _HEAP_VS_SUBSEGMENT
}

// SegmentLFHContext summarises the Segment Heap's native LFH backend --
// a completely different structure from NT Heap's _LFH_HEAP, with its own
// bucket scheme. Each active bucket entry (Ptr64 != 0 and != a scheduling
// stub) is reported with its TotalBlockCount.
type SegmentLFHContext struct {
	Addr         uint64 // address of _HEAP_LFH_CONTEXT inside the _SEGMENT_HEAP
	ActiveBuckets []SegmentLFHBucket
}

// SegmentLFHBucket is one active bucket entry in _HEAP_LFH_CONTEXT.Buckets.
// Index is the zero-based slot in the 129-entry array; TotalBlockCount is
// the cumulative allocation count across all subsegments ever created for
// this size class.
type SegmentLFHBucket struct {
	Index           int
	Ptr             uint64
	TotalBlockCount uint64
}

// ReadSegmentHeap decodes addr's _SEGMENT_HEAP. Returns an error if addr's
// signature isn't the Segment Heap one (use DetectHeapKind first, or
// ReadHeap for NT Heap handles).
func ReadSegmentHeap(r io.ReaderAt, addr uint64) (*SegmentHeap, error) {
	kind, err := DetectHeapKind(r, addr)
	if err != nil {
		return nil, err
	}
	if kind != HeapKindSegment {
		return nil, fmt.Errorf("0x%x is a %s, not a Segment Heap", addr, kind)
	}

	h := &SegmentHeap{Addr: addr}
	if h.Signature, err = readUint32AtValue(r, int64(addr)+segHeapOffSignature); err != nil {
		return nil, fmt.Errorf("Signature: %w", err)
	}
	if h.GlobalFlags, err = readUint32AtValue(r, int64(addr)+segHeapOffGlobalFlags); err != nil {
		return nil, fmt.Errorf("GlobalFlags: %w", err)
	}

	vsCtxAddr := addr + segHeapOffVsContext
	h.VS.Addr = vsCtxAddr
	if h.VS.CommittedUnits, err = readUint64At(r, int64(vsCtxAddr)+vsCtxOffTotalCommitted); err != nil {
		return nil, fmt.Errorf("VS.TotalCommittedUnits: %w", err)
	}
	if h.VS.FreeUnits, err = readUint64At(r, int64(vsCtxAddr)+vsCtxOffFreeCommitted); err != nil {
		return nil, fmt.Errorf("VS.FreeCommittedUnits: %w", err)
	}
	if h.VS.Subsegments, err = walkVSSubsegmentList(r, vsCtxAddr); err != nil {
		return nil, fmt.Errorf("VS subsegment list: %w", err)
	}
	h.VS.SubsegmentCount = len(h.VS.Subsegments)

	lfhCtxAddr := addr + segHeapOffLfhContext
	h.LFH.Addr = lfhCtxAddr
	if h.LFH.ActiveBuckets, err = readSegmentLFHBuckets(r, lfhCtxAddr); err != nil {
		return nil, fmt.Errorf("LFH buckets: %w", err)
	}

	return h, nil
}

// walkVSSubsegmentList enumerates the _HEAP_VS_SUBSEGMENT addresses by
// following the SubsegmentList LIST_ENTRY chain in the VS context.
func walkVSSubsegmentList(r io.ReaderAt, vsCtxAddr uint64) ([]uint64, error) {
	headAddr := vsCtxAddr + vsCtxOffSubsegmentList
	flink, err := readUint64At(r, int64(headAddr))
	if err != nil {
		return nil, fmt.Errorf("reading SubsegmentList head: %w", err)
	}

	var subsegments []uint64
	cur := flink
	for cur != headAddr && cur != 0 {
		subsegments = append(subsegments, cur) // ListEntry is at offset 0, so cur == subsegment addr
		if len(subsegments) > 4096 {
			return subsegments, fmt.Errorf("VS SubsegmentList exceeded 4096 entries (corrupted?)")
		}
		next, err := readUint64At(r, int64(cur)) // Flink is LIST_ENTRY's first field
		if err != nil || next == cur {
			break
		}
		cur = next
	}
	return subsegments, nil
}

// readSegmentLFHBuckets scans the 129-entry _HEAP_LFH_CONTEXT.Buckets array
// and returns every active entry (non-null pointer that isn't a scheduling
// stub, identified by low bit clear in the pointer value).
func readSegmentLFHBuckets(r io.ReaderAt, lfhCtxAddr uint64) ([]SegmentLFHBucket, error) {
	var buckets []SegmentLFHBucket
	bucketsBase := int64(lfhCtxAddr) + segLfhCtxOffBuckets
	for i := 0; i < segLfhBucketCount; i++ {
		ptr, err := readUint64At(r, bucketsBase+int64(i)*8)
		if err != nil {
			return buckets, fmt.Errorf("reading bucket[%d]: %w", i, err)
		}
		// Low bit set means this entry is a scheduler stub, not a real bucket pointer
		if ptr == 0 || ptr&1 != 0 {
			continue
		}
		total, err := readUint64At(r, int64(ptr)+segLfhBucketOffTotalBlocks)
		if err != nil {
			continue
		}
		if total == 0 {
			continue
		}
		buckets = append(buckets, SegmentLFHBucket{Index: i, Ptr: ptr, TotalBlockCount: total})
	}
	return buckets, nil
}

// AdjacentAddressPairs finds all pairs in addrs where the difference is
// exactly step bytes -- the structural-equivalent finder for "which two
// same-size allocations landed adjacent" that examples/heap_segment's
// spray loop discovers by trial and error. On Segment Heap, same-size
// allocations in the same subsegment page are packed step bytes apart
// (step == sizeof(Allocation), before any chunk-header overhead, which
// the Segment Heap's LFH accounts for separately from user data unlike
// NT Heap's HeapEntrySize scheme). Returns all (lo, hi) pairs in
// address order with hi == lo+step.
//
// CAUTION: step is the ALLOCATION GRANULARITY visible at the HeapAlloc
// caller level (e.g. sizeof(Profile)=32 in heap_segment.c), not
// sizeof(struct) + sizeof(chunk_header) -- Segment Heap's metadata
// isolation places chunk headers on a separate metadata page, so the
// gap between two adjacent user payloads really is sizeof(Allocation).
// Verify empirically for your specific build if this doesn't match.
func AdjacentAddressPairs(addrs []uint64, step uint64) [][2]uint64 {
	set := make(map[uint64]struct{}, len(addrs))
	for _, a := range addrs {
		set[a] = struct{}{}
	}
	var pairs [][2]uint64
	for _, a := range addrs {
		if _, ok := set[a+step]; ok {
			pairs = append(pairs, [2]uint64{a, a + step})
		}
	}
	return pairs
}

// FindAdjacentPair returns the first pair where hi == lo+step, or
// (0, 0, false) if none exists. Convenience wrapper over AdjacentAddressPairs
// for the common "give me any adjacent pair" case.
func FindAdjacentPair(addrs []uint64, step uint64) (lo, hi uint64, found bool) {
	pairs := AdjacentAddressPairs(addrs, step)
	if len(pairs) == 0 {
		return 0, 0, false
	}
	return pairs[0][0], pairs[0][1], true
}