//go:build windows
// +build windows

/*
** Job Arranger for ZABBIX
** Copyright (C) 2025 Daiwa Institute of Research Ltd. All Rights Reserved.
**
** This program is free software; you can redistribute it and/or modify
** it under the terms of the GNU General Public License as published by
** the Free Software Foundation; either version 2 of the License, or
** (at your option) any later version.
**
** This program is distributed in the hope that it will be useful,
** but WITHOUT ANY WARRANTY; without even the implied warranty of
** MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
** GNU General Public License for more details.
**
** You should have received a copy of the GNU General Public License
** along with this program; if not, write to the Free Software
** Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
**/

package forker

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

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

const (
	LOGON32_LOGON_INTERACTIVE = 2
	LOGON32_PROVIDER_DEFAULT  = 0
)

// Win32 API functions
var (
	kernel32             = syscall.NewLazyDLL("kernel32.dll")
	procGetConsoleWindow = kernel32.NewProc("GetConsoleWindow")
	procAllocConsole     = kernel32.NewProc("AllocConsole")

	modUserenv            = windows.NewLazySystemDLL("userenv.dll")
	procLoadUserProfile   = modUserenv.NewProc("LoadUserProfileW")
	procUnloadUserProfile = modUserenv.NewProc("UnloadUserProfile")
	procCreateEnvBlock    = modUserenv.NewProc("CreateEnvironmentBlock")
	procDestroyEnvBlock   = modUserenv.NewProc("DestroyEnvironmentBlock")
)

type PROFILEINFO struct {
	dwSize        uint32
	dwFlags       uint32
	lpUserName    *uint16
	lpProfilePath *uint16
	lpDefaultPath *uint16
	lpServerName  *uint16
	lpPolicyPath  *uint16
	hProfile      windows.Handle
}

func getConsoleWindow() uintptr {
	ret, _, _ := procGetConsoleWindow.Call()
	return ret
}

func allocConsole() error {
	r, _, err := procAllocConsole.Call()
	if r == 0 {
		return err
	}
	return nil
}

// Convert Windows handle to *os.File
func handleToFile(handle windows.Handle, name string) *os.File {
	fd := int(handle)

	file := os.NewFile(uintptr(fd), name)

	return file
}

func (p *Forker) StartNewProcess() (*exec.Cmd, error) {
	if err := validateExecutable(p.data.ExecPath); err != nil {
		return nil, fmt.Errorf("%w for '%s': %w", ErrExecCheck, p.data.ExecPath, err)
	}

	var stdOut, stdErr *os.File
	// Default stdout and stderr targets
	stdOut = os.Stdout
	stdErr = os.Stderr

	// Create output files [stdout, stderr, ret]
	if p.data.StdoutPath != "" {
		stdoutFile, err := os.Create(p.data.StdoutPath)
		if err != nil {
			return nil, fmt.Errorf("failed to create stdout file: %w", err)
		}
		stdOut = stdoutFile
		defer stdoutFile.Close()
	}

	if p.data.StderrPath != "" {
		stderrFile, err := os.Create(p.data.StderrPath)
		if err != nil {
			return nil, fmt.Errorf("failed to create stderr file: %w", err)
		}
		stdErr = stderrFile
		defer stderrFile.Close()
	}

	if p.data.RetPath != "" {
		retPath := p.data.RetPath

		retFile, err := os.Create(retPath)
		if err != nil {
			return nil, fmt.Errorf("failed to create ret file: %w", err)
		}
		retFile.Close()
	}

	// Escape each parameter for Windows command line
	var cmdParts []string
	cmdParts = append(cmdParts, shellEscapeWindows(p.data.ExecPath))
	for _, param := range p.data.ExecParams {
		cmdParts = append(cmdParts, shellEscapeWindows(param))
	}

	fullCmd := strings.Join(cmdParts, " ")

	// Use SysProcAttr.CmdLine for exact command line control
	var cmd *exec.Cmd
	var cmdLine string
	if p.data.DirectExec {
		cmd = exec.Command(p.data.ExecPath, p.data.ExecParams...)
	} else {
		cmd = exec.Command("cmd.exe")
		cmdLine = fmt.Sprintf(`/q /k "%s" & exit !errorlevel!`, fullCmd)
	}

	cmd.SysProcAttr = &syscall.SysProcAttr{
		CmdLine:       cmdLine,
		CreationFlags: 0,
	}

	// Set standard handles
	cmd.Stdout = stdOut
	cmd.Stderr = stdErr

	if p.data.UsePipeStdin {
		// setup stdinpipe
		stdinPipe, err := cmd.StdinPipe()
		if err != nil {
			return nil, fmt.Errorf("failed to get stdin pipe: %w", err)
		}

		p.data.StdinPipe = &stdinPipe
	}

	if p.data.UseConsoleStdin {
		if getConsoleWindow() == 0 {
			if err := allocConsole(); err != nil {
				return nil, fmt.Errorf("allocConsole failed: %w", err)
			}
		}

		// Create a real usable console input handle
		hConsoleIn, err := windows.CreateFile(
			windows.StringToUTF16Ptr("CONIN$"),
			windows.GENERIC_READ|windows.GENERIC_WRITE,
			windows.FILE_SHARE_READ|windows.FILE_SHARE_WRITE,
			nil,
			windows.OPEN_EXISTING,
			0,
			0,
		)
		if err != nil {
			return nil, fmt.Errorf("CreateFile(CONIN$) failed: %w", err)
		}
		windows.SetHandleInformation(hConsoleIn, windows.HANDLE_FLAG_INHERIT, windows.HANDLE_FLAG_INHERIT)

		consoleIn := handleToFile(hConsoleIn, "input_console")
		cmd.Stdin = consoleIn
	}

	// accept extra files
	cmd.ExtraFiles = p.data.ExtraFiles

	cmd.Env = os.Environ()
	cmd.Env = append(cmd.Env, p.data.Env...)

	// Check for logon user
	if p.data.Username != "" {
		if p.data.Domain == "" {
			// "." for local machine
			p.data.Domain = "."
		}

		// First logon the user
		hToken, err := logonUser(p.data.Username, p.data.Domain, p.data.Password)
		if err != nil {
			return nil, fmt.Errorf("failed to logon user: %w", err)
		}
		defer hToken.Close()

		// Load profile
		profile := PROFILEINFO{
			dwSize:     uint32(unsafe.Sizeof(PROFILEINFO{})),
			lpUserName: windows.StringToUTF16Ptr(p.data.Username),
		}

		ret, _, err := procLoadUserProfile.Call(
			uintptr(hToken),
			uintptr(unsafe.Pointer(&profile)),
		)

		if ret == 0 {
			return nil, fmt.Errorf("LoadUserProfileW failed: %w", err)
		}

		defer procUnloadUserProfile.Call(uintptr(hToken), uintptr(profile.hProfile))

		// Create environment block
		var envBlock uintptr
		ret, _, err = procCreateEnvBlock.Call(
			uintptr(unsafe.Pointer(&envBlock)),
			uintptr(hToken),
			uintptr(1), // TRUE: inherit from current process
		)
		if ret == 0 {
			return nil, fmt.Errorf("CreateEnvironmentBlock failed: %w", err)
		}
		defer procDestroyEnvBlock.Call(envBlock)

		env := envBlockToSlice(envBlock)
		cmd.Env = append(cmd.Env, env...)

		// Set the token
		cmd.SysProcAttr.Token = syscall.Token(hToken)
	}

	// Start the command
	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("failed to start process: %w", err)
	}

	// Make child dies when parent exits by using job object
	if !p.data.Detached {
		// Create Job Object
		job, err := windows.CreateJobObject(nil, nil)
		if err != nil {
			return nil, fmt.Errorf("failed to create job object: %w", err)
		}

		var info windows.JOBOBJECT_EXTENDED_LIMIT_INFORMATION
		info.BasicLimitInformation.LimitFlags = windows.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE
		size := uint32(unsafe.Sizeof(info))
		_, err = windows.SetInformationJobObject(job, windows.JobObjectExtendedLimitInformation, uintptr(unsafe.Pointer(&info)), size)
		if err != nil {
			return nil, fmt.Errorf("failed to set information job object: %w", err)
		}

		// Get PROCESS_HANDLE, not PID
		procHandle, err := windows.OpenProcess(windows.PROCESS_ALL_ACCESS, false, uint32(cmd.Process.Pid))
		if err != nil {
			return nil, fmt.Errorf("failed to open process: %w", err)
		}
		defer windows.CloseHandle(procHandle)

		err = windows.AssignProcessToJobObject(
			job, procHandle,
		)
		if err != nil {
			return nil, fmt.Errorf("failed to assign process to job object: %w", err)
		}

	}

	return cmd, nil
}

func (p *Forker) WriteStdin(bytes []byte) error {
	if p.data.StdinPipe != nil {
		_, err := (*p.data.StdinPipe).Write(bytes)
		(*p.data.StdinPipe).Close()
		return err
	}

	return nil
}

func logonUser(username, domain, password string) (windows.Token, error) {
	advapi32 := windows.NewLazySystemDLL("advapi32.dll")
	logonUserProc := advapi32.NewProc("LogonUserW")

	var token windows.Token

	// Convert strings to UTF16 pointers
	userPtr, err := windows.UTF16PtrFromString(username)
	if err != nil {
		return 0, err
	}
	username = strings.TrimSpace(username)
	domainPtr, err := windows.UTF16PtrFromString(domain)
	if err != nil {
		return 0, err
	}

	password = strings.TrimSpace(password)
	passPtr, err := windows.UTF16PtrFromString(password)
	if err != nil {
		return 0, err
	}

	// Call LogonUserW
	ret, _, err := logonUserProc.Call(
		uintptr(unsafe.Pointer(userPtr)),
		uintptr(unsafe.Pointer(domainPtr)),
		uintptr(unsafe.Pointer(passPtr)),
		uintptr(LOGON32_LOGON_INTERACTIVE),
		uintptr(LOGON32_PROVIDER_DEFAULT),
		uintptr(unsafe.Pointer(&token)),
	)

	if ret == 0 {
		return 0, fmt.Errorf("LogonUser failed: %v", err)
	}

	return token, nil
}

func envBlockToSlice(envBlock uintptr) []string {
	var env []string
	ptr := envBlock
	for {
		str := windows.UTF16PtrToString((*uint16)(unsafe.Pointer(ptr)))
		if len(str) == 0 {
			break
		}
		env = append(env, str)
		ptr = uintptr(unsafe.Pointer((*uint16)(unsafe.Pointer(ptr)))) + uintptr((len(str)+1)*2)
	}
	return env
}

func shellEscapeWindows(s string) string {
	if s == "" {
		return `""`
	}
	if strings.ContainsAny(s, " \t\"") {
		// Double-up quotes inside the string and wrap with quotes
		s = strings.ReplaceAll(s, `"`, `""`)
		return `"` + s + `"`
	}
	return s
}
