//go:build windows

package compose

import (
	"errors"
	"os/exec"
	"sync"
	"unsafe"

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

var ownedJobs = struct {
	sync.Mutex
	m map[*exec.Cmd]windows.Handle
}{m: make(map[*exec.Cmd]windows.Handle)}

// CREATE_SUSPENDED closes the registration race: no child code can create a
// descendant until the real process HANDLE has been assigned to the job.
func configureOwnedProcess(c *exec.Cmd) {
	c.SysProcAttr = &windows.SysProcAttr{CreationFlags: windows.CREATE_NEW_PROCESS_GROUP | windows.CREATE_SUSPENDED}
}

func registerOwnedProcess(c *exec.Cmd) error {
	if c.Process == nil {
		return errors.New("owned process has no process handle")
	}
	job, err := windows.CreateJobObject(nil, nil)
	if err != nil {
		return err
	}
	closeJob := true
	defer func() {
		if closeJob {
			windows.CloseHandle(job)
		}
	}()
	limits := windows.JOBOBJECT_EXTENDED_LIMIT_INFORMATION{}
	limits.BasicLimitInformation.LimitFlags = windows.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE
	if _, err = windows.SetInformationJobObject(job, windows.JobObjectExtendedLimitInformation, uintptr(unsafe.Pointer(&limits)), uint32(unsafe.Sizeof(limits))); err != nil {
		return err
	}
	// c.Process.Pid is used only to obtain a genuine process HANDLE; a PID is
	// never passed to AssignProcessToJobObject.
	process, err := windows.OpenProcess(windows.PROCESS_SET_QUOTA|windows.PROCESS_TERMINATE|windows.PROCESS_QUERY_LIMITED_INFORMATION, false, uint32(c.Process.Pid))
	if err != nil {
		return err
	}
	defer windows.CloseHandle(process)
	if err = windows.AssignProcessToJobObject(job, process); err != nil {
		return err
	}
	thread, err := suspendedPrimaryThread(uint32(c.Process.Pid))
	if err != nil {
		_ = windows.TerminateJobObject(job, 1)
		return err
	}
	_, resumeErr := windows.ResumeThread(thread)
	windows.CloseHandle(thread)
	if resumeErr != nil {
		_ = windows.TerminateJobObject(job, 1)
		return resumeErr
	}
	ownedJobs.Lock()
	ownedJobs.m[c] = job
	ownedJobs.Unlock()
	closeJob = false
	return nil
}

func suspendedPrimaryThread(pid uint32) (windows.Handle, error) {
	snapshot, err := windows.CreateToolhelp32Snapshot(windows.TH32CS_SNAPTHREAD, 0)
	if err != nil {
		return 0, err
	}
	defer windows.CloseHandle(snapshot)
	entry := windows.ThreadEntry32{Size: uint32(unsafe.Sizeof(windows.ThreadEntry32{}))}
	err = windows.Thread32First(snapshot, &entry)
	for err == nil {
		if entry.OwnerProcessID == pid {
			return windows.OpenThread(windows.THREAD_SUSPEND_RESUME, false, entry.ThreadID)
		}
		err = windows.Thread32Next(snapshot, &entry)
	}
	return 0, err
}

func releaseOwnedProcess(c *exec.Cmd) {
	ownedJobs.Lock()
	h := ownedJobs.m[c]
	delete(ownedJobs.m, c)
	ownedJobs.Unlock()
	if h != 0 {
		windows.CloseHandle(h)
	}
}
func terminateOwnedProcess(c *exec.Cmd) {
	ownedJobs.Lock()
	h := ownedJobs.m[c]
	ownedJobs.Unlock()
	if h != 0 {
		_ = windows.TerminateJobObject(h, 1)
	} else if c.Process != nil {
		_ = c.Process.Kill()
	}
}