//go:build windows

package subprocess

import (
	"errors"
	"fmt"
	"os/exec"
	"syscall"
	"unsafe"

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

type windowsProcessTree struct {
	job windows.Handle
}

func newOwnedProcessTree() (ownedProcessTree, error) {
	return &windowsProcessTree{}, nil
}

func (tree *windowsProcessTree) Start(cmd *exec.Cmd) error {
	job, err := windows.CreateJobObject(nil, nil)
	if err != nil {
		return fmt.Errorf("create job object: %w", err)
	}

	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 {
		_ = windows.CloseHandle(job)
		return fmt.Errorf("configure job object: %w", err)
	}

	cmd.SysProcAttr = &syscall.SysProcAttr{CreationFlags: windows.CREATE_SUSPENDED}
	if err := cmd.Start(); err != nil {
		_ = windows.CloseHandle(job)
		return err
	}

	process, err := windows.OpenProcess(windows.PROCESS_SET_QUOTA|windows.PROCESS_TERMINATE, false, uint32(cmd.Process.Pid))
	if err == nil {
		err = windows.AssignProcessToJobObject(job, process)
		_ = windows.CloseHandle(process)
	}
	if err == nil {
		err = resumeInitialThread(uint32(cmd.Process.Pid))
	}
	if err != nil {
		killErr := cmd.Process.Kill()
		waitErr := cmd.Wait()
		_ = windows.CloseHandle(job)
		return joinErrors(fmt.Errorf("assign process to job object: %w", err), killErr, waitErr)
	}

	tree.job = job
	return nil
}

func resumeInitialThread(processID uint32) error {
	snapshot, err := windows.CreateToolhelp32Snapshot(windows.TH32CS_SNAPTHREAD, 0)
	if err != nil {
		return fmt.Errorf("snapshot initial thread: %w", err)
	}
	defer func() { _ = windows.CloseHandle(snapshot) }()

	entry := windows.ThreadEntry32{Size: uint32(unsafe.Sizeof(windows.ThreadEntry32{}))}
	if err := windows.Thread32First(snapshot, &entry); err != nil {
		return fmt.Errorf("find initial thread: %w", err)
	}
	for {
		if entry.OwnerProcessID != processID {
			// Keep enumerating until the suspended process's only initial thread
			// is found.
		} else {
			thread, openErr := windows.OpenThread(windows.THREAD_SUSPEND_RESUME, false, entry.ThreadID)
			if openErr != nil {
				return fmt.Errorf("open initial thread: %w", openErr)
			}
			defer func() { _ = windows.CloseHandle(thread) }()
			if _, resumeErr := windows.ResumeThread(thread); resumeErr != nil {
				return fmt.Errorf("resume initial thread: %w", resumeErr)
			}
			return nil
		}

		if err := windows.Thread32Next(snapshot, &entry); err != nil {
			if errors.Is(err, windows.ERROR_NO_MORE_FILES) {
				break
			}
			return fmt.Errorf("find initial thread: %w", err)
		}
	}
	return fmt.Errorf("find initial thread: no thread found for process %d", processID)
}

func (tree *windowsProcessTree) TerminateGracefully() error {
	// Windows jobs have no portable graceful signal. Terminating the owned job
	// is the safe fallback and prevents a descendant from escaping cleanup.
	return tree.terminate()
}

func (tree *windowsProcessTree) TerminateForcefully() error {
	return tree.terminate()
}

func (tree *windowsProcessTree) Dispose() error {
	if tree.job == 0 {
		return nil
	}
	err := windows.CloseHandle(tree.job)
	tree.job = 0
	return err
}

func (tree *windowsProcessTree) terminate() error {
	if tree.job == 0 {
		return nil
	}
	return windows.TerminateJobObject(tree.job, 1)
}