path: root/internal/color/color.go

package color

import (
	"fmt"
	"math"
	"regexp"
	"strconv"
	"strings"
	"unicode/utf8"
)

const Reset = "\033[0m"

var ansiRe = regexp.MustCompile(`\x1b\[[0-9;]*[a-zA-Z]`)

var ansi16RGB = [16][3]int{
	{0, 0, 0},       // 0  black
	{170, 0, 0},     // 1  red
	{0, 170, 0},     // 2  green
	{170, 170, 0},   // 3  yellow
	{0, 0, 170},     // 4  blue
	{170, 0, 170},   // 5  magenta
	{0, 170, 170},   // 6  cyan
	{170, 170, 170}, // 7  white
	{85, 85, 85},    // 8  bright black
	{255, 85, 85},   // 9  bright red
	{85, 255, 85},   // 10 bright green
	{255, 255, 85},  // 11 bright yellow
	{85, 85, 255},   // 12 bright blue
	{255, 85, 255},  // 13 bright magenta
	{85, 255, 255},  // 14 bright cyan
	{255, 255, 255}, // 15 bright white
}

var cubeValues = [6]int{0, 95, 135, 175, 215, 255}

func Xterm256ToRGB(index int) (int, int, int) {
	if index < 16 {
		return ansi16RGB[index][0], ansi16RGB[index][1], ansi16RGB[index][2]
	}
	if index < 232 {
		i := index - 16
		b := cubeValues[i%6]
		g := cubeValues[(i/6)%6]
		r := cubeValues[(i/36)%6]
		return r, g, b
	}
	v := 8 + (index-232)*10
	return v, v, v
}

func colorDist(r1, g1, b1, r2, g2, b2 int) int {
	dr := r1 - r2
	dg := g1 - g2
	db := b1 - b2
	return dr*dr + dg*dg + db*db
}

func nearestANSI(index int) int {
	if index < 16 {
		if index < 8 {
			return 30 + index
		}
		return 90 + (index - 8)
	}
	r, g, b := Xterm256ToRGB(index)
	bestDist := math.MaxInt
	bestCode := 37
	for i := 0; i < 16; i++ {
		dist := colorDist(r, g, b, ansi16RGB[i][0], ansi16RGB[i][1], ansi16RGB[i][2])
		if dist < bestDist {
			bestDist = dist
			if i < 8 {
				bestCode = 30 + i
			} else {
				bestCode = 90 + (i - 8)
			}
		}
	}
	return bestCode
}

func nearestANSIBg(index int) int {
	fg := nearestANSI(index)
	if fg >= 90 {
		return fg - 90 + 100
	}
	return fg - 30 + 40
}

func nearestCube(v int) int {
	if v < 48 {
		return 0
	}
	if v < 115 {
		return 1
	}
	if v < 155 {
		return 2
	}
	if v < 195 {
		return 3
	}
	if v < 235 {
		return 4
	}
	return 5
}

func NearestXterm256(r, g, b int) int {
	ri := nearestCube(r)
	gi := nearestCube(g)
	bi := nearestCube(b)
	cubeIdx := 16 + ri*36 + gi*6 + bi
	cubeDist := colorDist(r, g, b, cubeValues[ri], cubeValues[gi], cubeValues[bi])

	avg := (r + g + b) / 3
	grayStep := (avg - 8 + 5) / 10
	if grayStep < 0 {
		grayStep = 0
	}
	if grayStep > 23 {
		grayStep = 23
	}
	grayIdx := 232 + grayStep
	gv := 8 + grayStep*10
	grayDist := colorDist(r, g, b, gv, gv, gv)

	bestIdx := cubeIdx
	bestDist := cubeDist
	if grayDist < bestDist {
		bestIdx = grayIdx
		bestDist = grayDist
	}

	for i := 0; i < 16; i++ {
		d := colorDist(r, g, b, ansi16RGB[i][0], ansi16RGB[i][1], ansi16RGB[i][2])
		if d < bestDist {
			bestDist = d
			bestIdx = i
		}
	}

	return bestIdx
}

func FgCode(mode string, index int) string {
	if mode == "none" || mode == "" || index < 0 || index > 255 {
		return ""
	}
	if mode == "xterm256" {
		return fmt.Sprintf("\033[38;5;%dm", index)
	}
	return fmt.Sprintf("\033[%dm", nearestANSI(index))
}

func BgCode(mode string, index int) string {
	if mode == "none" || mode == "" || index < 0 || index > 255 {
		return ""
	}
	if mode == "xterm256" {
		return fmt.Sprintf("\033[48;5;%dm", index)
	}
	return fmt.Sprintf("\033[%dm", nearestANSIBg(index))
}

func StyleCode(name string) string {
	switch name {
	case "bold":
		return "\033[1m"
	case "dim":
		return "\033[2m"
	case "underline":
		return "\033[4m"
	}
	return ""
}

func VisibleLen(s string) int {
	return utf8.RuneCountInString(ansiRe.ReplaceAllString(s, ""))
}

// Ellipsis returns a Unicode ellipsis character when unicode is true,
// or the ASCII "..." otherwise.
func Ellipsis(unicode bool) string {
	if unicode {
		return "\u2026"
	}
	return "..."
}

// TruncateVisible truncates s to at most maxVisible visible characters,
// preserving ANSI SGR codes. Assumes cells follow the common pattern
// "\x1b[...mTEXT\x1b[0m" or are plain text. Appends an ellipsis when truncated.
func TruncateVisible(s string, maxVisible int, unicode bool) string {
	ellipsis := Ellipsis(unicode)
	ellipLen := utf8.RuneCountInString(ellipsis)
	if maxVisible < ellipLen {
		maxVisible = ellipLen
	}
	if VisibleLen(s) <= maxVisible {
		return s
	}
	var pre, suf string
	rest := s
	if idx := strings.IndexByte(rest, 'm'); idx >= 0 && strings.HasPrefix(rest, "\x1b[") {
		pre = rest[:idx+1]
		rest = rest[idx+1:]
	}
	if strings.HasSuffix(rest, Reset) {
		suf = Reset
		rest = rest[:len(rest)-len(Reset)]
	}
	runes := []rune(rest)
	budget := maxVisible - ellipLen
	if budget < 0 {
		budget = 0
	}
	if len(runes) > budget {
		rest = string(runes[:budget]) + ellipsis
	}
	return pre + rest + suf
}

// WrapANSI wraps text to the given visible width, ignoring ANSI color codes
// when measuring. Each input line (split on "\n") is handled on its own: a line
// that already fits is returned untouched, so indentation and aligned UI are
// preserved. Only over-long lines are word-wrapped. The result is joined with
// "\r\n". A width <= 0 disables wrapping.
func WrapANSI(s string, width int) string {
	if width <= 0 {
		return s
	}
	var out []string
	for _, line := range strings.Split(s, "\n") {
		line = strings.TrimRight(line, "\r")
		out = append(out, wrapLine(line, width)...)
	}
	return strings.Join(out, "\r\n")
}

func wrapLine(line string, width int) []string {
	if VisibleLen(line) <= width {
		return []string{line}
	}
	words := strings.Fields(line)
	if len(words) == 0 {
		return []string{line}
	}
	cur := words[0]
	curWidth := VisibleLen(words[0])
	var lines []string
	for _, w := range words[1:] {
		ww := VisibleLen(w)
		if curWidth+1+ww <= width {
			cur += " " + w
			curWidth += 1 + ww
		} else {
			lines = append(lines, cur)
			cur = w
			curWidth = ww
		}
	}
	return append(lines, cur)
}

func ContrastFg(mode string, bgIndex int) string {
	r, g, b := Xterm256ToRGB(bgIndex)
	if r+g+b > 384 {
		return FgCode(mode, 0)
	}
	return FgCode(mode, 15)
}

type ColorSpec struct {
	Fg        int
	Bg        int
	Bold      bool
	Dim       bool
	Underline bool
	Gradient  []int
}

func NoColor() ColorSpec {
	return ColorSpec{Fg: -1, Bg: -1}
}

// Average returns a ColorSpec whose foreground is the RGB midpoint of the two
// inputs' foregrounds. If only one input has a foreground, that one is used; if
// neither does, NoColor is returned. Styles and background are not carried over.
func Average(a, b ColorSpec) ColorSpec {
	if a.Fg < 0 && b.Fg < 0 {
		return NoColor()
	}
	if a.Fg < 0 {
		return ColorSpec{Fg: b.Fg, Bg: -1}
	}
	if b.Fg < 0 {
		return ColorSpec{Fg: a.Fg, Bg: -1}
	}
	r1, g1, b1 := Xterm256ToRGB(a.Fg)
	r2, g2, b2 := Xterm256ToRGB(b.Fg)
	return ColorSpec{Fg: NearestXterm256((r1+r2)/2, (g1+g2)/2, (b1+b2)/2), Bg: -1}
}

func Parse(input string) ColorSpec {
	spec := ColorSpec{Fg: -1, Bg: -1}
	for _, token := range strings.Fields(input) {
		switch {
		case token == "bold":
			spec.Bold = true
		case token == "dim":
			spec.Dim = true
		case token == "underline":
			spec.Underline = true
		case strings.HasPrefix(token, "bg:"):
			if n, ok := parseColorIndex(strings.TrimPrefix(token, "bg:")); ok {
				spec.Bg = n
			}
		case strings.HasPrefix(token, "g:"):
			parts := strings.Split(token[2:], ",")
			var stops []int
			for _, p := range parts {
				if n, ok := parseColorIndex(strings.TrimSpace(p)); ok {
					stops = append(stops, n)
				}
			}
			if len(stops) >= 2 {
				spec.Gradient = stops
			}
		default:
			if n, ok := parseColorIndex(token); ok {
				spec.Fg = n
			}
		}
	}
	return spec
}

// parseColorIndex parses a 1-2 digit hexadecimal color index (00-FF, case
// insensitive) into its 0-255 value. Out-of-range or non-hex input fails.
func parseColorIndex(s string) (int, bool) {
	n, err := strconv.ParseUint(s, 16, 0)
	if err != nil || n > 255 {
		return 0, false
	}
	return int(n), true
}

func (s ColorSpec) Empty() bool {
	return s.Fg < 0 && s.Bg < 0 && !s.Bold && !s.Dim && !s.Underline && len(s.Gradient) == 0
}

func stylePrefix(spec ColorSpec) string {
	var codes []string
	if spec.Bold {
		codes = append(codes, StyleCode("bold"))
	}
	if spec.Dim {
		codes = append(codes, StyleCode("dim"))
	}
	if spec.Underline {
		codes = append(codes, StyleCode("underline"))
	}
	return strings.Join(codes, "")
}

func Render(mode string, spec ColorSpec, text string) string {
	if mode == "none" || mode == "" || spec.Empty() {
		return text
	}
	if len(spec.Gradient) >= 2 {
		return renderGradient(mode, spec, text)
	}
	var codes []string
	prefix := stylePrefix(spec)
	if prefix != "" {
		codes = append(codes, prefix)
	}
	if spec.Fg >= 0 {
		codes = append(codes, FgCode(mode, spec.Fg))
	}
	if spec.Bg >= 0 {
		codes = append(codes, BgCode(mode, spec.Bg))
	}
	if len(codes) == 0 {
		return text
	}
	return strings.Join(codes, "") + text + Reset
}

func interpolateGradient(stops []int, t float64) int {
	if t <= 0 {
		return stops[0]
	}
	if t >= 1 {
		return stops[len(stops)-1]
	}
	seg := t * float64(len(stops)-1)
	i := int(seg)
	if i >= len(stops)-1 {
		i = len(stops) - 2
	}
	frac := seg - float64(i)
	r1, g1, b1 := Xterm256ToRGB(stops[i])
	r2, g2, b2 := Xterm256ToRGB(stops[i+1])
	r := r1 + int(frac*float64(r2-r1))
	g := g1 + int(frac*float64(g2-g1))
	b := b1 + int(frac*float64(b2-b1))
	return NearestXterm256(r, g, b)
}

func renderGradient(mode string, spec ColorSpec, text string) string {
	runes := []rune(text)
	n := len(runes)
	if n == 0 {
		return ""
	}
	var sb strings.Builder
	prefix := stylePrefix(spec)
	if spec.Bg >= 0 {
		sb.WriteString(BgCode(mode, spec.Bg))
	}
	for i, r := range runes {
		t := 0.0
		if n > 1 {
			t = float64(i) / float64(n-1)
		}
		idx := interpolateGradient(spec.Gradient, t)
		if prefix != "" {
			sb.WriteString(prefix)
		}
		sb.WriteString(FgCode(mode, idx))
		sb.WriteRune(r)
	}
	sb.WriteString(Reset)
	return sb.String()
}

func ExpandTags(mode string, text string) string {
	return ExpandTagsDefault(mode, NoColor(), text)
}

func ExpandTagsDefault(mode string, def ColorSpec, text string) string {
	return ExpandTagsDefaultNamed(mode, def, nil, text)
}

// ExpandTagsNamed expands {tags} in text, additionally resolving a tag whose
// spec matches a key in named (a color category name) to that category's
// ColorSpec. Named hits are treated as real tags even when they resolve to an
// empty spec, so a category disabled with "off" renders its text plain instead
// of leaking the literal tag. Any tag not present in named falls back to the
// ordinary hex/style spec parser.
func ExpandTagsNamed(mode string, named map[string]ColorSpec, text string) string {
	return ExpandTagsDefaultNamed(mode, NoColor(), named, text)
}

func ExpandTagsDefaultNamed(mode string, def ColorSpec, named map[string]ColorSpec, text string) string {
	if !strings.Contains(text, "{") {
		if !def.Empty() {
			return Render(mode, def, text)
		}
		return text
	}

	var sb strings.Builder
	pos := 0
	for pos < len(text) {
		start := strings.Index(text[pos:], "{")
		if start < 0 {
			remaining := text[pos:]
			if remaining != "" {
				if !def.Empty() {
					sb.WriteString(Render(mode, def, remaining))
				} else {
					sb.WriteString(remaining)
				}
			}
			break
		}
		start += pos

		end := strings.Index(text[start+1:], "}")
		if end < 0 {
			remaining := text[pos:]
			if !def.Empty() {
				sb.WriteString(Render(mode, def, remaining))
			} else {
				sb.WriteString(remaining)
			}
			break
		}
		end += start + 1

		specStr := text[start+1 : end]

		if specStr == "" || specStr == "/" {
			before := text[pos : start+end-start+1]
			if !def.Empty() {
				sb.WriteString(Render(mode, def, before))
			} else {
				sb.WriteString(before)
			}
			pos = end + 1
			continue
		}

		spec, isNamed := NoColor(), false
		if named != nil {
			if s, ok := named[specStr]; ok {
				spec, isNamed = s, true
			}
		}
		if !isNamed {
			spec = Parse(specStr)
		}
		if spec.Empty() && !isNamed {
			before := text[pos : end+1]
			if !def.Empty() {
				sb.WriteString(Render(mode, def, before))
			} else {
				sb.WriteString(before)
			}
			pos = end + 1
			continue
		}

		closeStart := strings.Index(text[end+1:], "{/}")
		if closeStart < 0 {
			remaining := text[pos:]
			if !def.Empty() {
				sb.WriteString(Render(mode, def, remaining))
			} else {
				sb.WriteString(remaining)
			}
			break
		}
		closeStart += end + 1

		before := text[pos:start]
		if before != "" {
			if !def.Empty() {
				sb.WriteString(Render(mode, def, before))
			} else {
				sb.WriteString(before)
			}
		}

		inner := text[end+1 : closeStart]
		sb.WriteString(Render(mode, spec, inner))

		pos = closeStart + 3
	}

	return sb.String()
}

func ExpandDialogTags(mode string, dialogSpec ColorSpec, text string) string {
	return ExpandDialogTagsNamed(mode, dialogSpec, nil, text)
}

// ExpandDialogTagsNamed is ExpandDialogTags with named-color tag resolution.
func ExpandDialogTagsNamed(mode string, dialogSpec ColorSpec, named map[string]ColorSpec, text string) string {
	parts := strings.Split(text, `"`)
	var sb strings.Builder
	for i, part := range parts {
		if i%2 == 0 {
			if part != "" {
				sb.WriteString(ExpandTagsDefaultNamed(mode, NoColor(), named, part))
			}
		} else {
			if i < len(parts)-1 {
				sb.WriteString(ExpandTagsDefaultNamed(mode, dialogSpec, named, `"`+part+`"`))
			} else {
				sb.WriteString(ExpandTagsDefaultNamed(mode, dialogSpec, named, `"`+part))
			}
		}
	}
	return sb.String()
}