package memory

import (
	"encoding/binary"
	"errors"
	"fmt"
	"sync"
	"syscall"
	"unsafe"

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

type objectAttributes struct {
	Length                   uint32
	RootDirectory            windows.Handle
	ObjectName               uintptr
	Attributes               uint32
	SecurityDescriptor       uintptr
	SecurityQualityOfService uintptr
}

type clientID struct {
	UniqueProcess uintptr
	UniqueThread  uintptr
}

var (
	dsInitOnce sync.Once
	dsInitErr  error

	ntOpenProcessStub uintptr
	ntReadVMStub      uintptr
	ntWriteVMStub     uintptr
	ntCloseStub       uintptr
	ntAllocVMStub     uintptr
	ntProtectVMStub   uintptr
	ntFreeVMStub      uintptr
)

func initDirectSyscalls() {
	var err error
	ntOpenProcessStub, err = buildNtStub("NtOpenProcess")
	if err != nil {
		dsInitErr = err
		return
	}
	ntReadVMStub, err = buildNtStub("NtReadVirtualMemory")
	if err != nil {
		dsInitErr = err
		return
	}
	ntWriteVMStub, err = buildNtStub("NtWriteVirtualMemory")
	if err != nil {
		dsInitErr = err
		return
	}
	ntCloseStub, err = buildNtStub("NtClose")
	if err != nil {
		dsInitErr = err
		return
	}
	ntAllocVMStub, err = buildNtStub("NtAllocateVirtualMemory")
	if err != nil {
		dsInitErr = err
		return
	}
	ntProtectVMStub, err = buildNtStub("NtProtectVirtualMemory")
	if err != nil {
		dsInitErr = err
		return
	}
	ntFreeVMStub, err = buildNtStub("NtFreeVirtualMemory")
	if err != nil {
		dsInitErr = err
		return
	}
}

func ensureDirectSyscalls() error {
	dsInitOnce.Do(initDirectSyscalls)
	return dsInitErr
}

func buildNtStub(name string) (uintptr, error) {
	ntdll := windows.NewLazySystemDLL("ntdll.dll")
	proc := ntdll.NewProc(name)
	if err := ntdll.Load(); err != nil {
		return 0, fmt.Errorf("load ntdll.dll: %w", err)
	}

	addr := proc.Addr()
	if addr == 0 {
		return 0, fmt.Errorf("%s address is zero", name)
	}
	sysID, err := extractSysID(addr)
	if err != nil {
		return 0, fmt.Errorf("extract %s sysid: %w", name, err)
	}
	return makeSyscallStub(sysID)
}

func extractSysID(procAddr uintptr) (uint32, error) {
	b := unsafe.Slice((*byte)(unsafe.Pointer(procAddr)), 64)
	if len(b) < 11 {
		return 0, errors.New("short Nt stub")
	}
	if b[0] != 0x4C || b[1] != 0x8B || b[2] != 0xD1 || b[3] != 0xB8 {
		return 0, errors.New("unexpected Nt stub prefix")
	}
	found := false
	for i := 0; i+2 < len(b); i++ {
		if b[i] == 0x0F && b[i+1] == 0x05 && b[i+2] == 0xC3 {
			found = true
			break
		}
	}
	if !found {
		return 0, errors.New("syscall;ret not found")
	}
	return binary.LittleEndian.Uint32(b[4:8]), nil
}

func makeSyscallStub(sysID uint32) (uintptr, error) {
	code := []byte{0x4C, 0x8B, 0xD1, 0xB8, 0, 0, 0, 0, 0x0F, 0x05, 0xC3}
	binary.LittleEndian.PutUint32(code[4:8], sysID)
	mem, err := windows.VirtualAlloc(0, uintptr(len(code)), windows.MEM_COMMIT|windows.MEM_RESERVE, windows.PAGE_READWRITE)
	if err != nil {
		return 0, err
	}
	copy(unsafe.Slice((*byte)(unsafe.Pointer(mem)), len(code)), code)
	var oldProtect uint32
	err = windows.VirtualProtect(mem, uintptr(len(code)), windows.PAGE_EXECUTE_READ, &oldProtect)
	if err != nil {
		return 0, err
	}

	return mem, nil
}

func ntStatusErr(op string, status uintptr) error {
	return fmt.Errorf("%s failed: NTSTATUS=0x%08X", op, uint32(status))
}

func ntOpenProcess(pid uint32, access uint32) (windows.Handle, error) {
	if err := ensureDirectSyscalls(); err != nil {
		return 0, err
	}
	var h windows.Handle
	oa := objectAttributes{Length: uint32(unsafe.Sizeof(objectAttributes{}))}
	cid := clientID{UniqueProcess: uintptr(pid)}
	status, _, _ := syscall.SyscallN(
		ntOpenProcessStub,
		uintptr(unsafe.Pointer(&h)),
		uintptr(access),
		uintptr(unsafe.Pointer(&oa)),
		uintptr(unsafe.Pointer(&cid)),
	)
	if status != 0 {
		return 0, ntStatusErr("NtOpenProcess", status)
	}
	return h, nil
}

func ntReadProcessMemory(handle windows.Handle, address uintptr, out []byte) error {
	if len(out) == 0 {
		return nil
	}
	if err := ensureDirectSyscalls(); err != nil {
		return err
	}
	var bytesRead uintptr
	status, _, _ := syscall.SyscallN(
		ntReadVMStub,
		uintptr(handle),
		address,
		uintptr(unsafe.Pointer(&out[0])),
		uintptr(len(out)),
		uintptr(unsafe.Pointer(&bytesRead)),
	)
	if status != 0 {
		return ntStatusErr("NtReadVirtualMemory", status)
	}
	return nil
}

func ntWriteProcessMemory(handle windows.Handle, address uintptr, in []byte) error {
	if len(in) == 0 {
		return nil
	}
	if err := ensureDirectSyscalls(); err != nil {
		return err
	}
	var bytesWritten uintptr
	status, _, _ := syscall.SyscallN(
		ntWriteVMStub,
		uintptr(handle),
		address,
		uintptr(unsafe.Pointer(&in[0])),
		uintptr(len(in)),
		uintptr(unsafe.Pointer(&bytesWritten)),
	)
	if status != 0 {
		return ntStatusErr("NtWriteVirtualMemory", status)
	}
	return nil
}

func ntCloseHandle(handle windows.Handle) error {
	if handle == 0 {
		return nil
	}
	if err := ensureDirectSyscalls(); err != nil {
		return err
	}
	status, _, _ := syscall.SyscallN(ntCloseStub, uintptr(handle))
	if status != 0 {
		return ntStatusErr("NtClose", status)
	}
	return nil
}

func ntAllocateVirtualMemory(handle windows.Handle, size uintptr, allocType, protect uint32) (uintptr, error) {
	if err := ensureDirectSyscalls(); err != nil {
		return 0, err
	}
	base := uintptr(0)
	regionSize := size
	status, _, _ := syscall.SyscallN(
		ntAllocVMStub,
		uintptr(handle),
		uintptr(unsafe.Pointer(&base)),
		0,
		uintptr(unsafe.Pointer(&regionSize)),
		uintptr(allocType),
		uintptr(protect),
	)
	if status != 0 {
		return 0, ntStatusErr("NtAllocateVirtualMemory", status)
	}
	return base, nil
}

func ntProtectVirtualMemory(handle windows.Handle, address uintptr, size uintptr, newProtect uint32) (uint32, error) {
	if err := ensureDirectSyscalls(); err != nil {
		return 0, err
	}
	base := address
	regionSize := size
	var oldProtect uint32
	status, _, _ := syscall.SyscallN(
		ntProtectVMStub,
		uintptr(handle),
		uintptr(unsafe.Pointer(&base)),
		uintptr(unsafe.Pointer(&regionSize)),
		uintptr(newProtect),
		uintptr(unsafe.Pointer(&oldProtect)),
	)
	if status != 0 {
		return 0, ntStatusErr("NtProtectVirtualMemory", status)
	}
	return oldProtect, nil
}

func ntFreeVirtualMemory(handle windows.Handle, address uintptr, freeType uint32) error {
	if address == 0 {
		return nil
	}
	if err := ensureDirectSyscalls(); err != nil {
		return err
	}
	base := address
	regionSize := uintptr(0) // required for MEM_RELEASE
	status, _, _ := syscall.SyscallN(
		ntFreeVMStub,
		uintptr(handle),
		uintptr(unsafe.Pointer(&base)),
		uintptr(unsafe.Pointer(&regionSize)),
		uintptr(freeType),
	)
	if status != 0 {
		return ntStatusErr("NtFreeVirtualMemory", status)
	}
	return nil
}