package tools

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"strings"
	"sync"

	"agentu/pkg/llm"
)

type ShellRunTool struct {
	workingDir string
}

func NewShellRunTool(workingDir string) *ShellRunTool {
	return &ShellRunTool{workingDir: workingDir}
}
func (*ShellRunTool) Mutates() bool { return true }

func (t *ShellRunTool) Definition() llm.Tool {
	return llm.Tool{
		Type: "function",
		Function: llm.ToolFunction{
			Name:        "shell_run",
			Description: "Run a non-interactive shell command on the local machine. Use for builds, tests, local status inspection, and local automation. Avoid commands that require an interactive TTY, credentials, or open long-running sessions.",
			Parameters: JSONSchema(`{
  "type": "object",
  "properties": {
    "command": {"type": "string", "description": "Command to execute."},
    "workdir": {"type": "string", "description": "Optional working directory. Relative paths are resolved from the configured agent working directory."}
  },
  "required": ["command"],
  "additionalProperties": false
}`),
		},
	}
}

type shellRunArgs struct {
	Command string `json:"command"`
	Workdir string `json:"workdir"`
}

func (t *ShellRunTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
	return t.ExecuteStream(ctx, raw, nil)
}

func (t *ShellRunTool) ExecuteStream(ctx context.Context, raw json.RawMessage, emit func(string) error) (string, error) {
	var args shellRunArgs
	if err := json.Unmarshal(raw, &args); err != nil {
		return "", fmt.Errorf("parse arguments: %w", err)
	}
	if strings.TrimSpace(args.Command) == "" {
		return "", errors.New("command is required")
	}

	workdir, err := resolvePath(t.workingDir, args.Workdir)
	if err != nil {
		return "", err
	}

	shell := os.Getenv("SHELL")
	if shell == "" {
		shell = "/bin/sh"
	}
	cmd := exec.CommandContext(ctx, shell, "-lc", args.Command)
	cmd.Dir = workdir

	var mu sync.Mutex
	var combined bytes.Buffer
	writer := &shellStreamWriter{mu: &mu, buf: &combined, emit: emit}
	cmd.Stdout = writer
	cmd.Stderr = writer

	runErr := cmd.Run()
	text := FormatResult(combined.String())
	if errors.Is(ctx.Err(), context.DeadlineExceeded) {
		return text, fmt.Errorf("command timed out: %w", ctx.Err())
	}
	if runErr != nil {
		if text != "" {
			return text, fmt.Errorf("command failed: %w", runErr)
		}
		return "", fmt.Errorf("command failed: %w", runErr)
	}
	return text, nil
}

// shellStreamWriter captures command output for the final result while
// forwarding every write to the streaming callback as it happens.
type shellStreamWriter struct {
	mu   *sync.Mutex
	buf  *bytes.Buffer
	emit func(string) error
}

func (w *shellStreamWriter) Write(p []byte) (int, error) {
	w.mu.Lock()
	n, writeErr := w.buf.Write(p)
	w.mu.Unlock()
	if writeErr != nil {
		return n, writeErr
	}
	if w.emit != nil {
		if err := w.emit(string(p)); err != nil {
			return 0, err
		}
	}
	return len(p), nil
}