// package termcolor carries the terminal color query flow: parsing a real
// terminal's osc 4/10/11/12 replies, the attach payload section that carries
// them to the owner, and the agreement rule the owner applies before
// answering a query from an app inside the session.
package termcolor

import (
	"encoding/binary"
	"fmt"
	"maps"
	"slices"
	"strconv"
	"strings"
)

// color keys follow the xterm index space: 0-255 is the palette, 256 and up
// are the special slots.
const (
	Foreground = 256
	Background = 257
	Cursor     = 258
)

type Color struct{ R, G, B uint16 }

// ProbeQuery asks the terminal for its foreground, background, cursor, the
// first 16 palette entries, and a da1 sentinel to bound the wait.
var ProbeQuery = buildProbeQuery()

func buildProbeQuery() string {
	q := "\x1b]10;?\x1b\\\x1b]11;?\x1b\\\x1b]12;?\x1b\\"
	pairs := make([]string, 0, 32)
	for i := range 16 {
		pairs = append(pairs, strconv.Itoa(i), "?")
	}
	return q + "\x1b]4;" + strings.Join(pairs, ";") + "\x1b\\\x1b[c"
}

// ParseReplies extracts color reports from a terminal's replies to ProbeQuery.
// malformed or unrelated escape sequences are ignored.
func ParseReplies(b []byte) map[int]Color {
	colors := map[int]Color{}
	for i := 0; i+1 < len(b); {
		if b[i] != 0x1b || b[i+1] != ']' {
			i++
			continue
		}
		body, next, ok := oscSpan(b, i)
		if !ok {
			if next >= len(b) {
				break // incomplete tail
			}
			i = next
			continue
		}
		if k, c, ok := parseReply(body); ok {
			colors[k] = c
		}
		i = next
	}
	return colors
}

// oscSpan reads one osc string starting at b[i] (esc ]) and returns its body
// plus the index just past it. BEL and ST both terminate.
func oscSpan(b []byte, i int) (body string, next int, ok bool) {
	for j := i + 2; j < len(b); j++ {
		switch {
		case b[j] == '\a':
			return string(b[i+2 : j]), j + 1, true
		case b[j] == 0x1b && j+1 < len(b) && b[j+1] == '\\':
			return string(b[i+2 : j]), j + 2, true
		case b[j] == 0x1b:
			return "", j, false
		}
	}
	return "", len(b), false
}

func parseReply(body string) (int, Color, bool) {
	parts := strings.Split(body, ";")
	switch parts[0] {
	case "4":
		if len(parts) != 3 {
			return 0, Color{}, false
		}
		i, err := strconv.Atoi(parts[1])
		if err != nil || i < 0 || i >= Foreground {
			return 0, Color{}, false
		}
		c, ok := parseColor(parts[2])
		return i, c, ok
	case "10", "11", "12":
		if len(parts) != 2 {
			return 0, Color{}, false
		}
		n, _ := strconv.Atoi(parts[0]) // one of the case labels
		c, ok := parseColor(parts[1])
		return Foreground + n - 10, c, ok
	}
	return 0, Color{}, false
}

func parseColor(s string) (Color, bool) {
	spec, ok := strings.CutPrefix(s, "rgb:")
	if !ok {
		return Color{}, false
	}
	parts := strings.Split(spec, "/")
	if len(parts) != 3 {
		return Color{}, false
	}
	var ch [3]uint16
	for i, p := range parts {
		if len(p) < 1 || len(p) > 4 {
			return Color{}, false
		}
		v, err := strconv.ParseUint(p, 16, 32)
		if err != nil {
			return Color{}, false
		}
		max := uint64(1)<<(4*len(p)) - 1
		ch[i] = uint16(v * 0xffff / max)
	}
	return Color{ch[0], ch[1], ch[2]}, true
}

// DA1End returns the index just past a device attributes reply, the probe's
// sentinel, or -1 when none has arrived yet.
func DA1End(b []byte) int {
	for i := 0; i+1 < len(b); i++ {
		if b[i] != 0x1b || b[i+1] != '[' {
			continue
		}
		j := i + 2
		for j < len(b) && b[j] >= 0x20 && b[j] <= 0x3f {
			j++
		}
		if j < len(b) && b[j] == 'c' {
			return j + 1
		}
	}
	return -1
}

const colorsVersion = 1

// EncodeColors packs a color map as the attach payload section that follows
// the four size bytes. an empty map encodes to nothing, which leaves the
// payload in the legacy four byte shape.
func EncodeColors(colors map[int]Color) []byte {
	if len(colors) == 0 {
		return nil
	}
	keys := slices.Sorted(maps.Keys(colors))
	b := []byte{colorsVersion, byte(len(keys))}
	for _, k := range keys {
		c := colors[k]
		for _, v := range []uint16{uint16(k), c.R, c.G, c.B} {
			b = binary.BigEndian.AppendUint16(b, v)
		}
	}
	return b
}

// DecodeColors reads the colors section of an attach payload. ok is false for
// legacy payloads that carry no section.
func DecodeColors(b []byte) (map[int]Color, bool) {
	if len(b) < 2 || b[0] != colorsVersion {
		return nil, false
	}
	n := int(b[1])
	if len(b) != 2+8*n {
		return nil, false
	}
	colors := make(map[int]Color, n)
	for i := range n {
		o := 2 + 8*i
		colors[int(binary.BigEndian.Uint16(b[o:]))] = Color{
			R: binary.BigEndian.Uint16(b[o+2:]),
			G: binary.BigEndian.Uint16(b[o+4:]),
			B: binary.BigEndian.Uint16(b[o+6:]),
		}
	}
	return colors, true
}

// Agreed answers a color query only when every viewer reported the color and
// all reported the same value. zero viewers, a missing color, or a mismatch
// all leave the query unanswered.
func Agreed(index int, sets []map[int]Color) (Color, bool) {
	var first Color
	for i, set := range sets {
		c, ok := set[index]
		if !ok {
			return Color{}, false
		}
		if i > 0 && c != first {
			return Color{}, false
		}
		first = c
	}
	return first, len(sets) > 0
}

// Reply formats the osc response an owner writes for a color query.
func Reply(index int, c Color) string {
	if index >= Foreground {
		return fmt.Sprintf("\x1b]%d;rgb:%04x/%04x/%04x\x1b\\", 10+index-Foreground, c.R, c.G, c.B)
	}
	return fmt.Sprintf("\x1b]4;%d;rgb:%04x/%04x/%04x\x1b\\", index, c.R, c.G, c.B)
}