package exec

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"os"
	"os/exec"
	"runtime"
	"sort"
	"strings"
	"sync"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
)

// TerminalExecutor is the default command Provider. It runs allowlisted commands
// inside a grant-controlled terminal sandbox (DevDocs/design/04-运行时.md §8):
// command allowlist, working-directory confinement, environment allowlist,
// timeout, and output truncation. GUI and virtual-machine providers are reserved
// for later milestones and surface only through the Provider interface.
type TerminalExecutor struct {
	cfg   config.ExecSandbox
	log   *slog.Logger
	roots []string // absolute, clean allowed roots for workdir confinement
}

// New builds a TerminalExecutor from cfg and log. The sandbox is closed unless
// cfg.Enabled is true; a closed executor still constructs fine and refuses every
// request with ErrExecutionDisabled. When enabled, at least one allowed root is
// required so workdir confinement has an anchor.
func New(cfg config.ExecSandbox, log *slog.Logger) (*TerminalExecutor, error) {
	if log == nil {
		return nil, ErrNilLogger
	}
	roots := absCleaners(cfg.AllowedRoots)
	if cfg.Enabled && len(roots) == 0 {
		return nil, fmt.Errorf("%w: an enabled sandbox needs at least one allowed root", ErrNoAllowedRoot)
	}
	return &TerminalExecutor{
		cfg:   cfg,
		log:   logger.WithSource(log, "runtime/exec"),
		roots: roots,
	}, nil
}

// Execute runs the command in req under the sandbox. It returns an
// *[ExecResult] on success (including a non-zero exit code) or a sentinel error
// (wrapped) on refusal or an underlying context failure.
func (t *TerminalExecutor) Execute(ctx context.Context, req ExecRequest) (*ExecResult, error) {
	if !t.cfg.Enabled {
		t.log.Warn("sandbox disabled; refusing to execute command")
		return nil, ErrExecutionDisabled
	}
	if len(req.Command) == 0 {
		return nil, fmt.Errorf("%w: no argv", ErrEmptyCommand)
	}
	if err := t.checkCommand(req.Command); err != nil {
		return nil, err
	}
	workdir, err := t.resolveWorkdir(req.WorkDir)
	if err != nil {
		return nil, err
	}
	env, err := t.buildEnv(req.Env)
	if err != nil {
		return nil, err
	}
	return t.runCmd(ctx, req, workdir, env)
}

// runCmd launches the already-validated command, streams its output through a
// capped writer, honors the deadline, and classifies the outcome.
func (t *TerminalExecutor) runCmd(ctx context.Context, req ExecRequest, workdir string, env []string) (*ExecResult, error) {
	timeout := req.Timeout
	if timeout <= 0 {
		timeout = DefaultTimeout
	}
	maxOut := req.MaxOutputBytes
	if maxOut <= 0 {
		maxOut = DefaultMaxOutputBytes
	}

	runCtx := ctx
	cancel := func() {}
	if timeout > 0 {
		runCtx, cancel = context.WithTimeout(ctx, timeout)
	}
	defer cancel()

	cmd := exec.CommandContext(runCtx, req.Command[0], req.Command[1:]...)
	cmd.Dir = workdir
	cmd.Env = env
	if len(req.Stdin) > 0 {
		cmd.Stdin = bytes.NewReader(req.Stdin)
	}
	out := &cappedOutput{max: maxOut}
	cmd.Stdout = out.stdoutWriter()
	cmd.Stderr = out.stderrWriter()

	t.log.Debug("starting command", "cmd", strings.Join(req.Command, " "),
		"workdir", workdir, "timeout", timeout.String())
	startedAt := time.Now()
	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("%w: %s: %w", ErrExecFailed, strings.Join(req.Command, " "), err)
	}
	result := &ExecResult{StartedAt: startedAt, PID: cmd.Process.Pid}

	waitErr := cmd.Wait()
	result.EndedAt = time.Now()
	result.Stdout = out.stdout.String()
	result.Stderr = out.stderr.String()
	result.Truncated = out.truncated

	if waitErr != nil {
		switch {
		case errors.Is(runCtx.Err(), context.DeadlineExceeded):
			t.log.Warn("command timed out; killed", "cmd", strings.Join(req.Command, " "), "pid", result.PID)
			return nil, fmt.Errorf("%w: %s", ErrTimeout, strings.Join(req.Command, " "))
		case errors.Is(runCtx.Err(), context.Canceled):
			return nil, fmt.Errorf("%w: %s", context.Canceled, strings.Join(req.Command, " "))
		}
		var exitErr *exec.ExitError
		if errors.As(waitErr, &exitErr) {
			result.ExitCode = exitErr.ExitCode()
		} else {
			return nil, fmt.Errorf("%w: %s: %w", ErrExecFailed, strings.Join(req.Command, " "), waitErr)
		}
	}
	t.log.Debug("command finished", "cmd", strings.Join(req.Command, " "),
		"exit", result.ExitCode, "truncated", result.Truncated)
	return result, nil
}

// checkCommand enforces the command allowlist and the per-argument injection
// rules. The command runs via exec.Command (never a shell), so metacharacters
// are rejected rather than interpreted.
func (t *TerminalExecutor) checkCommand(argv []string) error {
	if !t.allowed(argv[0]) {
		return fmt.Errorf("%w: %q", ErrCommandNotAllowed, argv[0])
	}
	for _, a := range argv {
		if err := validateArg(a); err != nil {
			return fmt.Errorf("%w: %q: %w", ErrCommandInjection, a, err)
		}
	}
	return nil
}

func (t *TerminalExecutor) allowed(name string) bool {
	for _, c := range t.cfg.CommandList {
		if c == name {
			return true
		}
	}
	return false
}

func (t *TerminalExecutor) envAllowed(name string) bool {
	for _, a := range t.cfg.AllowedEnv {
		if a == name {
			return true
		}
	}
	return false
}

// resolveWorkdir confines the process working directory to an allowed root. An
// empty request uses the first root; a relative path is resolved against it.
func (t *TerminalExecutor) resolveWorkdir(workdir string) (string, error) {
	if workdir == "" {
		if len(t.roots) == 0 {
			return "", ErrNoAllowedRoot
		}
		real, err := resolveReal(t.roots[0])
		if err != nil {
			return "", fmt.Errorf("%w: %s", ErrWorkDirOutsideRoot, t.roots[0])
		}
		return real, nil
	}
	abs, err := resolveInsideRoot(workdir, t.roots)
	if err != nil {
		if errors.Is(err, ErrPathOutsideRoot) || errors.Is(err, ErrSymlinkEscape) {
			return "", fmt.Errorf("%w: %s", ErrWorkDirOutsideRoot, workdir)
		}
		return "", err
	}
	return abs, nil
}

// terminalEnvAllowlist 返回运行终端子进程所必需的最小宿主环境变量键集合。
//
// 为什么是这些（审计口径）：终端命令可由模型/插件触发，若继承 os.Environ() 全集，
// 子进程就能读到宿主的 API key、token、DB 口令等一切凭据。但子进程又必须拿到一
// 撮由操作系统/运行时提供的定位与查找变量，否则连可执行文件都解析不了。这里只保
// 留「去哪找可执行文件、临时目录在哪、用户目录在哪、常规系统信息」这类不含凭据的
// 键：
//   - Windows：PATH/PATHEXT 定位可执行文件；SYSTEMROOT/WINDIR/COMSPEC 是 cmd 与
//     系统 DLL 加载的前提；TEMP/TMP 是临时目录；USERPROFILE/HOMEDRIVE/HOMEPATH/
//     APPDATA/LOCALAPPDATA 是用户目录；NUMBER_OF_PROCESSORS/OS/
//     PROCESSOR_ARCHITECTURE 是常规系统信息（部分程序依赖）。
//   - Unix：PATH 定位可执行文件；HOME/TMPDIR 是用户与临时目录；USER/SHELL/LANG/
//     LC_ALL 是常规终端身份与本地化变量。
//
// 其余变量（如 *_API_KEY、*_TOKEN、DB_*、JWT_*）一律不透传；调用方确需某个键时，
// 由 cfg.AllowedEnv 显式放行并经 req.Env 覆盖（见 envAllowed）。
func terminalEnvAllowlist() []string {
	if runtime.GOOS == "windows" {
		windowsKeys := []string{
			"PATH", "PATHEXT", "SYSTEMROOT", "WINDIR", "COMSPEC",
			"TEMP", "TMP", "USERPROFILE", "HOMEDRIVE", "HOMEPATH",
			"APPDATA", "LOCALAPPDATA", "NUMBER_OF_PROCESSORS", "OS",
			"PROCESSOR_ARCHITECTURE",
		}
		return windowsKeys
	}
	unixKeys := []string{"PATH", "HOME", "TMPDIR", "LANG", "LC_ALL", "USER", "SHELL"}
	return unixKeys
}

// envKeyEqual 比较两个环境变量名是否指向同一变量：Windows 的名字大小写不敏感
// （宿主常写成 Path，不能因此漏掉 PATH），Unix 敏感。
func envKeyEqual(a, b string) bool {
	if runtime.GOOS == "windows" {
		return strings.EqualFold(a, b)
	}
	return a == b
}

// envKeyIn 判断 name 是否属于 keys（按平台口径比较）。
func envKeyIn(name string, keys []string) bool {
	for _, k := range keys {
		if envKeyEqual(k, name) {
			return true
		}
	}
	return false
}

// envKeyOf 返回 m 中与 name 等价的键的原始拼写，用于覆盖时避免同名重复键。
func envKeyOf(m map[string]string, name string) (string, bool) {
	if _, ok := m[name]; ok {
		return name, true
	}
	if runtime.GOOS != "windows" {
		return "", false
	}
	for k := range m {
		if strings.EqualFold(k, name) {
			return k, true
		}
	}
	return "", false
}

// buildEnv 组装子进程环境：只透传运行终端所必需的最小宿主键集合
// （terminalEnvAllowlist），再叠加 req.Env 中经 envAllowed 显式放行的覆盖值。
// 宿主其余变量（含密钥）一律不透传；req.Env 里非白名单的键直接拒绝。
func (t *TerminalExecutor) buildEnv(extra map[string]string) ([]string, error) {
	overlays := make(map[string]string, len(extra))
	for k, v := range extra {
		if !t.envAllowed(k) {
			return nil, fmt.Errorf("%w: %s", ErrEnvNotAllowed, k)
		}
		overlays[k] = v
	}

	allow := terminalEnvAllowlist()
	kept := make(map[string]string)
	for _, e := range os.Environ() {
		k, v, ok := strings.Cut(e, "=")
		if !ok || k == "" || !envKeyIn(k, allow) {
			continue
		}
		kept[k] = v
	}

	for k, v := range overlays {
		if existing, ok := envKeyOf(kept, k); ok {
			delete(kept, existing)
		}
		kept[k] = v
	}

	out := make([]string, 0, len(kept))
	for k, v := range kept {
		out = append(out, k+"="+v)
	}
	sort.Strings(out)
	return out, nil
}

// validateArg rejects anything unsafe in a command argument: empty value, shell
// metacharacters, control characters, or a path-traversal ("..") component.
func validateArg(s string) error {
	if s == "" {
		return errors.New("empty argument")
	}
	if strings.ContainsAny(s, `;&|`+"`$()<>") {
		return errors.New("contains shell metacharacters")
	}
	if strings.ContainsAny(s, "\x00\r\n") {
		return errors.New("contains control characters")
	}
	for _, sep := range []string{"/", `\`} {
		for _, part := range strings.Split(s, sep) {
			if part == ".." {
				return errors.New("contains path traversal")
			}
		}
	}
	return nil
}

// cappedOutput captures stdout and stderr while enforcing a combined byte cap.
// Writes past the cap are dropped and the truncated flag is raised. The shared
// mutex makes the two streams safe to write concurrently by os/exec.
type cappedOutput struct {
	mu        sync.Mutex
	stdout    bytes.Buffer
	stderr    bytes.Buffer
	max       int
	used      int
	truncated bool
}

func (o *cappedOutput) stdoutWriter() io.Writer { return &capWriter{o: o, buf: &o.stdout} }
func (o *cappedOutput) stderrWriter() io.Writer { return &capWriter{o: o, buf: &o.stderr} }

type capWriter struct {
	o   *cappedOutput
	buf *bytes.Buffer
}

func (w *capWriter) Write(p []byte) (int, error) {
	w.o.mu.Lock()
	defer w.o.mu.Unlock()
	n := len(p)
	remaining := w.o.max - w.o.used
	if remaining <= 0 {
		w.o.truncated = true
		return n, nil
	}
	take := remaining
	if take > n {
		take = n
	}
	w.buf.Write(p[:take])
	w.o.used += take
	if take < n {
		w.o.truncated = true
	}
	return n, nil
}