//go:build !windows

package agent

import (
	"context"
	"errors"
	"log/slog"
	"os"
	"os/exec"
	"syscall"
	"time"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/runtimeproc"
)

const processGroupKillGracePeriod = 2 * time.Second
const processGroupWaitDelay = processGroupKillGracePeriod + 500*time.Millisecond

type processGroupKillScheduler func(time.Duration, func()) *time.Timer

// configureRuntimeCommand sets up process group creation so the runtime and
// all its children can be signalled as a group, and routes ctx cancellation
// through terminateRuntimeCommand so the whole group shuts down gracefully
// instead of Go's default SIGKILL to just the binary PID.
func configureRuntimeCommand(cmd *exec.Cmd) {
	if cmd == nil {
		return
	}
	if cmd.SysProcAttr == nil {
		cmd.SysProcAttr = &syscall.SysProcAttr{}
	}
	cmd.SysProcAttr.Setpgid = true
	cmd.WaitDelay = processGroupWaitDelay
	cmd.Cancel = func() error {
		terminateRuntimeCommand(cmd)
		return nil
	}
}

// terminateRuntimeCommand sends SIGTERM to the entire process group.
// If the group is still alive after the grace period, SIGKILL is sent.
func terminateRuntimeCommand(cmd *exec.Cmd) {
	terminateRuntimeCommandWithKillScheduler(cmd, time.AfterFunc)
}

func terminateRuntimeCommandWithKillScheduler(cmd *exec.Cmd, scheduleKill processGroupKillScheduler) {
	if cmd == nil || cmd.Process == nil {
		return
	}

	descendants := snapshotRuntimeDescendants(cmd.Process.Pid)
	if err := signalRuntimeCommand(cmd, syscall.SIGTERM); err != nil {
		return
	}
	signalRuntimeDescendants(descendants, syscall.SIGTERM)

	scheduleKill(processGroupKillGracePeriod, func() {
		if runtimeProcessGroupAlive(cmd) {
			if err := signalRuntimeCommand(cmd, syscall.SIGKILL); err != nil {
				slog.Warn("failed to force-kill runtime process group",
					"pid", cmd.Process.Pid,
					"error", err,
				)
			}
		}
		signalRuntimeDescendants(descendants, syscall.SIGKILL)
	})
}

func snapshotRuntimeDescendants(rootPID int) []runtimeproc.ProcessInfo {
	processes, err := runtimeproc.NewPlatformManager().ListProcesses(context.Background())
	if err != nil {
		slog.Warn("failed to snapshot runtime descendants", "pid", rootPID, "error", err)
		return nil
	}
	children := make(map[int][]runtimeproc.ProcessInfo)
	for _, process := range processes {
		if process.PID > 0 && process.PPID > 0 {
			children[process.PPID] = append(children[process.PPID], process)
		}
	}
	var descendants []runtimeproc.ProcessInfo
	var visit func(int)
	visit = func(parentPID int) {
		for _, child := range children[parentPID] {
			visit(child.PID)
			descendants = append(descendants, child)
		}
	}
	visit(rootPID)
	return descendants
}

func signalRuntimeDescendants(descendants []runtimeproc.ProcessInfo, signal syscall.Signal) {
	if len(descendants) == 0 {
		return
	}
	processes, err := runtimeproc.NewPlatformManager().ListProcesses(context.Background())
	if err != nil {
		slog.Warn("failed to verify runtime descendants", "error", err)
		return
	}
	live := make(map[int]runtimeproc.ProcessInfo, len(processes))
	for _, process := range processes {
		live[process.PID] = process
	}
	for _, expected := range descendants {
		if expected.PID <= 0 || expected.PID == os.Getpid() {
			continue
		}
		current, ok := live[expected.PID]
		if !ok || !sameRuntimeProcessIdentity(expected, current) {
			continue
		}
		if err := syscall.Kill(expected.PID, signal); err != nil && !errors.Is(err, syscall.ESRCH) {
			slog.Warn("failed to signal escaped runtime descendant",
				"pid", expected.PID,
				"signal", signal,
				"error", err,
			)
		}
	}
}

func sameRuntimeProcessIdentity(expected, current runtimeproc.ProcessInfo) bool {
	if expected.PID != current.PID {
		return false
	}
	if !expected.StartedAt.IsZero() && !current.StartedAt.IsZero() {
		return expected.StartedAt.Equal(current.StartedAt)
	}
	if expected.Exe != "" && current.Exe != "" {
		return expected.Exe == current.Exe && expected.Cmdline == current.Cmdline
	}
	return expected.Cmdline != "" && expected.Cmdline == current.Cmdline
}

// signalRuntimeCommand sends a signal to the process group first, falling
// back to the single process if group signalling fails.
func signalRuntimeCommand(cmd *exec.Cmd, signal syscall.Signal) error {
	if cmd == nil || cmd.Process == nil {
		return nil
	}

	if err := syscall.Kill(-cmd.Process.Pid, signal); err == nil || errors.Is(err, syscall.ESRCH) {
		return nil
	}
	if err := cmd.Process.Signal(signal); err == nil || errors.Is(err, os.ErrProcessDone) {
		return nil
	} else {
		return err
	}
}

// runtimeCommandExited reports whether the runtime process has exited.
func runtimeCommandExited(cmd *exec.Cmd) bool {
	if cmd == nil || cmd.Process == nil {
		return true
	}
	if state := cmd.ProcessState; state != nil && state.Exited() {
		return true
	}

	err := cmd.Process.Signal(syscall.Signal(0))
	return errors.Is(err, os.ErrProcessDone) || errors.Is(err, syscall.ESRCH)
}

// runtimeProcessGroupAlive reports whether the runtime process group still has
// any live members, even if the original parent process has already exited.
func runtimeProcessGroupAlive(cmd *exec.Cmd) bool {
	if cmd == nil || cmd.Process == nil {
		return false
	}

	err := syscall.Kill(-cmd.Process.Pid, syscall.Signal(0))
	return err == nil || errors.Is(err, syscall.EPERM)
}