package tools

import (
	"context"
	"fmt"
	"regexp"
	"strings"
	"time"
	"xbot/llm"

	log "xbot/logger"
)

const defaultShellTimeout = 120 * time.Second

// ShellTool 执行命令工具
type ShellTool struct{}

func (t *ShellTool) Name() string {
	return "Shell"
}

func (t *ShellTool) Description() string {
	return `Execute a command and return its output.
The command will be executed in the agent's working directory.
IMPORTANT: Commands are executed non-interactively with a timeout. Do NOT run interactive commands (e.g. vim, top, htop) or commands that require manual input. For commands that might prompt for input, use non-interactive flags (e.g. "apt-get -y", "yes |", "ssh -o BatchMode=yes"). For sudo, use NOPASSWD or "echo password | sudo -S".

PROCESS CLEANUP: Non-background commands are killed (including all child processes) when they return. Do NOT use nohup, disown, or trailing & — they create orphaned processes that waste resources and cause confusion. If a command needs to outlive the tool call, use "background": true instead.

BACKGROUND MODE: Set "background": true to run long-running commands (dev servers, build processes) without blocking. Returns a task ID immediately. The agent continues working while the command runs in the background. When the command finishes, its output is automatically injected into the conversation. To check progress, use task_status — but do NOT poll it repeatedly. If status is "running", do other work or sleep 3+ seconds before checking again.

AUTO-BACKGROUND: If a command times out, it is automatically converted to a background task so no work is lost. The agent receives the task ID and can continue. Same polling rule applies: do NOT call task_status in rapid succession.

Parameters (JSON):
  - command: string, the command to execute
  - timeout: number (optional), timeout in seconds (default: 120, max: 600)
  - background: boolean (optional), run in background mode

Environment Variables:
- Commands run in a login shell (detected from container's /etc/passwd), which automatically sources /etc/profile, ~/.bash_profile, ~/.bashrc, etc.`
}

func (t *ShellTool) Parameters() []llm.ToolParam {
	return []llm.ToolParam{
		{Name: "command", Type: "string", Description: "The command to execute", Required: true},
		{Name: "timeout", Type: "number", Description: "Timeout in seconds (default: 120, max: 600)", Required: false},
		{Name: "background", Type: "boolean", Description: "Run command in background (for long-running tasks like dev servers). Returns task ID immediately.", Required: false},
		{Name: "run_as", Type: "string", Description: "OS username to execute as. Requires permission control to be enabled. Only effective in none sandbox mode.", Required: false},
		{Name: "reason", Type: "string", Description: "Optional human-readable reason shown in approval requests when approval is required.", Required: false},
	}
}

func (t *ShellTool) Execute(toolCtx *ToolContext, input string) (*ToolResult, error) {
	params, err := parseToolArgs[struct {
		Command    string  `json:"command"`
		Timeout    float64 `json:"timeout"`
		Background bool    `json:"background"`
		RunAs      string  `json:"run_as"`
		Reason     string  `json:"reason"`
	}](input)
	if err != nil {
		return nil, err
	}

	if params.Command == "" {
		return nil, fmt.Errorf("command is required")
	}

	// When permission control is disabled, ignore any stale run_as/reason
	// the LLM might send from cached context.
	if toolCtx.Ctx == nil || !isPermControlActiveFromCtx(toolCtx.Ctx) {
		params.RunAs = ""
		params.Reason = ""
	}

	if err := validateRunAsReason(params.RunAs, params.Reason); err != nil {
		return nil, err
	}

	// 检测命令中的控制字符和 null bytes
	if strings.ContainsAny(params.Command, "\x00\x01\x02\x03\x04\x05\x06\x07\x08\x0b\x0c\x0e\x0f\x10\x11\x12\x13\x14\x15\x16\x17\x18\x19\x1a\x1b\x1c\x1d\x1e\x1f") {
		return nil, fmt.Errorf("command contains control characters (null bytes or other non-printable characters)")
	}

	// 安全预检：拦截危险命令
	// - run_as 模式下禁止任何形式的 sudo（用户切换由框架通过 cmdbuilder 处理）
	// - permission control 启用时，即使未设置 run_as，也禁止 LLM 直接使用 sudo
	if blocked, reason := checkDangerousCommand(toolCtx.Ctx, params.Command, params.RunAs != ""); blocked {
		return nil, fmt.Errorf("command blocked by safety check: %s", reason)
	}

	const maxShellTimeout = 600 * time.Second

	timeout := defaultShellTimeout
	if params.Timeout > 0 {
		timeout = time.Duration(params.Timeout) * time.Second
		if timeout > maxShellTimeout {
			log.WithFields(log.Fields{
				"requested": timeout,
				"max":       maxShellTimeout,
			}).Warn("Shell timeout exceeds maximum, capping")
			timeout = maxShellTimeout
		}
	}

	// 使用传入的 context 作为父 context，支持外部取消（如用户 stop）
	parentCtx := context.Background()
	if toolCtx != nil && toolCtx.Ctx != nil {
		parentCtx = toolCtx.Ctx
	}

	userID := ""
	workspaceRoot := ""
	execDir := ""
	if toolCtx != nil {
		workspaceRoot = toolCtx.WorkspaceRoot
		if toolCtx.CurrentDir != "" {
			execDir = toolCtx.CurrentDir
		} else if toolCtx.WorkspaceRoot != "" {
			execDir = toolCtx.WorkspaceRoot
		} else {
			execDir = toolCtx.WorkingDir
		}
		userID = toolCtx.OriginUserID
		if userID == "" {
			userID = toolCtx.SenderID // fallback
		}
	}

	// 沙箱模式：workspace 必须用宿主机路径（用于 bind mount / 容器查找），
	// 不能用容器内路径（CurrentDir），否则会导致容器 mount 校验失败并重建。
	sandboxWorkspace := workspaceRoot
	if sandboxWorkspace == "" {
		sandboxWorkspace = execDir
	}

	// 使用 ToolContext 中的沙箱实例（由 SandboxRouter 按用户路由注入）
	sandbox := toolCtx.Sandbox
	if sandbox == nil {
		sandbox = GetSandbox()
	}

	// 获取容器默认 shell 并使用 login shell 执行命令
	shell, err := sandbox.GetShell(userID, sandboxWorkspace)
	if err != nil {
		return nil, fmt.Errorf("failed to get shell: %w", err)
	}

	// 构建登录 shell 命令
	shellCmd := params.Command

	// 审计日志：记录每次 shell 执行
	log.WithFields(log.Fields{
		"command":    params.Command,
		"timeout":    timeout,
		"background": params.Background,
	}).Debug("Shell command executing")

	// Build ExecSpec based on sandbox mode
	buildSpec := func() ExecSpec {
		switch sandbox.Name() {
		case "docker":
			dir := ""
			if toolCtx != nil && toolCtx.CurrentDir != "" {
				dir = toolCtx.CurrentDir
			} else if toolCtx != nil && toolCtx.Sandbox != nil && toolCtx.Sandbox.Name() != "none" {
				dir = toolCtx.Sandbox.Workspace(toolCtx.OriginUserID)
			}
			return ExecSpec{
				Command:   shell,
				Args:      []string{shell, "-l", "-c", shellCmd},
				Shell:     false,
				Dir:       dir,
				Timeout:   timeout,
				Workspace: sandboxWorkspace,
				UserID:    userID,
			}
		case "remote":
			remoteDir := ""
			if toolCtx != nil && toolCtx.CurrentDir != "" {
				remoteDir = toolCtx.CurrentDir
			} else if rs, ok := sandbox.(*RemoteSandbox); ok {
				remoteDir = rs.Workspace(userID)
			}
			return ExecSpec{
				Command: shell,
				Args:    []string{shell, "-l", "-c", shellCmd},
				Shell:   false,
				Dir:     remoteDir,
				Timeout: timeout,
				UserID:  userID,
			}
		default:
			// None sandbox: use platform-aware shell args.
			// Unix: bash -l -c "command" (login shell, loads profile)
			// Windows: powershell.exe -Command "command" (loads profile by default)
			args := LoginShellArgs(shell, shellCmd)
			return ExecSpec{
				Command:   shell,
				Args:      args,
				Shell:     false,
				Dir:       execDir,
				Timeout:   timeout,
				UserID:    userID,
				RunAsUser: params.RunAs,
			}
		}
	}

	// Background mode: launch task in goroutine, return task ID immediately
	if params.Background {
		return t.executeBackground(toolCtx, shellCmd, sandbox, buildSpec)
	}

	// Foreground mode: synchronous execution with auto-promote on timeout
	return t.executeForeground(toolCtx, shellCmd, sandbox, parentCtx, timeout, buildSpec)
}

// executeBackground launches a command as a background task.
func (t *ShellTool) executeBackground(
	toolCtx *ToolContext,
	command string,
	sandbox Sandbox,
	buildSpec func() ExecSpec,
) (*ToolResult, error) {
	if toolCtx == nil || toolCtx.BgTaskManager == nil {
		return nil, fmt.Errorf("background tasks not supported (BgTaskManager not configured)")
	}

	sessionKey := toolCtx.BgSessionKey
	if sessionKey == "" {
		sessionKey = toolCtx.Channel + ":" + toolCtx.ChatID
	}
	senderID := toolCtx.OriginUserID

	task := toolCtx.BgTaskManager.Start(sessionKey, senderID, command,
		func(ctx context.Context, outputBuf func(string)) (int, error) {
			spec := buildSpec()
			spec.Timeout = 0 // no timeout for background
			return sandboxExecAsync(ctx, sandbox, spec, outputBuf)
		},
	)

	result := fmt.Sprintf(
		"Background task started [id: %s]\nCommand: %s\n\n"+
			"The task is running in the background. You can continue working.\n"+
			"When it completes, the output will be automatically injected into the conversation.\n"+
			"- Use task_status to check current progress (but do NOT poll — if running, wait or do other work first)\n"+
			"- Use task_kill to terminate the task",
		task.ID, task.Command,
	)

	return NewResultWithTips(result, fmt.Sprintf("Background task running: bg:%s", task.ID)), nil
}

// executeForeground runs a command synchronously. If it times out, auto-promotes
// to a background task (like Claude Code) so no work is lost.
func (t *ShellTool) executeForeground(
	toolCtx *ToolContext,
	command string,
	sandbox Sandbox,
	parentCtx context.Context,
	timeout time.Duration,
	buildSpec func() ExecSpec,
) (*ToolResult, error) {
	// Build spec with KeepAlive for none-sandbox so timeout doesn't kill the process.
	// The process can then be adopted by BgTaskManager on timeout.
	spec := buildSpec()
	if sandbox.Name() == "none" && toolCtx != nil && toolCtx.BgTaskManager != nil {
		spec.KeepAlive = true
	}
	result, err := sandbox.Exec(parentCtx, spec)

	if err != nil {
		return nil, fmt.Errorf("sandbox exec: %w", err)
	}

	// 合并输出
	var resultBuilder strings.Builder
	if result.Stdout != "" {
		resultBuilder.WriteString(result.Stdout)
	}
	if result.Stderr != "" {
		if resultBuilder.Len() > 0 {
			resultBuilder.WriteString("\n")
		}
		resultBuilder.WriteString("[stderr] ")
		resultBuilder.WriteString(result.Stderr)
	}
	output := strings.TrimSpace(resultBuilder.String())

	if result.TimedOut {
		// AUTO-PROMOTE: convert timed-out command to background task
		if toolCtx != nil && toolCtx.BgTaskManager != nil {
			sessionKey := toolCtx.BgSessionKey
			if sessionKey == "" {
				sessionKey = toolCtx.Channel + ":" + toolCtx.ChatID
			}
			senderID := toolCtx.OriginUserID

			var task *BackgroundTask

			// If the sandbox supports KeepAlive and returned a live process,
			// adopt it (no re-execution) — the original process continues running.
			if result.Process != nil {
				partialOutput := output
				var ongoingFn func() string
				if result.OngoingOutput != nil {
					ongoingFn = result.OngoingOutput
				}
				task = toolCtx.BgTaskManager.Adopt(sessionKey, senderID, command, result.Process, partialOutput, result.ExitCodeCh, ongoingFn)
				log.WithFields(log.Fields{
					"command": command,
					"timeout": timeout,
					"task_id": task.ID,
				}).Info("Timed-out command adopted as background task (no re-exec)")
			} else {
				// Sandbox doesn't support KeepAlive (docker/remote) — fall back to re-execution.
				partialOutput := output
				task = toolCtx.BgTaskManager.Start(sessionKey, senderID, command,
					func(ctx context.Context, outputBuf func(string)) (int, error) {
						spec := buildSpec()
						spec.Timeout = 0
						if partialOutput != "" {
							outputBuf(partialOutput + "\n\n--- [restarted after timeout] ---\n")
						}
						return sandboxExecAsync(ctx, sandbox, spec, outputBuf)
					},
				)
				log.WithFields(log.Fields{
					"command": command,
					"timeout": timeout,
					"task_id": task.ID,
				}).Info("Timed-out command auto-promoted to background task (re-exec)")
			}

			timeoutMsg := fmt.Sprintf(
				"[TIMEOUT after %s] Command timed out. Auto-promoted to background task [id: %s]\n"+
					"Partial output before timeout:\n%s\n\n"+
					"The command continues running in the background. Its output will be injected when done.\n"+
					"- Use task_status to check progress (but do NOT poll — if running, wait or do other work first)\n"+
					"- Use task_kill to terminate",
				timeout, task.ID, output,
			)
			return NewResultWithTips(timeoutMsg, fmt.Sprintf("Auto-promoted to background: bg:%s", task.ID)), nil
		}

		// No BgTaskManager — fall back to old behavior
		timeoutErr := fmt.Sprintf("[TIMEOUT after %s] Command timed out", timeout)
		if output != "" {
			timeoutErr = fmt.Sprintf("[TIMEOUT after %s] Partial output:\n%s", timeout, output)
		}
		log.WithFields(log.Fields{
			"command": command,
			"timeout": timeout,
			"output":  output,
		}).Warn("Shell command timed out")
		return NewErrorResult(timeoutErr), nil
	}

	if result.ExitCode != 0 {
		var errMsg string
		if output != "" {
			errMsg = fmt.Sprintf("[EXIT %d] %s\n%s", result.ExitCode, command, output)
		} else if result.Stderr != "" {
			errMsg = fmt.Sprintf("[EXIT %d] %s\n[stderr] %s", result.ExitCode, command, result.Stderr)
		} else {
			errMsg = fmt.Sprintf("[EXIT %d] %s (no output)", result.ExitCode, command)
		}

		log.WithFields(log.Fields{
			"command":  command,
			"exitCode": result.ExitCode,
			"stderr":   result.Stderr,
		}).Warn("Shell command failed")

		return NewErrorResult(errMsg), nil
	}

	if output == "" {
		return NewResult("Command executed successfully (no output)"), nil
	}

	res := NewResult(output)
	if tip := detectCdTip(command); tip != "" {
		res = res.WithTips(tip)
	}
	return res, nil
}

// sandboxExecAsync runs a sandbox command asynchronously, streaming output via outputBuf.
func sandboxExecAsync(
	ctx context.Context,
	sandbox Sandbox,
	spec ExecSpec,
	outputBuf func(string),
) (int, error) {
	switch sandbox.Name() {
	case "none":
		return noneSandboxExecAsync(ctx, spec, outputBuf)
	case "remote":
		return remoteSandboxExecAsync(ctx, sandbox, spec, outputBuf)
	default:
		// Docker: synchronous fallback (timeout=0 means no timeout)
		result, err := sandbox.Exec(ctx, spec)
		if outputBuf != nil && result != nil {
			if result.Stdout != "" {
				outputBuf(result.Stdout)
			}
			if result.Stderr != "" {
				outputBuf("[stderr] " + result.Stderr)
			}
		}
		if err != nil {
			if result != nil {
				return result.ExitCode, err
			}
			return -1, err
		}
		return result.ExitCode, nil
	}
}

// remoteSandboxExecAsync runs a command on a remote runner asynchronously.
// It starts the command via bg_exec protocol, then polls status until completion.
func remoteSandboxExecAsync(
	ctx context.Context,
	sandbox Sandbox,
	spec ExecSpec,
	outputBuf func(string),
) (int, error) {
	rs, ok := sandbox.(*RemoteSandbox)
	if !ok {
		return -1, fmt.Errorf("remote sandbox type assertion failed")
	}

	// Generate a unique task ID for the runner.
	taskID := "remote-" + generateID()

	// Start the background task on the runner.
	if err := rs.ExecBg(ctx, spec, taskID); err != nil {
		return -1, fmt.Errorf("remote bg_exec: %w", err)
	}

	// Poll until the task completes or context is cancelled.
	const pollInterval = 2 * time.Second
	for {
		select {
		case <-ctx.Done():
			// Try to kill the task on the runner before returning.
			killCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
			rs.KillBg(killCtx, spec.UserID, taskID)
			cancel()
			return -1, ctx.Err()
		case <-time.After(pollInterval):
		}

		status, err := rs.StatusBg(ctx, spec.UserID, taskID)
		if err != nil {
			return -1, fmt.Errorf("remote bg_status: %w", err)
		}

		// Stream any new output.
		if outputBuf != nil {
			if status.Stdout != "" {
				outputBuf(status.Stdout)
			}
			if status.Stderr != "" {
				outputBuf("[stderr] " + status.Stderr)
			}
		}

		switch status.Status {
		case "completed":
			return status.ExitCode, nil
		case "failed":
			return status.ExitCode, fmt.Errorf("remote task failed with exit code %d", status.ExitCode)
		case "killed":
			return -1, fmt.Errorf("remote task was killed")
		case "running":
			// Continue polling.
		default:
			return -1, fmt.Errorf("unknown remote task status: %s", status.Status)
		}
	}
}

// cdPattern detects standalone cd commands (not inside subshells, comments, or strings).
// Matches: "cd foo", "cd /path", "cd ..", "cd ~", as well as "cd foo && ls" etc.
var cdPattern = regexp.MustCompile(`(?:^|&&|\|\||;)\s*cd\s+`)

// detectCdTip returns a tip string if the command contains a cd that won't persist.
func detectCdTip(command string) string {
	if !cdPattern.MatchString(command) {
		return ""
	}
	return `NOTE: "cd" inside Shell only affects this single command — the working directory resets on the next tool call. Use the Cd tool to persistently change directory.`
}

// dangerPatterns 定义绝对禁止执行的命令模式（黑名单拦截，直接拒绝）
var dangerPatterns = []struct {
	pattern *regexp.Regexp
	reason  string
}{
	{regexp.MustCompile(`rm\s+-[^\s]*rf\s+/\s*`), "rm -rf / is destructive and will wipe the entire filesystem"},
	{regexp.MustCompile(`mkfs\b`), "mkfs will destroy filesystem data"},
	{regexp.MustCompile(`dd\s+.*(/dev/zero|/dev/random|/dev/null)\s+.*of=/dev/`), "dd writing to device is destructive"},
	{regexp.MustCompile(`:\(\)\s*\{.*\}\s*;`), "fork bomb detected"},
	{regexp.MustCompile(`chmod\s+777\s+/\s*`), "chmod 777 / is a security risk"},
	{regexp.MustCompile(`mv\s+/\s+/dev/null`), "mv / /dev/null is destructive"},
}

// warningPatterns 定义高危命令（告警但允许执行）
var warningPatterns = []*regexp.Regexp{
	regexp.MustCompile(`\brm\s+(-[^\s]*rf|-rf)\b`),
	regexp.MustCompile(`\bdd\b`),
	regexp.MustCompile(`\bmkfs\b`),
	regexp.MustCompile(`\bchmod\s+777\b`),
	regexp.MustCompile(`\b(format| FORMAT)\b`),
}

// checkDangerousCommand 检查命令是否包含危险模式
// 返回 (blocked, reason)，如果 blocked=true 则应拒绝执行
// disallowSudo: 当为 true 时（run_as 模式），任何形式的 sudo 都被禁止
func checkDangerousCommand(ctx context.Context, cmd string, disallowSudo bool) (bool, string) {
	// 检查绝对禁止模式
	for _, dp := range dangerPatterns {
		if dp.pattern.MatchString(cmd) {
			return true, dp.reason
		}
	}

	// sudo 检查
	sudoRe := regexp.MustCompile(`\bsudo\b`)
	if sudoRe.MatchString(cmd) {
		if disallowSudo {
			// run_as 模式：用户切换由框架控制，禁止 LLM 使用任何形式的 sudo
			return true, "sudo is not allowed when run_as is set (user switching is handled by the framework)"
		}
		defaultUser, privilegedUser := PermUsersFromContext(ctx)
		if defaultUser != "" || privilegedUser != "" {
			// Permission control enabled: all LLM-authored sudo is forbidden.
			// User switching must go through run_as so approval/default-user policy can be enforced.
			return true, "sudo is not allowed when permission control is enabled (use run_as instead so the framework can enforce approval policy)"
		}
		// No additional sudo restrictions in default mode.
	}

	// 检查高危告警模式（仅日志记录，不拦截）
	for _, wp := range warningPatterns {
		if wp.MatchString(cmd) {
			log.WithField("command", cmd).Warn("Dangerous command detected (allowed with warning)")
			break
		}
	}

	return false, ""
}
