//go:build !windows

package agent

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
)

func TestTerminateRuntimeCommand_KillsRemainingProcessGroupAfterParentExit(t *testing.T) {
	t.Parallel()

	pidFile := fmt.Sprintf("%s/child.pid", t.TempDir())
	script := fmt.Sprintf("sh -c 'trap \"\" TERM; while :; do sleep 1; done' & echo $! > %s; wait", shellQuote(pidFile))

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	cmd := exec.CommandContext(ctx, "sh", "-c", script)
	configureRuntimeCommand(cmd)
	require.NoError(t, cmd.Start())

	childPID := waitForPIDFile(t, pidFile, 5*time.Second)
	defer killProcessIfAlive(childPID)

	terminateRuntimeCommand(cmd)
	cancel()

	require.Eventually(t, func() bool {
		return !processAlive(childPID)
	}, 10*time.Second, 100*time.Millisecond)

	err := cmd.Wait()
	if err != nil {
		var exitErr *exec.ExitError
		require.ErrorAs(t, err, &exitErr)
	}
}

func TestConfigureRuntimeCommand_ContextCancelDrainsTermIgnoringGroup(t *testing.T) {
	type scheduledKill struct {
		delay    time.Duration
		callback func()
	}
	killScheduled := make(chan scheduledKill, 1)

	pidFile := fmt.Sprintf("%s/child.pid", t.TempDir())
	script := fmt.Sprintf(
		"trap '' TERM; sh -c 'trap \"\" TERM; while :; do sleep 1; done' & echo $! > %s; wait",
		shellQuote(pidFile),
	)

	ctx, cancel := context.WithCancel(context.Background())
	cmd := exec.CommandContext(ctx, "sh", "-c", script)
	configureRuntimeCommand(cmd)
	cmd.Cancel = func() error {
		terminateRuntimeCommandWithKillScheduler(cmd, func(delay time.Duration, callback func()) *time.Timer {
			killScheduled <- scheduledKill{delay: delay, callback: callback}
			return nil
		})
		return nil
	}
	require.NoError(t, cmd.Start())

	parentPID := cmd.Process.Pid
	childPID := waitForPIDFile(t, pidFile, 5*time.Second)
	defer killProcessIfAlive(parentPID)
	defer killProcessIfAlive(childPID)

	waitDone := make(chan error, 1)
	go func() { waitDone <- cmd.Wait() }()

	cancel()
	var forceKill scheduledKill
	select {
	case forceKill = <-killScheduled:
	case <-time.After(10 * time.Second):
		t.Fatal("context cancellation did not schedule the process-group force kill")
	}
	require.Equal(t, processGroupKillGracePeriod, forceKill.delay)
	forceKill.callback()

	select {
	case err := <-waitDone:
		require.Error(t, err)
	case <-time.After(5 * time.Second):
		t.Fatal("cmd.Wait did not return after the scheduled process-group force kill")
	}
	require.Eventually(t, func() bool {
		return !processGroupRunning(parentPID) && !processAlive(parentPID) && !processAlive(childPID)
	}, 2*time.Second, 50*time.Millisecond)
}

func TestConfigureRuntimeCommand_CancelDrainsSetsidDescendant(t *testing.T) {
	tempDir := t.TempDir()
	helper := filepath.Join(tempDir, "setsid-child.py")
	require.NoError(t, os.WriteFile(helper, []byte(`
import os, signal, time
os.setsid()
signal.signal(signal.SIGTERM, signal.SIG_IGN)
os.write(3, f"{os.getpid()}\n".encode())
while True:
    time.sleep(1)
`), 0o600))
	script := fmt.Sprintf("trap '' TERM; python3 %s & wait", shellQuote(helper))

	readyReader, readyWriter, err := os.Pipe()
	require.NoError(t, err)
	t.Cleanup(func() { _ = readyReader.Close() })
	t.Cleanup(func() { _ = readyWriter.Close() })
	// The inherited pipe is both a readiness signal and a lifecycle signal:
	// EOF proves that the TERM-ignoring escaped descendant has been killed.
	ready := bufio.NewReader(readyReader)

	type scheduledKill struct {
		delay    time.Duration
		callback func()
	}
	killScheduled := make(chan scheduledKill, 1)

	cmd := exec.CommandContext(context.Background(), "sh", "-c", script)
	cmd.ExtraFiles = []*os.File{readyWriter}
	configureRuntimeCommand(cmd)
	// The sibling test covers CommandContext's asynchronous context watcher.
	// Drive the configured Cancel function directly here and capture its delayed
	// force-kill callback so this test does not depend on the production grace
	// period or runner scheduling latency.
	cmd.Cancel = func() error {
		terminateRuntimeCommandWithKillScheduler(cmd, func(delay time.Duration, callback func()) *time.Timer {
			killScheduled <- scheduledKill{delay: delay, callback: callback}
			return nil
		})
		return nil
	}
	require.NoError(t, cmd.Start())
	require.NoError(t, readyWriter.Close())

	parentPID := cmd.Process.Pid
	childPIDText, err := ready.ReadString('\n')
	require.NoError(t, err)
	childPID, err := strconv.Atoi(strings.TrimSpace(childPIDText))
	require.NoError(t, err)
	require.Positive(t, childPID)
	defer killProcessIfAlive(parentPID)
	defer killProcessIfAlive(childPID)
	childPGID, err := syscall.Getpgid(childPID)
	require.NoError(t, err)
	require.NotEqual(t, parentPID, childPGID)

	waitDone := make(chan error, 1)
	go func() { waitDone <- cmd.Wait() }()
	descendantsClosed := make(chan error, 1)
	go func() {
		_, readErr := io.ReadAll(ready)
		descendantsClosed <- readErr
	}()

	watchdog := time.NewTimer(30 * time.Second)
	defer watchdog.Stop()
	require.NoError(t, cmd.Cancel())
	var forceKill scheduledKill
	select {
	case forceKill = <-killScheduled:
	case <-watchdog.C:
		t.Fatal("configured cancellation did not schedule the force-kill callback")
	}
	require.Equal(t, processGroupKillGracePeriod, forceKill.delay)
	forceKill.callback()

	select {
	case err := <-waitDone:
		require.Error(t, err)
	case <-watchdog.C:
		t.Fatal("cmd.Wait did not return after the force-kill callback")
	}
	select {
	case err := <-descendantsClosed:
		require.NoError(t, err)
	case <-watchdog.C:
		t.Fatal("runtime descendants did not close the inherited readiness pipe")
	}
}

func waitForPIDFile(t *testing.T, path string, timeout time.Duration) int {
	t.Helper()

	ticker := time.NewTicker(50 * time.Millisecond)
	defer ticker.Stop()
	deadline := time.NewTimer(timeout)
	defer deadline.Stop()

	for {
		data, err := os.ReadFile(path)
		if err == nil {
			pid, convErr := strconv.Atoi(strings.TrimSpace(string(data)))
			require.NoError(t, convErr)
			require.Positive(t, pid)
			return pid
		}
		if !os.IsNotExist(err) {
			require.NoError(t, err)
		}
		select {
		case <-ticker.C:
		case <-deadline.C:
			t.Fatalf("timed out waiting for child pid file %s", path)
		}
	}
}

func processAlive(pid int) bool {
	if pid <= 0 {
		return false
	}
	err := syscall.Kill(pid, syscall.Signal(0))
	if err != nil && err != syscall.EPERM {
		return false
	}
	// kill(2) succeeds for zombies, which may remain indefinitely when a
	// container's PID 1 does not reap orphaned grandchildren.
	state, ok := processStatus(pid)
	return !ok || !strings.HasPrefix(state, "Z")
}

func processGroupRunning(pgid int) bool {
	if pgid <= 0 {
		return false
	}
	out, err := exec.Command("ps", "-axo", "pgid=,stat=").Output()
	if err != nil {
		err = syscall.Kill(-pgid, syscall.Signal(0))
		return err == nil || err == syscall.EPERM
	}
	for _, line := range strings.Split(string(out), "\n") {
		fields := strings.Fields(line)
		if len(fields) < 2 {
			continue
		}
		processGroupID, err := strconv.Atoi(fields[0])
		if err == nil && processGroupID == pgid && !strings.HasPrefix(fields[1], "Z") {
			return true
		}
	}
	return false
}

func processStatus(pid int) (string, bool) {
	out, err := exec.Command("ps", "-o", "stat=", "-p", strconv.Itoa(pid)).Output()
	if err != nil {
		return "", false
	}
	state := strings.TrimSpace(string(out))
	return state, state != ""
}

func killProcessIfAlive(pid int) {
	if !processAlive(pid) {
		return
	}
	_ = syscall.Kill(pid, syscall.SIGKILL)
}

func shellQuote(value string) string {
	return "'" + strings.ReplaceAll(value, "'", "'\"'\"'") + "'"
}