package main

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"math/bits"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"
	"time"
	"unicode/utf8"
)

const (
	// clipCharacters is the most characters of a tool output: its head and
	// its tail, half each.
	clipCharacters = 30000

	// clipBytes is the most bytes that a tool keeps of each end of an
	// output, enough for clipCharacters/2 characters of 4 bytes.
	clipBytes = 4 * clipCharacters / 2

	// shellWaitDelay is how long a command can keep its output open after it
	// ends or after its kill.
	shellWaitDelay = 2 * time.Second
)

// chatLog is the log that the tools write their calls and results to.
type chatLog interface {
	Append(kind, text string) (Message, error)
	Flush(ctx context.Context) error
}

// chatReader gives the messages and the nodes that zoom and date read.
type chatReader interface {
	Len() int
	Message(id int) (Message, error)
	Node(level, index int) (Node, bool, error)
}

// toolbox runs the tool calls of one agent. Each call and each result goes to
// the log as it happens: the call as a tool message, the result as an echo
// message. The call is durable before its tool runs, so the log never misses
// an action outside the agent. Commands and files run as the user in the
// working directory, with no sandbox.
type toolbox struct {
	log     chatLog
	chat    chatReader
	workdir string
}

// shellArguments are the arguments of shell.
type shellArguments struct {
	Command string `json:"command"`
}

// pathArguments are the arguments of read.
type pathArguments struct {
	Path string `json:"path"`
}

// writeArguments are the arguments of write.
type writeArguments struct {
	Path string `json:"path"`
	Text string `json:"text"`
}

// editArguments are the arguments of edit.
type editArguments struct {
	Path string `json:"path"`
	Old  string `json:"old"`
	New  string `json:"new"`
}

// zoomArguments are the arguments of zoom.
type zoomArguments struct {
	ID int `json:"id"`
	N  int `json:"n"`
}

// dateArguments are the arguments of date.
type dateArguments struct {
	ID int `json:"id"`
}

// Call logs the call of the tool name with its JSON arguments, makes the log
// durable, runs the tool, and logs its clipped output, which it returns. A
// failed tool gives its error as the output. An error of Call itself is an
// error of the log. When ctx ends before the flush, the tool does not run.
func (b *toolbox) Call(ctx context.Context, name, arguments string) (string, error) {
	_, err := b.log.Append("tool", callText(name, arguments))
	if err != nil {
		return "", err
	}

	err = b.log.Flush(ctx)
	if err != nil && ctx.Err() == nil {
		return "", err
	}

	output := "not run: the turn was interrupted"
	if err == nil {
		output = clip(b.run(ctx, name, arguments))
	}

	_, err = b.log.Append("echo", output)
	if err != nil {
		return "", err
	}

	return output, nil
}

// run runs one tool and returns its output, or its error as text.
func (b *toolbox) run(ctx context.Context, name, arguments string) string {
	var output string
	var err error
	switch name {
	case "zoom":
		var args zoomArguments
		err = decodeArguments(arguments, &args)
		if err == nil {
			output, err = b.zoom(args.ID, args.N)
		}
	case "date":
		var args dateArguments
		err = decodeArguments(arguments, &args)
		if err == nil {
			output, err = b.date(args.ID)
		}
	case "shell":
		var args shellArguments
		err = decodeArguments(arguments, &args)
		if err == nil {
			output = b.shell(ctx, args.Command)
		}
	case "read":
		var args pathArguments
		err = decodeArguments(arguments, &args)
		if err == nil {
			output, err = b.read(args.Path)
		}
	case "write":
		var args writeArguments
		err = decodeArguments(arguments, &args)
		if err == nil {
			output, err = b.write(args.Path, args.Text)
		}
	case "edit":
		var args editArguments
		err = decodeArguments(arguments, &args)
		if err == nil {
			output, err = b.edit(args.Path, args.Old, args.New)
		}
	default:
		err = fmt.Errorf("unknown tool %q", name)
	}

	if err != nil {
		return "error: " + err.Error()
	}

	return output
}

// zoom returns message id whole for n = 1, or else the two lines under the line
// id+n.
func (b *toolbox) zoom(id, n int) (string, error) {
	if n < 1 || bits.OnesCount(uint(n)) != 1 || id < 0 || id%n != 0 {
		return "", fmt.Errorf("line %d+%d does not exist: n must be a power of 2 and id a multiple of n", id, n)
	}

	if id+n > b.chat.Len() {
		return "", fmt.Errorf("line %d+%d goes past the last message, %d", id, n, b.chat.Len()-1)
	}

	if n == 1 {
		message, err := b.chat.Message(id)
		if err != nil {
			return "", err
		}

		return leafSource(message), nil
	}

	level := bits.TrailingZeros(uint(n)) - 1
	var lines []string
	for _, index := range []int{2 * id / n, 2*id/n + 1} {
		node, built, err := b.chat.Node(level, index)
		if err != nil {
			return "", err
		}

		text := pending
		if built {
			text = strings.ReplaceAll(node.Text, "\n", " ")
		}

		lines = append(lines, fmt.Sprintf("%d+%d|%s", index<<level, 1<<level, text))
	}

	return strings.Join(lines, "\n"), nil
}

// date returns the date and time of message id, with its weekday.
func (b *toolbox) date(id int) (string, error) {
	if id < 0 || id >= b.chat.Len() {
		return "", fmt.Errorf("message %d does not exist", id)
	}

	message, err := b.chat.Message(id)
	if err != nil {
		return "", err
	}

	date, err := time.Parse(time.RFC3339Nano, message.Date)
	if err != nil {
		return "", err
	}

	return date.Local().Format("Monday 2006-01-02 15:04:05 -07:00"), nil
}

// shell runs command with bash in the working directory and returns its output
// and error output. A failed command adds its exit status. When ctx ends, the
// whole process group of the command is killed.
func (b *toolbox) shell(ctx context.Context, command string) string {
	cmd := exec.CommandContext(ctx, "bash", "-c", command)
	cmd.Dir = b.workdir
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Setpgid: true,
	}
	cmd.Cancel = func() error {
		return syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
	}
	cmd.WaitDelay = shellWaitDelay

	output := &headTail{}
	cmd.Stdout = output
	cmd.Stderr = output
	err := cmd.Run()

	text := output.String()
	var exit *exec.ExitError
	switch {
	case ctx.Err() != nil:
		text += "\n[interrupted]"
	case errors.As(err, &exit):
		text += "\n[" + exit.String() + "]"
	case err != nil:
		text += "\n[error: " + err.Error() + "]"
	case text == "":
		text = "(no output)"
	}

	return text
}

// read returns the text of the file at path.
func (b *toolbox) read(path string) (text string, err error) {
	file, err := os.Open(b.path(path))
	if err != nil {
		return "", err
	}

	defer func() {
		err = errors.Join(err, file.Close())
	}()

	output := &headTail{}
	_, err = io.Copy(output, file)
	if err != nil {
		return "", err
	}

	return output.String(), nil
}

// write replaces the file at path with text, and makes its missing
// directories.
func (b *toolbox) write(path, text string) (string, error) {
	path = b.path(path)
	err := os.MkdirAll(filepath.Dir(path), 0o755)
	if err != nil {
		return "", err
	}

	err = os.WriteFile(path, []byte(text), 0o644)
	if err != nil {
		return "", err
	}

	return fmt.Sprintf("wrote %d bytes to %s", len(text), path), nil
}

// edit replaces the one occurrence of old in the file at path with
// replacement. It keeps the mode of the file.
func (b *toolbox) edit(path, old, replacement string) (string, error) {
	path = b.path(path)
	data, err := os.ReadFile(path)
	if err != nil {
		return "", err
	}

	if old == "" {
		return "", errors.New("old is empty")
	}

	count := strings.Count(string(data), old)
	if count != 1 {
		return "", fmt.Errorf("old occurs %d times in %s, not once", count, path)
	}

	info, err := os.Stat(path)
	if err != nil {
		return "", err
	}

	err = os.WriteFile(path, []byte(strings.Replace(string(data), old, replacement, 1)), info.Mode().Perm())
	if err != nil {
		return "", err
	}

	return "edited " + path, nil
}

// path returns path, relative to the working directory when it is not
// absolute.
func (b *toolbox) path(path string) string {
	if filepath.IsAbs(path) {
		return path
	}

	return filepath.Join(b.workdir, path)
}

// headTail keeps the first and the last clipBytes bytes written to it, and
// counts the bytes between them that it drops.
type headTail struct {
	head    []byte
	tail    []byte
	dropped int
}

// Write keeps the bytes of p at the ends of the output.
func (h *headTail) Write(p []byte) (int, error) {
	n := len(p)
	room := min(clipBytes-len(h.head), len(p))
	h.head = append(h.head, p[:room]...)
	p = p[room:]

	h.tail = append(h.tail, p...)
	if len(h.tail) > clipBytes {
		cut := len(h.tail) - clipBytes
		h.dropped += cut
		h.tail = bytes.Clone(h.tail[cut:])
	}

	return n, nil
}

// String returns the output, clipped when it is too long. Invalid UTF-8 becomes
// the replacement character.
func (h *headTail) String() string {
	replacement := string(utf8.RuneError)
	if h.dropped == 0 {
		return clip(strings.ToValidUTF8(string(h.head)+string(h.tail), replacement))
	}

	head := firstCharacters(strings.ToValidUTF8(string(h.head), replacement), clipCharacters/2)
	tail := lastCharacters(strings.ToValidUTF8(string(h.tail), replacement), clipCharacters/2)
	cut := len(h.head) + h.dropped + len(h.tail) - len(head) - len(tail)

	return joinClipped(head, tail, cut)
}

// callText returns the log text of a tool call: the name and the JSON
// arguments.
func callText(name, arguments string) string {
	return name + " " + arguments
}

// decodeArguments decodes the JSON arguments of a call into args. An unknown
// field is an error.
func decodeArguments(arguments string, args any) error {
	decoder := json.NewDecoder(strings.NewReader(arguments))
	decoder.DisallowUnknownFields()
	err := decoder.Decode(args)
	if err != nil {
		return fmt.Errorf("arguments: %w", err)
	}

	return nil
}

// clip returns text when it has at most clipCharacters characters. Otherwise it
// returns the first and the last clipCharacters/2 characters, with a mark of
// the cut between them.
func clip(text string) string {
	if utf8.RuneCountInString(text) <= clipCharacters {
		return text
	}

	head := firstCharacters(text, clipCharacters/2)
	tail := lastCharacters(text, clipCharacters/2)

	return joinClipped(head, tail, len(text)-len(head)-len(tail))
}

// joinClipped joins the head and the tail of a clipped output with a mark of
// the cut bytes.
func joinClipped(head, tail string, cut int) string {
	return head + "\n[... " + strconv.Itoa(cut) + " bytes cut ...]\n" + tail
}

// firstCharacters returns the first count characters of text.
func firstCharacters(text string, count int) string {
	end := 0
	for range count {
		if end >= len(text) {
			break
		}

		_, size := utf8.DecodeRuneInString(text[end:])
		end += size
	}

	return text[:end]
}

// lastCharacters returns the last count characters of text.
func lastCharacters(text string, count int) string {
	start := len(text)
	for range count {
		if start <= 0 {
			break
		}

		_, size := utf8.DecodeLastRuneInString(text[:start])
		start -= size
	}

	return text[start:]
}