package tools

import (
	"context"
	"fmt"
	"os/exec"
	"runtime"
	"strings"
	"time"

	llm "gitea.stevedudenhoeffer.com/steve/go-llm/v2"
)

// ExecParams defines parameters for the exec tool.
type ExecParams struct {
	Command string `json:"command" description:"The shell command to execute"`
}

// ExecOption configures the exec tool.
type ExecOption func(*execConfig)

type execConfig struct {
	allowedCommands []string
	workDir         string
	timeout         time.Duration
}

// WithAllowedCommands restricts which commands can be executed.
// If empty, all commands are allowed.
func WithAllowedCommands(cmds []string) ExecOption {
	return func(c *execConfig) { c.allowedCommands = cmds }
}

// WithWorkDir sets the working directory for command execution.
func WithWorkDir(dir string) ExecOption {
	return func(c *execConfig) { c.workDir = dir }
}

// WithExecTimeout sets the maximum execution time.
func WithExecTimeout(d time.Duration) ExecOption {
	return func(c *execConfig) { c.timeout = d }
}

// Exec creates a shell command execution tool.
//
// Example:
//
//	tools := llm.NewToolBox(
//	    tools.Exec(tools.WithAllowedCommands([]string{"ls", "cat", "grep"})),
//	)
func Exec(opts ...ExecOption) llm.Tool {
	cfg := &execConfig{
		timeout: 30 * time.Second,
	}
	for _, opt := range opts {
		opt(cfg)
	}

	return llm.Define[ExecParams]("exec", "Execute a shell command and return its output",
		func(ctx context.Context, p ExecParams) (string, error) {
			// Check allowed commands
			if len(cfg.allowedCommands) > 0 {
				parts := strings.Fields(p.Command)
				if len(parts) == 0 {
					return "", fmt.Errorf("empty command")
				}
				allowed := false
				for _, cmd := range cfg.allowedCommands {
					if parts[0] == cmd {
						allowed = true
						break
					}
				}
				if !allowed {
					return "", fmt.Errorf("command %q is not in the allowed list", parts[0])
				}
			}

			ctx, cancel := context.WithTimeout(ctx, cfg.timeout)
			defer cancel()

			var cmd *exec.Cmd
			if runtime.GOOS == "windows" {
				cmd = exec.CommandContext(ctx, "cmd", "/C", p.Command)
			} else {
				cmd = exec.CommandContext(ctx, "sh", "-c", p.Command)
			}

			if cfg.workDir != "" {
				cmd.Dir = cfg.workDir
			}

			output, err := cmd.CombinedOutput()
			if err != nil {
				return fmt.Sprintf("Error: %s\nOutput: %s", err.Error(), string(output)), nil
			}

			return string(output), nil
		},
	)
}