package tui

import (
	"strings"
)

// stripEscapes removes terminal control sequences: CSI (colors, cursor
// moves), OSC (titles, hyperlinks, clipboard), DCS and the other string
// sequences, and two-byte ESC codes. With keepColors it keeps the CSI
// sequences that only set colors and styles (SGR, ESC [ ... m) and erase
// to the end of the line (ESC [ K).
//
// A string sequence ends at its terminator, or at an ESC, as in a
// terminal; but at the end of its line at the latest: one left open (or
// closed lines later) would otherwise take all the text after it out of
// sight, so then only its opening is removed.
func stripEscapes(s string, keepColors bool) string {
	if !strings.Contains(s, "\x1b") {
		return s
	}
	var b strings.Builder
	for i := 0; i < len(s); {
		if s[i] != 0x1b {
			b.WriteByte(s[i])
			i++
			continue
		}
		start := i
		i++ // ESC
		if i >= len(s) {
			break
		}
		switch c := s[i]; c {
		case '[': // CSI: parameters and intermediates, then a final byte
			i++
			for i < len(s) && s[i] >= 0x20 && s[i] <= 0x3f {
				i++
			}
			if i < len(s) && s[i] >= 0x40 && s[i] <= 0x7e {
				params := s[start+2 : i]
				color := s[i] == 'm' && strings.Trim(params, "0123456789;:") == "" || s[i] == 'K' && (params == "" || params == "0")
				i++
				if keepColors && color {
					b.WriteString(s[start:i])
				}
			}
		case ']', 'P', 'X', '^', '_': // strings ended by BEL (OSC) or ESC \
			i++
			for j := i; j < len(s) && s[j] != '\n'; j++ {
				if c == ']' && s[j] == 0x07 {
					i = j + 1
					break
				}
				if s[j] == 0x1b {
					i = j // a new sequence, ST or not, ends the string
					if j+1 < len(s) && s[j+1] == '\\' {
						i = j + 2
					}
					break
				}
			}
		default: // a two-byte sequence
			if c >= 0x20 && c <= 0x7e {
				i++
			} // a control character after ESC acts on its own, as in a terminal
		}
	}
	return b.String()
}

// Sanitize makes text from the model, a file or a command safe to show:
// without terminal control sequences or control characters, so it cannot
// move the cursor, recolor the screen or draw a fake prompt. Tabs become
// spaces; newlines stay.
func Sanitize(s string) string {
	s = strings.ReplaceAll(stripEscapes(s, false), "\r\n", "\n")
	var b strings.Builder
	b.Grow(len(s))
	for _, r := range s {
		switch {
		case r == '\n':
			b.WriteRune(r)
		case r == '\t':
			b.WriteString("    ")
		case r < 0x20 || r == 0x7f || (r >= 0x80 && r < 0xa0):
			// Control characters, C1 controls included.
		case r >= '\u202a' && r <= '\u202e', r >= '\u2066' && r <= '\u2069', r == '\u200e', r == '\u200f', r == '\u061c',
			r >= '\U000e0000' && r <= '\U000e007f':
			// Bidirectional controls, which can make text read in a
			// different order than it runs, and tag characters, which
			// spell invisible text.
			b.WriteRune('\ufffd')
		default:
			b.WriteRune(r)
		}
	}
	return b.String()
}

// SanitizeColors is Sanitize for the output of a program the user chose to
// color text, such as a diff renderer: it keeps the sequences that only
// set colors and styles (SGR, ESC [ ... m) and erase to the end of the line
// (ESC [ K, which fills a line's background), and removes everything else
// as Sanitize does.
func SanitizeColors(s string) string {
	s = stripEscapes(s, true)
	// The only ESCs left start the kept sequences: Sanitize the text
	// between them.
	var b strings.Builder
	for {
		i := strings.IndexByte(s, 0x1b)
		if i < 0 {
			break
		}
		j := i + 2
		for j < len(s) && s[j] < 0x40 {
			j++
		}
		if j >= len(s) {
			break // cannot happen: stripEscapes keeps whole sequences only
		}
		b.WriteString(Sanitize(s[:i]))
		b.WriteString(s[i : j+1])
		s = s[j+1:]
	}
	b.WriteString(Sanitize(s))
	return b.String()
}