package workspace

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

	iop "iop/proto/gen/iop"
)

func TestMain(m *testing.M) {
	if handled, exitCode := RunCommandShim(os.Args); handled {
		os.Exit(exitCode)
	}
	os.Exit(m.Run())
}

func TestWorkspaceCommandHelperProcess(t *testing.T) {
	mode := os.Getenv("IOP_WORKSPACE_HELPER")
	if mode == "" {
		return
	}
	switch mode {
	case "success":
		_, _ = fmt.Fprint(os.Stdout, "command-stdout")
		_, _ = fmt.Fprint(os.Stderr, "command-stderr")
	case "nonzero":
		os.Exit(7)
	case "environment":
		_, _ = fmt.Fprintf(os.Stdout, "%s|%s", os.Getenv("IOP_TEST_VALUE"), os.Getenv("IOP_AMBIENT_SECRET"))
	case "output":
		_, _ = fmt.Fprint(os.Stdout, strings.Repeat("o", 128<<10))
		_, _ = fmt.Fprint(os.Stderr, strings.Repeat("e", 128<<10))
	case "cwd":
		identity, err := os.ReadFile("identity.txt")
		if err != nil {
			os.Exit(8)
		}
		cwd, err := os.Getwd()
		if err != nil {
			os.Exit(9)
		}
		_, _ = fmt.Fprintf(os.Stdout, "%s|%s", identity, cwd)
	case "block":
		if err := os.WriteFile(os.Getenv("IOP_START_FILE"), []byte("started"), 0o600); err != nil {
			os.Exit(10)
		}
		select {}
	case "group":
		cmd := exec.Command(os.Args[0], "-test.run=^TestWorkspaceCommandGrandchild$")
		cmd.Env = []string{"IOP_WORKSPACE_GRANDCHILD=1"}
		if err := cmd.Start(); err != nil {
			os.Exit(11)
		}
		if err := os.WriteFile(os.Getenv("IOP_CHILD_PID_FILE"), []byte(strconv.Itoa(cmd.Process.Pid)), 0o600); err != nil {
			_ = cmd.Process.Kill()
			os.Exit(12)
		}
		select {}
	case "sentinel":
		if err := os.WriteFile(os.Getenv("IOP_SENTINEL_FILE"), []byte("target-started"), 0o600); err != nil {
			os.Exit(13)
		}
	default:
		os.Exit(14)
	}
	os.Exit(0)
}

func TestWorkspaceCommandGrandchild(t *testing.T) {
	if os.Getenv("IOP_WORKSPACE_GRANDCHILD") == "" {
		return
	}
	select {}
}

func newCommandRuntime(t *testing.T, root string, outputLimit int64) *Runtime {
	t.Helper()
	executable, err := os.Executable()
	if err != nil {
		t.Fatal(err)
	}
	runtime, err := NewRuntime([]*iop.WorkspaceConfig{{
		Ref: "workspace-command", Platform: "darwin", Root: root,
		Operations: []iop.WorkspaceOperation{iop.WorkspaceOperation_WORKSPACE_OPERATION_COMMAND},
		Commands: []*iop.WorkspaceCommandConfig{{
			Id: "helper", Executable: executable,
			Args: []string{"-test.run=^TestWorkspaceCommandHelperProcess$"},
		}},
		EnvironmentAllowlist: []string{
			"IOP_WORKSPACE_HELPER", "IOP_TEST_VALUE", "IOP_START_FILE",
			"IOP_CHILD_PID_FILE", "IOP_SENTINEL_FILE",
		},
		MaxOutputBytes: outputLimit, MaxCommandTimeoutMs: 3000,
	}}, "darwin", nil)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = runtime.Close() })
	return runtime
}

func openCommandRequest(t *testing.T, runtime *Runtime, requestID string, outputLimit int64) {
	t.Helper()
	_, err := runtime.Open(RequestAuthority{
		RequestID: requestID, WorkspaceRef: "workspace-command",
		Operations: []iop.WorkspaceOperation{iop.WorkspaceOperation_WORKSPACE_OPERATION_COMMAND},
		CommandIDs: []string{"helper"}, MaxOutputBytes: outputLimit, MaxCommandTimeoutMS: 3000,
	})
	if err != nil {
		t.Fatal(err)
	}
}

func commandInput(requestID, toolCallID, mode string) CommandInput {
	return CommandInput{
		RequestID: requestID, ToolCallID: toolCallID, CommandID: "helper", TimeoutMS: 2000,
		Environment: map[string]string{"IOP_WORKSPACE_HELPER": mode},
	}
}

func TestCommandExecutorSuccessFailureAndEnvironment(t *testing.T) {
	t.Setenv("IOP_AMBIENT_SECRET", "must-not-be-inherited")
	runtime := newCommandRuntime(t, t.TempDir(), 256)
	openCommandRequest(t, runtime, "request-success", 256)

	success := runtime.ExecuteCommand(context.Background(), commandInput("request-success", "tool-success", "success"))
	if success.Status != iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS || success.ExitCode != 0 || string(success.Stdout) != "command-stdout" || string(success.Stderr) != "command-stderr" {
		t.Fatalf("success = %+v", success)
	}
	nonzero := runtime.ExecuteCommand(context.Background(), commandInput("request-success", "tool-nonzero", "nonzero"))
	if nonzero.Status != iop.WorkspaceStatus_WORKSPACE_STATUS_ERROR || nonzero.Code != iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_INTERNAL || nonzero.ExitCode != 7 {
		t.Fatalf("nonzero = %+v", nonzero)
	}
	environmentInput := commandInput("request-success", "tool-environment", "environment")
	environmentInput.Environment["IOP_TEST_VALUE"] = "approved"
	environment := runtime.ExecuteCommand(context.Background(), environmentInput)
	if environment.Status != iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS || string(environment.Stdout) != "approved|" {
		t.Fatalf("environment = %+v", environment)
	}

	unknown := commandInput("request-success", "tool-unknown", "success")
	unknown.CommandID = "not-approved"
	if result := runtime.ExecuteCommand(context.Background(), unknown); result.Code != iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_INVALID_REQUEST {
		t.Fatalf("unknown command = %+v", result)
	}
	unapprovedEnvironment := commandInput("request-success", "tool-env-denied", "success")
	unapprovedEnvironment.Environment["HOME"] = "/sensitive"
	if result := runtime.ExecuteCommand(context.Background(), unapprovedEnvironment); result.Code != iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_INVALID_REQUEST {
		t.Fatalf("unapproved environment = %+v", result)
	}
	oversizedTimeout := commandInput("request-success", "tool-timeout-denied", "success")
	oversizedTimeout.TimeoutMS = 3001
	if result := runtime.ExecuteCommand(context.Background(), oversizedTimeout); result.Code != iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_INVALID_REQUEST {
		t.Fatalf("oversized timeout = %+v", result)
	}
}

func TestCommandExecutorSharedOutputBound(t *testing.T) {
	runtime := newCommandRuntime(t, t.TempDir(), 64)
	openCommandRequest(t, runtime, "request-output", 64)
	result := runtime.ExecuteCommand(context.Background(), commandInput("request-output", "tool-output", "output"))
	if result.Status != iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS || !result.Truncated || len(result.Stdout)+len(result.Stderr) > 64 {
		t.Fatalf("bounded output = %+v stdout=%d stderr=%d", result, len(result.Stdout), len(result.Stderr))
	}
}

func TestCommandExecutorTimeoutAndContextCancel(t *testing.T) {
	root := t.TempDir()
	runtime := newCommandRuntime(t, root, 64)
	openCommandRequest(t, runtime, "request-timeout", 64)
	timeout := commandInput("request-timeout", "tool-timeout", "block")
	timeout.TimeoutMS = 50
	timeout.Environment["IOP_START_FILE"] = filepath.Join(root, "timeout-started")
	result := runtime.ExecuteCommand(context.Background(), timeout)
	if result.Status != iop.WorkspaceStatus_WORKSPACE_STATUS_TIMEOUT || result.Code != iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_TIMEOUT || result.ExitCode != -1 {
		t.Fatalf("timeout = %+v", result)
	}

	// Pre-cancelled context fast-path assertion.
	preCtx, preCancel := context.WithCancel(context.Background())
	preCancel()
	preInput := commandInput("request-timeout", "tool-pre-cancel", "success")
	result = runtime.ExecuteCommand(preCtx, preInput)
	if result.Status != iop.WorkspaceStatus_WORKSPACE_STATUS_CANCELLED || result.Code != iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_CANCELLED {
		t.Fatalf("pre-cancelled context = %+v", result)
	}

	// Live active context cancellation and process group termination assertion.
	pidFile := filepath.Join(root, "active-child.pid")
	activeCtx, activeCancel := context.WithCancel(context.Background())
	activeInput := commandInput("request-timeout", "tool-active-context", "group")
	activeInput.Environment["IOP_CHILD_PID_FILE"] = pidFile
	resultCh := make(chan Result, 1)
	go func() {
		resultCh <- runtime.ExecuteCommand(activeCtx, activeInput)
	}()
	waitForFile(t, pidFile)
	pidBytes, err := os.ReadFile(pidFile)
	if err != nil {
		t.Fatal(err)
	}
	pid, err := strconv.Atoi(strings.TrimSpace(string(pidBytes)))
	if err != nil {
		t.Fatal(err)
	}
	activeCancel()
	activeResult := <-resultCh
	if activeResult.Status != iop.WorkspaceStatus_WORKSPACE_STATUS_CANCELLED || activeResult.Code != iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_CANCELLED || activeResult.ExitCode != -1 {
		t.Fatalf("live active context cancel result = %+v", activeResult)
	}
	deadline := time.Now().Add(2 * time.Second)
	for processExists(pid) && time.Now().Before(deadline) {
		time.Sleep(10 * time.Millisecond)
	}
	if processExists(pid) {
		t.Fatalf("grandchild process %d survived live active context cancellation", pid)
	}
}

func TestCommandExecutorExplicitCancelAndRequestIsolation(t *testing.T) {
	root := t.TempDir()
	runtime := newCommandRuntime(t, root, 64)
	openCommandRequest(t, runtime, "request-a", 64)
	openCommandRequest(t, runtime, "request-b", 64)
	resultA := make(chan Result, 1)
	resultB := make(chan Result, 1)
	inputA := commandInput("request-a", "tool-shared", "block")
	inputA.Environment["IOP_START_FILE"] = filepath.Join(root, "started-a")
	inputB := commandInput("request-b", "tool-shared", "block")
	inputB.Environment["IOP_START_FILE"] = filepath.Join(root, "started-b")
	go func() { resultA <- runtime.ExecuteCommand(context.Background(), inputA) }()
	go func() { resultB <- runtime.ExecuteCommand(context.Background(), inputB) }()
	waitForFile(t, inputA.Environment["IOP_START_FILE"])
	waitForFile(t, inputB.Environment["IOP_START_FILE"])

	if wrong := runtime.Cancel("request-a", "tool-other"); wrong.Code != iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_NOT_FOUND {
		t.Fatalf("wrong cancel = %+v", wrong)
	}
	if cancelled := runtime.Cancel("request-a", "tool-shared"); cancelled.Status != iop.WorkspaceStatus_WORKSPACE_STATUS_CANCELLED {
		t.Fatalf("cancel a = %+v", cancelled)
	}
	if result := <-resultA; result.Status != iop.WorkspaceStatus_WORKSPACE_STATUS_CANCELLED {
		t.Fatalf("result a = %+v", result)
	}
	select {
	case result := <-resultB:
		t.Fatalf("cross-request cancel stopped b: %+v", result)
	case <-time.After(50 * time.Millisecond):
	}
	if cancelled := runtime.Cancel("request-b", "tool-shared"); cancelled.Status != iop.WorkspaceStatus_WORKSPACE_STATUS_CANCELLED {
		t.Fatalf("cancel b = %+v", cancelled)
	}
	if result := <-resultB; result.Status != iop.WorkspaceStatus_WORKSPACE_STATUS_CANCELLED {
		t.Fatalf("result b = %+v", result)
	}
}

func TestCommandExecutorCancelKillsProcessGroup(t *testing.T) {
	root := t.TempDir()
	runtime := newCommandRuntime(t, root, 64)
	openCommandRequest(t, runtime, "request-group", 64)
	pidFile := filepath.Join(root, "child.pid")
	input := commandInput("request-group", "tool-group", "group")
	input.Environment["IOP_CHILD_PID_FILE"] = pidFile
	resultChannel := make(chan Result, 1)
	go func() { resultChannel <- runtime.ExecuteCommand(context.Background(), input) }()
	waitForFile(t, pidFile)
	pidBytes, err := os.ReadFile(pidFile)
	if err != nil {
		t.Fatal(err)
	}
	pid, err := strconv.Atoi(string(pidBytes))
	if err != nil {
		t.Fatal(err)
	}
	if cancelled := runtime.Cancel("request-group", "tool-group"); cancelled.Status != iop.WorkspaceStatus_WORKSPACE_STATUS_CANCELLED {
		t.Fatalf("cancel = %+v", cancelled)
	}
	if result := <-resultChannel; result.Status != iop.WorkspaceStatus_WORKSPACE_STATUS_CANCELLED {
		t.Fatalf("result = %+v", result)
	}
	deadline := time.Now().Add(2 * time.Second)
	for processExists(pid) && time.Now().Before(deadline) {
		time.Sleep(10 * time.Millisecond)
	}
	if processExists(pid) {
		t.Fatalf("grandchild process %d survived group cancellation", pid)
	}
}

func TestCommandExecutorUsesOpenedRootAfterRenameReplacement(t *testing.T) {
	parent := t.TempDir()
	root := filepath.Join(parent, "workspace")
	if err := os.Mkdir(root, 0o700); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(root, "identity.txt"), []byte("original"), 0o600); err != nil {
		t.Fatal(err)
	}
	runtime := newCommandRuntime(t, root, 256)
	openCommandRequest(t, runtime, "request-cwd", 256)
	renamed := filepath.Join(parent, "workspace-renamed")
	if err := os.Rename(root, renamed); err != nil {
		t.Fatal(err)
	}
	foreign := filepath.Join(parent, "foreign")
	if err := os.Mkdir(foreign, 0o700); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(foreign, "identity.txt"), []byte("foreign"), 0o600); err != nil {
		t.Fatal(err)
	}
	if err := os.Symlink(foreign, root); err != nil {
		t.Fatal(err)
	}
	result := runtime.ExecuteCommand(context.Background(), commandInput("request-cwd", "tool-cwd", "cwd"))
	if result.Status != iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS || !strings.HasPrefix(string(result.Stdout), "original|") || strings.Contains(string(result.Stdout), "foreign") {
		t.Fatalf("cwd result = %+v", result)
	}
}

func TestCommandExecutorRejectsCorruptRootIdentityBeforeTarget(t *testing.T) {
	rootPath := t.TempDir()
	directory, err := os.Open(rootPath)
	if err != nil {
		t.Fatal(err)
	}
	defer directory.Close()
	info, err := directory.Stat()
	if err != nil {
		t.Fatal(err)
	}
	device, inode, ok := fileIdentity(info)
	if !ok {
		t.Fatal("root identity unavailable")
	}
	executable, err := os.Executable()
	if err != nil {
		t.Fatal(err)
	}
	sentinel := filepath.Join(rootPath, "target-started")
	output := newCommandOutput(64)
	process, err := startCommandProcess(commandLaunchRecord{
		Version: commandLaunchVersion, Executable: executable,
		Args:        []string{"-test.run=^TestWorkspaceCommandHelperProcess$"},
		Environment: []string{"IOP_SENTINEL_FILE=" + sentinel, "IOP_WORKSPACE_HELPER=sentinel"},
		Device:      device, Inode: inode + 1,
	}, directory, output)
	if err != nil {
		t.Fatal(err)
	}
	result := awaitCommand(context.Background(), newCommandExecution(), process, time.Second, time.Now(), output)
	if result.Status != iop.WorkspaceStatus_WORKSPACE_STATUS_ERROR || result.Code != iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_INTERNAL {
		t.Fatalf("corrupt identity result = %+v", result)
	}
	if _, err := os.Stat(sentinel); !errors.Is(err, os.ErrNotExist) {
		t.Fatalf("target sentinel exists or stat failed unexpectedly: %v", err)
	}
}

func waitForFile(t *testing.T, path string) {
	t.Helper()
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		if _, err := os.Stat(path); err == nil {
			return
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatalf("timed out waiting for %s", filepath.Base(path))
}

func processExists(pid int) bool {
	err := syscall.Kill(pid, 0)
	return err == nil || !errors.Is(err, syscall.ESRCH)
}