package runtime

import (
	"fmt"
	"strings"
	"unsafe"

	"golang.org/x/sys/windows"
)

// Object manager directory enumeration. NtOpenDirectoryObject and
// NtQueryDirectoryObject are undocumented NT APIs, so they are called directly
// through ntdll.dll (x/sys/windows has no wrappers for them).

var (
	modntdll                  = windows.NewLazySystemDLL("ntdll.dll")
	procNtOpenDirectoryObject = modntdll.NewProc("NtOpenDirectoryObject")
	procNtQueryDirectoryObject = modntdll.NewProc("NtQueryDirectoryObject")
)

// DIRECTORY_QUERY is the access right needed to enumerate a directory object.
const directoryQuery = 0x0001

// unicodeString mirrors the kernel UNICODE_STRING structure.
type unicodeString struct {
	Length        uint16
	MaximumLength uint16
	Buffer        *uint16
}

// objectAttributes mirrors OBJECT_ATTRIBUTES.
type objectAttributes struct {
	Length                   uint32
	RootDirectory            windows.Handle
	ObjectName               *unicodeString
	Attributes               uint32
	SecurityDescriptor       unsafe.Pointer
	SecurityQualityOfService unsafe.Pointer
}

// objectDirectoryInformation mirrors OBJECT_DIRECTORY_INFORMATION, the entry
// type returned by NtQueryDirectoryObject.
type objectDirectoryInformation struct {
	Name     unicodeString
	TypeName unicodeString
}

// ntOpenDirectoryObject opens an object manager directory by name.
func ntOpenDirectoryObject(handle *windows.Handle, access uint32, attrs *objectAttributes) error {
	r0, _, _ := procNtOpenDirectoryObject.Call(
		uintptr(unsafe.Pointer(handle)),
		uintptr(access),
		uintptr(unsafe.Pointer(attrs)),
	)
	if r0 != 0 {
		return windows.NTStatus(r0)
	}
	return nil
}

// ntQueryDirectoryObject enumerates a directory object into buf, returning the
// number of bytes written in retLen.
func ntQueryDirectoryObject(handle windows.Handle, buf *byte, length uint32, restart bool, context *uint32, retLen *uint32) error {
	var restartFlag uintptr
	if restart {
		restartFlag = 1
	}
	r0, _, _ := procNtQueryDirectoryObject.Call(
		uintptr(handle),
		uintptr(unsafe.Pointer(buf)),
		uintptr(length),
		0, // ReturnSingleEntry = false
		restartFlag,
		uintptr(unsafe.Pointer(context)),
		uintptr(unsafe.Pointer(retLen)),
	)
	if r0 != 0 {
		return windows.NTStatus(r0)
	}
	return nil
}

// EnumerateDirectory returns the names of all objects in an object manager
// directory (e.g. `\Device`, `\BaseNamedObjects`, `\GLOBAL??`).
func EnumerateDirectory(dir string) ([]string, error) {
	name, err := windows.UTF16PtrFromString(dir)
	if err != nil {
		return nil, err
	}
	us := &unicodeString{
		Length:        uint16(len(dir) * 2),
		MaximumLength: uint16((len(dir) + 1) * 2),
		Buffer:        name,
	}
	attrs := &objectAttributes{
		Length:     uint32(unsafe.Sizeof(objectAttributes{})),
		ObjectName: us,
	}

	var h windows.Handle
	if err := ntOpenDirectoryObject(&h, directoryQuery, attrs); err != nil {
		return nil, err
	}
	defer windows.CloseHandle(h)

	// Enumerate in passes. A directory can hold more entries than fit in one
	// buffer, in which case NtQueryDirectoryObject returns STATUS_MORE_ENTRIES
	// (or STATUS_BUFFER_TOO_SMALL for a single oversized entry). Grow the buffer
	// and restart the scan from the beginning in that case.
	bufSize := 64 * 1024
	for {
		buf := make([]byte, bufSize)
		var out []string
		var context uint32
		restart := true
		needGrow := false
		for {
			var retLen uint32
			err := ntQueryDirectoryObject(h, &buf[0], uint32(len(buf)), restart, &context, &retLen)
			if err == windows.STATUS_NO_MORE_ENTRIES {
				break
			}
			if err == windows.STATUS_MORE_ENTRIES || err == windows.STATUS_BUFFER_TOO_SMALL {
				needGrow = true
				break
			}
			if err != nil {
				return nil, err
			}
			out = append(out, parseDirectoryEntries(buf[:retLen])...)
			restart = false
		}
		if !needGrow {
			return out, nil
		}
		if bufSize >= 16*1024*1024 {
			return nil, fmt.Errorf("directory %q too large to enumerate", dir)
		}
		bufSize *= 2
	}
}

// parseDirectoryEntries decodes a buffer of OBJECT_DIRECTORY_INFORMATION
// entries. The fixed structures are packed at the start of the buffer; the Name
// and TypeName string data follows them, and each entry's Name.Buffer /
// TypeName.Buffer is a real pointer into that trailing region (not inline data).
func parseDirectoryEntries(buf []byte) []string {
	base := uintptr(unsafe.Pointer(&buf[0]))
	entrySize := int(unsafe.Sizeof(objectDirectoryInformation{}))
	if len(buf) < entrySize {
		return nil
	}

	// The first entry's Name.Buffer points at the start of the string-data
	// region, which begins immediately after the packed structures. That gives
	// the entry count without scanning into the string bytes.
	first := (*objectDirectoryInformation)(unsafe.Pointer(&buf[0]))
	stringStart := uintptr(unsafe.Pointer(first.Name.Buffer))
	if stringStart < base {
		return nil
	}
	numEntries := int(stringStart-base) / entrySize

	var out []string
	for i := 0; i < numEntries; i++ {
		info := (*objectDirectoryInformation)(unsafe.Pointer(&buf[i*entrySize]))
		nameLen := int(info.Name.Length)
		if nameLen == 0 {
			continue
		}
		namePtr := uintptr(unsafe.Pointer(info.Name.Buffer))
		if namePtr < base || namePtr+uintptr(nameLen) > base+uintptr(len(buf)) {
			continue
		}
		rel := namePtr - base
		out = append(out, utf16BytesToString(buf[rel:rel+uintptr(nameLen)]))
	}
	return out
}

// utf16BytesToString converts a UTF-16LE byte slice to a Go string.
func utf16BytesToString(b []byte) string {
	u := make([]uint16, len(b)/2)
	for i := range u {
		u[i] = uint16(b[i*2]) | uint16(b[i*2+1])<<8
	}
	return windows.UTF16ToString(u)
}

// DeviceNames returns the user-mode device paths (`\\.\X`) for every device and
// symbolic link visible in the object manager. It enumerates `\GLOBAL??` (the
// DosDevices namespace) and `\Device`.
func DeviceNames() []string {
	seen := map[string]bool{}
	var out []string
	add := func(p string) {
		if p == "" || seen[p] {
			return
		}
		seen[p] = true
		out = append(out, p)
	}

	// \GLOBAL?? holds the symbolic links (e.g. RTCore64 -> \Device\RTCore64).
	if names, err := EnumerateDirectory(`\GLOBAL??`); err == nil {
		for _, n := range names {
			add(`\\.\` + n)
		}
	}
	// \Device holds the device objects themselves.
	if names, err := EnumerateDirectory(`\Device`); err == nil {
		for _, n := range names {
			add(`\\.\` + n)
		}
	}
	return out
}

// SectionNames returns the names of named section objects in \BaseNamedObjects,
// mapped to the user-mode Global\ namespace.
func SectionNames() []string {
	names, err := EnumerateDirectory(`\BaseNamedObjects`)
	if err != nil {
		return nil
	}
	var out []string
	for _, n := range names {
		if strings.HasPrefix(n, "Global\\") {
			out = append(out, n)
		} else {
			out = append(out, `Global\`+n)
		}
	}
	return out
}