package main

import (
	"fmt"
	"io"
	"os"
	"os/exec"
	"strings"
	"time"
)

func main() {
	mode := "both"
	if len(os.Args) > 1 {
		mode = os.Args[1]
	}
	if mode != "both" && mode != "osc" && mode != "cpr" {
		fmt.Fprintf(os.Stderr, "usage: %s [both|osc|cpr]\n", os.Args[0])
		os.Exit(2)
	}

	tty, err := os.OpenFile("/dev/tty", os.O_RDWR, 0)
	if err != nil {
		fmt.Fprintf(os.Stderr, "open /dev/tty: %v\n", err)
		os.Exit(1)
	}
	defer tty.Close()

	oldState, err := stty(tty, "-g")
	if err != nil {
		fmt.Fprintf(os.Stderr, "stty -g: %v\n", err)
		os.Exit(1)
	}
	defer func() { _, _ = stty(tty, oldState) }()

	if _, err := stty(tty, "-echo", "-icanon", "min", "0", "time", "2"); err != nil {
		fmt.Fprintf(os.Stderr, "stty raw-ish: %v\n", err)
		os.Exit(1)
	}

	query, description := queryFor(mode)
	fmt.Printf("pid=%d TERM=%q TMUX_set=%v mode=%s\n", os.Getpid(), os.Getenv("TERM"), os.Getenv("TMUX") != "", mode)
	fmt.Println("sending " + description)

	if _, err := tty.WriteString(query); err != nil {
		fmt.Fprintf(os.Stderr, "write query: %v\n", err)
		os.Exit(1)
	}

	deadline := time.Now().Add(2 * time.Second)
	var got []byte
	buf := make([]byte, 256)
	for time.Now().Before(deadline) {
		n, err := tty.Read(buf)
		if n > 0 {
			got = append(got, buf[:n]...)
			if completeFor(mode, got) {
				break
			}
		}
		if err != nil && err != io.EOF {
			fmt.Fprintf(os.Stderr, "read response: %v\n", err)
			os.Exit(1)
		}
		time.Sleep(20 * time.Millisecond)
	}

	fmt.Printf("read %d bytes from /dev/tty\n", len(got))
	fmt.Printf("raw: %q\n", got)
	fmt.Printf("escaped: %s\n", escape(got))
	fmt.Printf("OSC 11 response count: %d\n", countOSC11(got))
	fmt.Printf("cursor-position response count: %d\n", countCPR(got))
	fmt.Printf("has OSC 11 response: %v\n", hasOSC11(got))
	fmt.Printf("has cursor-position response: %v\n", hasCPR(got))
	fmt.Println("after this exits, type Enter. If query text appears at the next prompt, the response leaked.")
}

func queryFor(mode string) (string, string) {
	switch mode {
	case "osc":
		return "\x1b]11;?\x1b\\", "OSC 11 background-color query only"
	case "cpr":
		return "\x1b[6n", "CSI 6n cursor-position query only"
	default:
		return "\x1b]11;?\x1b\\\x1b[6n", "OSC 11 background-color query and CSI 6n cursor-position query"
	}
}

func completeFor(mode string, b []byte) bool {
	switch mode {
	case "osc":
		return hasOSC11(b)
	case "cpr":
		return hasCPR(b)
	default:
		return hasOSC11(b) && hasCPR(b)
	}
}

func stty(tty *os.File, args ...string) (string, error) {
	cmd := exec.Command("stty", args...)
	cmd.Stdin = tty
	cmd.Stderr = os.Stderr
	out, err := cmd.Output()
	return strings.TrimSpace(string(out)), err
}

func hasOSC11(b []byte) bool {
	return countOSC11(b) > 0
}

func hasCPR(b []byte) bool {
	return countCPR(b) > 0
}

func countOSC11(b []byte) int {
	return strings.Count(string(b), "\x1b]11;")
}

func countCPR(b []byte) int {
	s := string(b)
	count := 0
	for i := 0; i < len(s); i++ {
		if !strings.HasPrefix(s[i:], "\x1b[") {
			continue
		}
		j := i + len("\x1b[")
		for j < len(s) && (s[j] == ';' || (s[j] >= '0' && s[j] <= '9')) {
			j++
		}
		if j < len(s) && s[j] == 'R' {
			count++
			i = j
		}
	}
	return count
}

func escape(b []byte) string {
	var out strings.Builder
	for _, c := range b {
		switch c {
		case '\x1b':
			out.WriteString("<ESC>")
		case '\a':
			out.WriteString("<BEL>")
		case '\r':
			out.WriteString("<CR>")
		case '\n':
			out.WriteString("<LF>")
		default:
			if c < 0x20 || c == 0x7f {
				fmt.Fprintf(&out, "<0x%02x>", c)
			} else {
				out.WriteByte(c)
			}
		}
	}
	return out.String()
}