kitty/tools/cmd/edit_in_kitty/main.go

// License: GPLv3 Copyright: 2022, Kovid Goyal, <kovid at kovidgoyal.net>

package edit_in_kitty

import (
	"encoding/base64"
	"fmt"
	"io"
	"io/fs"
	"os"
	"strconv"
	"strings"

	"golang.org/x/sys/unix"

	"kitty/tools/cli"
	"kitty/tools/tui"
	"kitty/tools/tui/loop"
	"kitty/tools/utils"
	"kitty/tools/utils/humanize"
)

var _ = fmt.Print

func encode(x string) string {
	return base64.StdEncoding.EncodeToString(utils.UnsafeStringToBytes(x))
}

type OnDataCallback = func(data_type string, data []byte) error

func edit_loop(data_to_send string, kill_if_signaled bool, on_data OnDataCallback) (err error) {
	lp, err := loop.New(loop.NoAlternateScreen, loop.NoRestoreColors, loop.NoMouseTracking)
	if err != nil {
		return
	}
	current_text := strings.Builder{}
	data := strings.Builder{}
	data.Grow(4096)
	started := false
	canceled := false
	update_type := ""

	handle_line := func(line string) error {
		if canceled {
			return nil
		}
		if started {
			if update_type == "" {
				update_type = line
			} else {
				if line == "KITTY_DATA_END" {
					lp.QueueWriteString(update_type + "\r\n")
					if update_type == "DONE" {
						lp.Quit(0)
						return nil
					}
					b, err := base64.StdEncoding.DecodeString(data.String())
					data.Reset()
					data.Grow(4096)
					started = false
					if err == nil {
						err = on_data(update_type, b)
					}
					update_type = ""
					if err != nil {
						return err
					}
				} else {
					data.WriteString(line)
				}
			}
		} else {
			if line == "KITTY_DATA_START" {
				started = true
				update_type = ""
			}
		}
		return nil
	}

	check_for_line := func() error {
		if canceled {
			return nil
		}
		s := current_text.String()
		for {
			idx := strings.Index(s, "\n")
			if idx < 0 {
				break
			}
			err = handle_line(s[:idx])
			if err != nil {
				return err
			}
			s = s[idx+1:]
		}
		current_text.Reset()
		current_text.Grow(4096)
		if s != "" {
			current_text.WriteString(s)
		}
		return nil
	}

	lp.OnInitialize = func() (string, error) {
		pos, chunk_num := 0, 0
		for {
			limit := utils.Min(pos+2048, len(data_to_send))
			if limit <= pos {
				break
			}
			lp.QueueWriteString("\x1bP@kitty-edit|" + strconv.Itoa(chunk_num) + ":")
			lp.QueueWriteString(data_to_send[pos:limit])
			lp.QueueWriteString("\x1b\\")
			chunk_num++
			pos = limit
		}
		lp.QueueWriteString("\x1bP@kitty-edit|\x1b\\")
		return "", nil
	}

	lp.OnText = func(text string, from_key_event bool, in_bracketed_paste bool) error {
		if !from_key_event {
			current_text.WriteString(text)
			err = check_for_line()
			if err != nil {
				return err
			}
		}
		return nil
	}

	const abort_msg = "\x1bP@kitty-edit|0:abort_signaled=interrupt\x1b\\\x1bP@kitty-edit|\x1b\\"

	lp.OnKeyEvent = func(event *loop.KeyEvent) error {
		if event.MatchesPressOrRepeat("ctrl+c") || event.MatchesPressOrRepeat("esc") {
			event.Handled = true
			canceled = true
			lp.QueueWriteString(abort_msg)
			if !started {
				return tui.Canceled
			}
		}
		return nil
	}

	err = lp.Run()
	if err != nil {
		return
	}
	if canceled {
		return tui.Canceled
	}

	ds := lp.DeathSignalName()
	if ds != "" {
		fmt.Print(abort_msg)
		if kill_if_signaled {
			lp.KillIfSignalled()
			return
		}
		return &tui.KilledBySignal{Msg: fmt.Sprint("Killed by signal: ", ds), SignalName: ds}
	}
	return
}

func edit_in_kitty(path string) (err error) {
	read_file, err := os.Open(path)
	if err != nil {
		return fmt.Errorf("Failed to open %s for reading with error: %w", path, err)
	}
	defer read_file.Close()
	var s unix.Stat_t
	err = unix.Fstat(int(read_file.Fd()), &s)
	if err != nil {
		return fmt.Errorf("Failed to stat %s with error: %w", path, err)
	}
	if s.Size > 8*1024*1024 {
		return fmt.Errorf("File size %s is too large for performant editing", humanize.Bytes(uint64(s.Size)))
	}

	file_data, err := io.ReadAll(read_file)
	if err != nil {
		return fmt.Errorf("Failed to read from %s with error: %w", path, err)
	}
	read_file.Close()
	data := strings.Builder{}
	data.Grow(len(file_data) * 4)

	add := func(key, val string) {
		if data.Len() > 0 {
			data.WriteString(",")
		}
		data.WriteString(key)
		data.WriteString("=")
		data.WriteString(val)
	}
	add_encoded := func(key, val string) { add(key, encode(val)) }

	if unix.Access(path, unix.R_OK|unix.W_OK) != nil {
		return fmt.Errorf("%s is not readable and writeable", path)
	}
	cwd, err := os.Getwd()
	if err != nil {
		return fmt.Errorf("Failed to get the current working directory with error: %w", err)
	}
	add_encoded("cwd", cwd)
	for _, arg := range os.Args[2:] {
		add_encoded("a", arg)
	}
	add("file_inode", fmt.Sprintf("%d:%d:%d", s.Dev, s.Ino, s.Mtim.Nano()))
	add_encoded("file_data", utils.UnsafeBytesToString(file_data))
	fmt.Println("Waiting for editing to be completed, press Esc to abort...")
	write_data := func(data_type string, rdata []byte) (err error) {
		err = utils.AtomicWriteFile(path, rdata, fs.FileMode(s.Mode).Perm())
		if err != nil {
			err = fmt.Errorf("Failed to write data to %s with error: %w", path, err)
		}
		return
	}
	err = edit_loop(data.String(), true, write_data)
	if err != nil {
		if err == tui.Canceled {
			return err
		}
		return fmt.Errorf("Failed to receive edited file back from terminal with error: %w", err)
	}
	return
}

func EntryPoint(parent *cli.Command) *cli.Command {
	sc := parent.AddSubCommand(&cli.Command{
		Name:             "edit-in-kitty",
		Usage:            "edit-in-kitty [options] file-to-edit",
		ShortDescription: "Edit a file in a kitty overlay window",
		HelpText: "Edit the specified file in a kitty overlay window. Works over SSH as well.\n\n" +
			"For usage instructions see: https://sw.kovidgoyal.net/kitty/shell-integration/#edit-file",
		Run: func(cmd *cli.Command, args []string) (ret int, err error) {
			if len(args) == 0 {
				fmt.Fprintln(os.Stderr, "Usage:", cmd.Usage)
				return 1, fmt.Errorf("No file to edit specified.")
			}
			if len(args) != 1 {
				fmt.Fprintln(os.Stderr, "Usage:", cmd.Usage)
				return 1, fmt.Errorf("Only one file to edit must be specified")
			}
			err = edit_in_kitty(args[0])
			return 0, err
		},
	})
	AddCloneSafeOpts(sc)
	return sc
}