package tools

import (
	"bytes"
	"context"
	"fmt"
	"os/exec"
	"strings"
	"syscall"
	"time"
)

const (
	defaultBashTimeout = 60 * time.Second
	maxBashTimeout     = 10 * time.Minute
	// Output is capped head and tail. Head-only truncation would cut off the
	// end of a test run, which is exactly where the failure is.
	bashHeadLines = 100
	bashTailLines = 100
)

// BashArgs is the bash input.
type BashArgs struct {
	Command   string `json:"command" jsonschema:"description=Shell command to run, from the working root"`
	TimeoutMS int    `json:"timeout_ms,omitempty" jsonschema:"description=Timeout in milliseconds; defaults to 60000, maximum 600000"`
}

// NewBash builds the bash tool.
func NewBash(root string) Tool {
	return New("bash", strings.TrimSpace(`
Run a shell command from the working root and return its exit code, stdout and
stderr. Use it to run the project's tests and build. Output is capped, keeping
the first and last hundred lines. Prefer glob and grep over find and cat: they
are faster and their output is bounded.`),
		func(ctx context.Context, in BashArgs) (string, error) {
			if strings.TrimSpace(in.Command) == "" {
				return "", fmt.Errorf("command is required")
			}
			timeout := defaultBashTimeout
			if in.TimeoutMS > 0 {
				timeout = time.Duration(in.TimeoutMS) * time.Millisecond
			}
			if timeout > maxBashTimeout {
				timeout = maxBashTimeout
			}

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

			cmd := exec.CommandContext(ctx, "bash", "-c", in.Command)
			cmd.Dir = resolveBase(root, "")
			// A new process group, so a command that spawns children can be
			// killed whole. Killing only the leader leaves orphans holding the
			// pipes open and the loop hangs waiting on them.
			cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
			cmd.Cancel = func() error {
				if cmd.Process == nil {
					return nil
				}
				return syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
			}
			cmd.WaitDelay = 2 * time.Second

			var stdout, stderr bytes.Buffer
			cmd.Stdout = &stdout
			cmd.Stderr = &stderr

			start := time.Now()
			runErr := cmd.Run()
			elapsed := time.Since(start)

			exitCode := cmd.ProcessState.ExitCode()
			timedOut := ctx.Err() == context.DeadlineExceeded

			var b strings.Builder
			switch {
			case timedOut:
				fmt.Fprintf(&b, "timed out after %s (killed)\n", timeout)
			case runErr != nil && exitCode < 0:
				fmt.Fprintf(&b, "failed to run: %v\n", runErr)
			default:
				fmt.Fprintf(&b, "exit code: %d (%s)\n", exitCode, elapsed.Round(time.Millisecond))
			}
			writeStream(&b, "stdout", stdout.String())
			writeStream(&b, "stderr", stderr.String())

			out := strings.TrimRight(b.String(), "\n")
			// A non-zero exit is reported as an error so the model sees it as a
			// failure it has to respond to, with the output still attached.
			if timedOut || exitCode != 0 {
				return "", fmt.Errorf("%s", out)
			}
			return out, nil
		})
}

func writeStream(b *strings.Builder, label, s string) {
	s = strings.TrimRight(s, "\n")
	if s == "" {
		return
	}
	fmt.Fprintf(b, "\n--- %s ---\n%s\n", label, headTail(s, bashHeadLines, bashTailLines))
}