//go:build windows

package windows

import (
	"fmt"
	"unsafe"

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

const (
	wtsCurrentServerHandle windows.Handle = 0
	invalidSessionID                      = ^uint32(0)
)

// DiscoverActiveSessions translates only verified active WTS user tokens into
// selector DTOs. It is intended for the LocalSystem service: WTSQueryUserToken
// requires LocalSystem and SeTcbPrivilege. A token-query failure leaves that
// WTS session unusable rather than guessing an identity.
func DiscoverActiveSessions() ([]SessionCandidate, error) {
	var sessions *windows.WTS_SESSION_INFO
	var count uint32
	if err := windows.WTSEnumerateSessions(wtsCurrentServerHandle, 0, 1, &sessions, &count); err != nil {
		return nil, fmt.Errorf("enumerate WTS sessions: %w", err)
	}
	if sessions == nil {
		return nil, nil
	}
	defer windows.WTSFreeMemory(uintptr(unsafe.Pointer(sessions)))

	console := windows.WTSGetActiveConsoleSessionId()
	result := make([]SessionCandidate, 0, count)
	for _, session := range unsafe.Slice(sessions, count) {
		if session.State != windows.WTSActive {
			continue
		}
		candidate, err := candidateFromSession(session.SessionID, console)
		if err == nil {
			result = append(result, candidate)
		}
	}
	return result, nil
}

func candidateFromSession(sessionID, console uint32) (SessionCandidate, error) {
	var token windows.Token
	if err := windows.WTSQueryUserToken(sessionID, &token); err != nil {
		return SessionCandidate{}, fmt.Errorf("query user token for session %d: %w", sessionID, err)
	}
	defer token.Close()

	user, err := token.GetTokenUser()
	if err != nil || user.User.Sid == nil {
		return SessionCandidate{}, fmt.Errorf("query token user for session %d: %w", sessionID, err)
	}
	logonSID, err := tokenLogonSID(token)
	if err != nil {
		return SessionCandidate{}, fmt.Errorf("query token logon SID for session %d: %w", sessionID, err)
	}
	facts := tokenFacts(token)
	if !facts.Usable {
		return SessionCandidate{}, fmt.Errorf("token for session %d is not usable", sessionID)
	}
	return SessionCandidate{
		SessionID: sessionID, Console: console != invalidSessionID && sessionID == console,
		UserSID: user.User.Sid.String(), LogonSID: logonSID, Token: facts,
	}, nil
}

func tokenFacts(token windows.Token) TokenFacts {
	facts := TokenFacts{Usable: true}
	if token.IsElevated() {
		facts.FullAdministrator = true
		return facts
	}
	// WTSQueryUserToken normally returns the user's standard/filtered primary
	// token. The launch adapter re-verifies integrity and can create a
	// restricted medium token if a legacy full-only token is encountered.
	facts.StandardOrFiltered = true
	linked, err := token.GetLinkedToken()
	if err == nil {
		facts.LinkedFullAvailable = linked.IsElevated()
		_ = linked.Close()
	}
	return facts
}

func tokenLogonSID(token windows.Token) (string, error) {
	groups, err := token.GetTokenGroups()
	if err != nil {
		return "", err
	}
	for _, group := range groups.AllGroups() {
		if group.Sid != nil && group.Attributes&windows.SE_GROUP_LOGON_ID == windows.SE_GROUP_LOGON_ID {
			return group.Sid.String(), nil
		}
	}
	return "", fmt.Errorf("token has no logon SID")
}