package win32api

import (
	"syscall"
	"unsafe"

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

func DuplicateTokenEx(hExistingToken HANDLE, dwDesiredAccess DWORD,
	lpTokenAttributes uintptr, ImpersonationLevel int,
	TokenType TOKEN_TYPE, phNewToken *TOKEN) error {

	Dup, err := getProcAddr("advapi32.dll", "DuplicateTokenEx")
	if err != nil {
		return err
	}
	r, _, errno := syscall.Syscall6(Dup, 6, uintptr(hExistingToken), uintptr(dwDesiredAccess), lpTokenAttributes, uintptr(ImpersonationLevel),
		uintptr(TokenType), uintptr(unsafe.Pointer(phNewToken)))
	if r == 0 {
		return error(errno)
	}
	return nil
}

func CreateProcessAsUser(hToken TOKEN, lpApplicationName, lpCommandLine string,
	lpProcessAttributes, lpThreadAttributes, bInheritHandles uintptr,
	dwCreationFlags DWORD, lpEnvironment HANDLE, lpCurrentDirectory string,
	lpStartupInfo *StartupInfo, lpProcessInformation *ProcessInformation) error {
	var (
		applicationName uintptr
		commandLine     uintptr
		workingDir      uintptr
	)
	CPAU, err := getProcAddr("advapi32.dll", "CreateProcessAsUserW")
	if err != nil {
		return err
	}
	if len(lpApplicationName) > 0 {
		applicationName = uintptr(unsafe.Pointer(windows.StringToUTF16Ptr(lpApplicationName)))
	}
	if len(lpCommandLine) > 0 {
		commandLine = uintptr(unsafe.Pointer(windows.StringToUTF16Ptr(lpCommandLine)))
	}
	if len(lpCurrentDirectory) > 0 {
		workingDir = uintptr(unsafe.Pointer(windows.StringToUTF16Ptr(lpCurrentDirectory)))
	}
	r, _, errno := syscall.Syscall12(CPAU, 11, uintptr(hToken), applicationName,
		commandLine, lpProcessAttributes, lpThreadAttributes, bInheritHandles, uintptr(dwCreationFlags), uintptr(lpEnvironment),
		workingDir, uintptr(unsafe.Pointer(lpStartupInfo)), uintptr(unsafe.Pointer(lpProcessInformation)), 0)
	if r == 0 {
		return error(errno)
	}
	return nil
}
func GetTokenInformation(TokenHandle HANDLE, TokenInformationClass, TokenInformation,
	TokenInformationLength uintptr, ReturnLength *uintptr) error {
	GTI, err := getProcAddr("advapi32.dll", "GetTokenInformation")
	if err != nil {
		return err
	}
	if r, _, errno := syscall.Syscall6(GTI, 5, uintptr(TokenHandle), TokenInformationClass,
		TokenInformation, TokenInformationLength, uintptr(unsafe.Pointer(ReturnLength)), 0); r == 0 {
		return error(errno)
	}
	return nil
}

func GetUserName() (string, error) {
	gun, err := getProcAddr("advapi32.dll", "GetUserNameW")
	if err != nil {
		return "", err
	}

	size := uint32(64)
	for {
		buf := make([]uint16, size)
		n := uint32(len(buf))
		r, _, errno := syscall.Syscall(gun, 2, uintptr(unsafe.Pointer(&buf[0])), uintptr(unsafe.Pointer(&n)), 0)
		if r != 0 {
			return syscall.UTF16ToString(buf), nil
		}
		if errno == syscall.ERROR_INSUFFICIENT_BUFFER {
			if n > size {
				size = n
			} else {
				size *= 2
			}
			continue
		}
		if errno != 0 {
			return "", error(errno)
		}
		return "", syscall.EINVAL
	}
}

func OpenProcessToken(ProcessHandle HANDLE, DesiredAccess DWORD, TokenHandle *TOKEN) error {
	if TokenHandle == nil {
		return syscall.EINVAL
	}
	proc, err := getProcAddr("advapi32.dll", "OpenProcessToken")
	if err != nil {
		return err
	}
	r, _, errno := syscall.Syscall(proc, 3, uintptr(ProcessHandle), uintptr(DesiredAccess), uintptr(unsafe.Pointer(TokenHandle)))
	if r == 0 {
		if errno != 0 {
			return error(errno)
		}
		return syscall.EINVAL
	}
	return nil
}

func CheckTokenMembership(tokenHandle HANDLE, sidToCheck unsafe.Pointer) (bool, error) {
	if sidToCheck == nil {
		return false, syscall.EINVAL
	}
	proc, err := getProcAddr("advapi32.dll", "CheckTokenMembership")
	if err != nil {
		return false, err
	}
	var isMember int32
	r, _, errno := syscall.Syscall(proc, 3,
		uintptr(tokenHandle),
		uintptr(sidToCheck),
		uintptr(unsafe.Pointer(&isMember)),
	)
	if r == 0 {
		if errno != 0 {
			return false, error(errno)
		}
		return false, syscall.EINVAL
	}
	return isMember != 0, nil
}

func LookupPrivilegeValue(lpSystemName, lpName string, lpLuid *LUID) error {
	if lpLuid == nil {
		return syscall.EINVAL
	}
	proc, err := getProcAddr("advapi32.dll", "LookupPrivilegeValueW")
	if err != nil {
		return err
	}

	var systemNamePtr uintptr
	if len(lpSystemName) > 0 {
		systemNamePtr = uintptr(unsafe.Pointer(windows.StringToUTF16Ptr(lpSystemName)))
	}
	r, _, errno := syscall.Syscall(proc, 3,
		systemNamePtr,
		uintptr(unsafe.Pointer(windows.StringToUTF16Ptr(lpName))),
		uintptr(unsafe.Pointer(lpLuid)),
	)
	if r == 0 {
		if errno != 0 {
			return error(errno)
		}
		return syscall.EINVAL
	}
	return nil
}

func AdjustTokenPrivileges(TokenHandle TOKEN, DisableAllPrivileges bool, NewState *TOKEN_PRIVILEGES, BufferLength DWORD,
	PreviousState *TOKEN_PRIVILEGES, ReturnLength *DWORD) error {
	proc, err := getProcAddr("advapi32.dll", "AdjustTokenPrivileges")
	if err != nil {
		return err
	}

	disableAll := uintptr(0)
	if DisableAllPrivileges {
		disableAll = 1
	}
	r, _, errno := syscall.Syscall6(proc, 6,
		uintptr(TokenHandle),
		disableAll,
		uintptr(unsafe.Pointer(NewState)),
		uintptr(BufferLength),
		uintptr(unsafe.Pointer(PreviousState)),
		uintptr(unsafe.Pointer(ReturnLength)),
	)
	if r == 0 {
		if errno != 0 {
			return error(errno)
		}
		return syscall.EINVAL
	}
	if errno == windows.ERROR_NOT_ALL_ASSIGNED {
		return error(errno)
	}
	return nil
}

func RevertToSelf() error {
	proc, err := getProcAddr("advapi32.dll", "RevertToSelf")
	if err != nil {
		return err
	}
	r, _, errno := syscall.Syscall(proc, 0, 0, 0, 0)
	if r == 0 {
		if errno != 0 {
			return error(errno)
		}
		return syscall.EINVAL
	}
	return nil
}

func CreateProcessWithToken(hToken TOKEN, dwLogonFlags DWORD, lpApplicationName, lpCommandLine string,
	dwCreationFlags DWORD, lpEnvironment HANDLE, lpCurrentDirectory string,
	lpStartupInfo *StartupInfo, lpProcessInformation *ProcessInformation) error {
	var (
		applicationName uintptr
		commandLine     uintptr
		currentDir      uintptr
	)
	if len(lpApplicationName) > 0 {
		applicationName = uintptr(unsafe.Pointer(windows.StringToUTF16Ptr(lpApplicationName)))
	}
	if len(lpCommandLine) > 0 {
		commandLine = uintptr(unsafe.Pointer(windows.StringToUTF16Ptr(lpCommandLine)))
	}
	if len(lpCurrentDirectory) > 0 {
		currentDir = uintptr(unsafe.Pointer(windows.StringToUTF16Ptr(lpCurrentDirectory)))
	}

	proc, err := getProcAddr("advapi32.dll", "CreateProcessWithTokenW")
	if err != nil {
		return err
	}
	r, _, errno := syscall.Syscall12(proc, 9,
		uintptr(hToken),
		uintptr(dwLogonFlags),
		applicationName,
		commandLine,
		uintptr(dwCreationFlags),
		uintptr(lpEnvironment),
		currentDir,
		uintptr(unsafe.Pointer(lpStartupInfo)),
		uintptr(unsafe.Pointer(lpProcessInformation)),
		0,
		0,
		0,
	)
	if r == 0 {
		if errno != 0 {
			return error(errno)
		}
		return syscall.EINVAL
	}
	return nil
}

func IsTokenElevated(token TOKEN) (bool, error) {
	var elevation TOKEN_ELEVATION
	var retLen uintptr
	if err := GetTokenInformation(
		HANDLE(token),
		TokenElevation,
		uintptr(unsafe.Pointer(&elevation)),
		uintptr(unsafe.Sizeof(elevation)),
		&retLen,
	); err != nil {
		return false, err
	}
	return elevation.TokenIsElevated != 0, nil
}

func IsCurrentProcessElevated() (bool, error) {
	processHandle, err := syscall.GetCurrentProcess()
	if err != nil {
		return false, err
	}
	var token TOKEN
	if err := OpenProcessToken(HANDLE(processHandle), TOKEN_QUERY, &token); err != nil {
		return false, err
	}
	defer func() {
		_ = CloseHandle(HANDLE(token))
	}()
	return IsTokenElevated(token)
}

func IsCurrentUserInAdminGroup() (bool, error) {
	adminSID, err := windows.CreateWellKnownSid(windows.WinBuiltinAdministratorsSid)
	if err != nil {
		return false, err
	}

	// Passing tokenHandle=0 lets Windows use the calling thread/process effective token.
	return CheckTokenMembership(0, unsafe.Pointer(adminSID))
}