guesswidth.go

// Package guesswidth handles the format as formatted by printf.
// Spaces exist as delimiters, but spaces are not always delimiters.
// The width seems to be a fixed length, but it doesn't always fit.
// guesswidth finds the column separation position
// from the reference line(header) and multiple lines(body).
package guesswidth

import (
	"bufio"
	"fmt"
	"io"
	"runtime/debug"
	"sort"
	"strings"
	"unicode"
	"unicode/utf8"

	"github.com/rivo/uniseg"
)

// GuessWidth reads records from printf-like output.
type GuessWidth struct {
	reader *bufio.Reader
	// pos is a list of separator positions.
	pos []int
	// Widths is the width of the column.
	Widths []Cols
	// preLines stores the lines read for scan.
	preLines []string
	// preCount is the number returned by read.
	preCount int

	// ScanNum is the number to scan to analyze.
	ScanNum int
	// Header is the base line number. It starts from 0.
	Header int
	// limitSplit is the maximum number of columns to split.
	LimitSplit int
	// MinLines is the minimum number of lines to recognize as a separator.
	// 1 if only the header, 2 or more if there is a blank in the body.
	MinLines int
	// TrimSpace is whether to trim the space in the value.
	TrimSpace bool
}

// Cols is the width and justification of the column.
type Cols struct {
	Width      int // Width is the width of the column.
	Justified  int // Justified indicates the justification of the column (Left or Right).
	rightCount int // rightCount is the number of times the column appears to be right-aligned.
}

// Justification constants for columns.
const (
	Left  = iota // Left indicates left alignment.
	Right        // Right indicates right alignment.

)

// NewReader returns a new Reader that reads from r.
func NewReader(r io.Reader) *GuessWidth {
	reader := bufio.NewReader(r)
	g := &GuessWidth{
		reader:     reader,
		ScanNum:    100,
		preCount:   0,
		Header:     0,
		LimitSplit: 0,
		MinLines:   2,
		TrimSpace:  true,
	}
	return g
}

// ReadAll reads all rows
// and returns a two-dimensional slice of rows and columns.
func (g *GuessWidth) ReadAll() [][]string {
	if len(g.preLines) == 0 {
		g.Scan(g.ScanNum)
	}

	g.Widths = make([]Cols, len(g.pos)+1)
	var rows [][]string
	for {
		columns, err := g.Read()
		if err != nil {
			break
		}
		g.UpdateMaxWidth(columns)
		rows = append(rows, columns)
	}

	rightThreshold := len(rows) / 2
	if rightThreshold == 0 {
		// Avoid deciding alignment from a single row.
		rightThreshold = len(rows)
	}
	g.SetJustified(rightThreshold)
	return rows
}

// UpdateMaxWidth updates the maximum width of the column.
func (g *GuessWidth) UpdateMaxWidth(columns []string) []Cols {
	if len(g.Widths) < len(columns) {
		for n := len(g.Widths); n < len(columns); n++ {
			g.Widths = append(g.Widths, Cols{})
		}
	}

	lastCol := len(g.Widths) - 1
	for n, col := range columns {
		width := uniseg.StringWidth(strings.TrimSpace(col))
		if width > g.Widths[n].Width {
			g.Widths[n].Width = width
		}
		if n != lastCol && isRightAlign(col) {
			g.Widths[n].rightCount++
		}
	}
	return g.Widths
}

// SetJustified marks a column as right-justified when its right-aligned row count
// exceeds threshold.
func (g *GuessWidth) SetJustified(threshold int) []Cols {
	for n, col := range g.Widths {
		if col.rightCount > threshold {
			col.Justified = Right
		}
		g.Widths[n] = col
	}
	return g.Widths
}

// isRightAlign checks if a string is right-aligned based on trailing spaces.
// A string is considered right-aligned if it has non-space characters followed by one or more spaces at the end.
func isRightAlign(str string) bool {
	if str == "" {
		return false
	}
	r := len(str) - 1
	for n := 0; n < len(str); n++ {
		if str[n] != ' ' {
			return false
		}
		if str[r-n] != ' ' {
			return true
		}
	}
	return false
}

// readLine reads a single line from the reader, handling line endings and EOF correctly.
func readLine(reader *bufio.Reader) (string, error) {
	line, err := reader.ReadString('\n')
	if err != nil && (err != io.EOF || len(line) == 0) {
		return "", err
	}
	if before, ok := strings.CutSuffix(line, "\n"); ok {
		line = before
		line = strings.TrimSuffix(line, "\r")
	}
	return line, nil
}

// Scan preReads and parses the lines.
func (g *GuessWidth) Scan(num int) {
	for range num {
		line, err := readLine(g.reader)
		if err != nil {
			break
		}
		g.preLines = append(g.preLines, line)
	}

	g.pos = Positions(g.preLines, g.Header, g.MinLines)

	if g.LimitSplit > 0 {
		if len(g.pos) > g.LimitSplit {
			g.pos = g.pos[:g.LimitSplit]
		}
	}
}

// Read reads one row and returns a slice of columns.
// Scan is executed first if it is not preRead.
func (g *GuessWidth) Read() ([]string, error) {
	if len(g.preLines) == 0 {
		g.Scan(g.ScanNum)
	}

	var line string
	if g.preCount < len(g.preLines) {
		line = g.preLines[g.preCount]
		g.preCount++
	} else {
		var err error
		line, err = readLine(g.reader)
		if err != nil {
			return nil, err
		}
	}

	return split(line, g.pos, g.TrimSpace), nil
}

// ToTable parses a slice of lines and returns a table.
func ToTable(lines []string, header int, trimSpace bool) [][]string {
	pos := Positions(lines, header, 2)
	return toRows(lines, pos, trimSpace)
}

// ToTableN parses a slice of lines and returns a table, but limits the number of splits.
func ToTableN(lines []string, header int, numSplit int, trimSpace bool) [][]string {
	if numSplit <= 0 {
		return ToTable(lines, header, trimSpace)
	}
	pos := Positions(lines, header, 2)
	if len(pos) > numSplit {
		pos = pos[:numSplit]
	}
	return toRows(lines, pos, trimSpace)
}

// Positions returns separator positions
// from multiple lines and header line number.
// Lines before the header line are ignored.
func Positions(lines []string, header int, minLines int) []int {
	var blanks []int

	if header < 0 {
		header = 0
	}
	for n, line := range lines {
		if n < header {
			continue
		}
		if n == header {
			blanks = lookupBlanks(strings.TrimSuffix(line, " "))
			continue
		}
		blanks = countBlanks(blanks, strings.TrimSuffix(line, " "))
	}

	pos := positions(blanks, minLines)
	if header < len(lines) {
		pos = ensureLastHeaderColumnPosition(pos, lines[header])
	}
	return pos
}

// ensureLastHeaderColumnPosition ensures that the last header column position
// is included in the separator positions if it exists.
func ensureLastHeaderColumnPosition(pos []int, headerLine string) []int {
	if p, ok := lastHeaderColumnPosition(headerLine); ok {
		found := false
		for _, existing := range pos {
			if existing >= p-1 && existing <= p {
				found = true
				break
			}
		}
		if !found {
			pos = append(pos, p)
			sort.Ints(pos)
		}
	}
	return pos
}

// lastHeaderColumnPosition returns the start position of the last column in
// the header line and a boolean indicating if it exists.
func lastHeaderColumnPosition(line string) (int, bool) {
	gr := uniseg.NewGraphemes(line)
	width := 0
	lastColumnStart := 0
	hasPreviousColumn := false
	previousWasSpace := true
	for gr.Next() {
		value := gr.Str()
		if value == " " {
			previousWasSpace = true
		} else {
			if previousWasSpace {
				if hasPreviousColumn {
					lastColumnStart = width
				}
				hasPreviousColumn = true
			}
			previousWasSpace = false
		}
		width += gr.Width()
	}
	return lastColumnStart, hasPreviousColumn && lastColumnStart > 0
}

// separatorPosition determines the best position for a separator in a line.
func separatorPosition(graphemes []string, widths []int, p int, start int, pos []int, n int, w int) int {
	r, _ := utf8.DecodeRuneInString(graphemes[p])
	if unicode.IsSpace(r) {
		return p
	}

	f := p
	fp := 0

	for ; f < len(graphemes); f++ {
		r, _ := utf8.DecodeRuneInString(graphemes[f])
		if unicode.IsSpace(r) {
			break
		}
		fp += widths[f]
	}

	b := p
	bp := 0
	for ; b > 0; b-- {
		r, _ := utf8.DecodeRuneInString(graphemes[b])
		if unicode.IsSpace(r) {
			break
		}
		bp += widths[b]
	}

	fWidth := w + fp
	bWidth := w - bp
	if n < len(pos)-1 {
		if separatorBeforeDoubleSpace(graphemes, widths, p, f, fWidth, pos[n+1]) {
			return f
		}
		if fWidth == pos[n+1] {
			return b
		}
		if bWidth == pos[n] {
			return f
		}
		if bWidth > pos[n] && bWidth < pos[n+1] {
			return b
		}
	}
	if fp > bp && b > start {
		return b
	}
	return f
}

// separatorBeforeDoubleSpace keeps a shifted value in its column when the
// following value starts after at least two spaces at the next column boundary.
func separatorBeforeDoubleSpace(graphemes []string, widths []int, p, f, fWidth, nextPosition int) bool {
	wordStart := p == 0
	if p > 0 {
		r, _ := utf8.DecodeRuneInString(graphemes[p-1])
		wordStart = unicode.IsSpace(r)
	}
	if !wordStart || f >= len(graphemes) {
		return false
	}

	next := f
	spaceWidth := 0
	for next < len(graphemes) {
		r, _ := utf8.DecodeRuneInString(graphemes[next])
		if !unicode.IsSpace(r) {
			break
		}
		spaceWidth += widths[next]
		next++
	}
	if spaceWidth < 2 || next >= len(graphemes) {
		return false
	}

	nextWidth := fWidth + spaceWidth
	for next < len(graphemes) {
		r, _ := utf8.DecodeRuneInString(graphemes[next])
		if unicode.IsSpace(r) {
			break
		}
		nextWidth += widths[next]
		next++
	}
	return nextWidth == nextPosition
}

// split splits a line into columns based on separator positions.
func split(line string, pos []int, trimSpace bool) []string {
	columns := make([]string, len(pos)+1)
	if len(pos) == 0 {
		if trimSpace {
			line = strings.TrimSpace(line)
		}
		columns[0] = line
		return columns
	}

	var graphemes []string
	var widths []int
	gr := uniseg.NewGraphemes(line)
	for gr.Next() {
		graphemes = append(graphemes, gr.Str())
		widths = append(widths, gr.Width())
	}

	n := 0
	start, end := 0, 0
	w := 0
	for p := 0; p < len(graphemes); p++ {
		if n >= len(pos) {
			start = end + 1
			break
		}
		if pos[n] <= w {
			end = separatorPosition(graphemes, widths, p, start, pos, n, w)
			if start > end {
				break
			}
			col := strings.Join(graphemes[start:end], "")
			if trimSpace {
				columns[n] = strings.TrimSpace(col)
			} else {
				columns[n] = col
			}
			n++
			start = end
		}
		w += widths[p]
	}
	if start >= len(graphemes) {
		return columns
	}

	// Handle the remaining part of the line after the last separator.
	col := strings.Join(graphemes[start:], "")
	if trimSpace {
		col = strings.TrimSpace(col)
	}
	columns[n] = col
	return columns
}

// toRows returns rows separated by columns.
func toRows(lines []string, pos []int, trimSpace bool) [][]string {
	rows := make([][]string, 0, len(lines))
	for _, line := range lines {
		columns := split(line, pos, trimSpace)
		rows = append(rows, columns)
	}
	return rows
}

// lookupBlanks returns a slice of counts of blank spaces in the line.
func lookupBlanks(line string) []int {
	blanks := make([]int, 0)
	first := true
	gr := uniseg.NewGraphemes(line)
	for gr.Next() {
		v := gr.Str()
		if v == " " {
			if first {
				blanks = append(blanks, 0)
				continue
			}
			blanks = append(blanks, 1)
			continue
		}

		first = false
		blanks = append(blanks, 0)
		if gr.Width() == 2 {
			blanks = append(blanks, 0)
		}
	}
	return blanks
}

// countBlanks counts the number of blank spaces in the line and updates the blanks slice.
func countBlanks(blanks []int, line string) []int {
	n := 0
	gr := uniseg.NewGraphemes(line)
	for gr.Next() {
		r := gr.Str()
		if n >= len(blanks) {
			break
		}
		if r == " " && blanks[n] > 0 {
			blanks[n] += 1
		}

		n++
		if gr.Width() == 2 {
			n++
		}
	}
	return blanks
}

// positions generates a list of separator positions from a blank slice.
func positions(blanks []int, minLines int) []int {
	max := minLines
	p := 0
	var pos []int
	for n, v := range blanks {
		if v >= max {
			max = v
			p = n
		}
		if v == 0 {
			max = minLines
			if p > 0 {
				pos = append(pos, p)
				p = 0
			}
		}
	}
	return pos
}

// DebugCountPrint is for debugging which prints the space count.
func DebugCountPrint(line string, blanks []int) {
	fmt.Println(line)
	for _, k := range blanks {
		fmt.Print(k)
	}
	fmt.Println()
}

var (
	version  string
	revision string
)

// Version returns the version of the package.
func Version() string {
	if version != "" {
		return version + " rev:" + revision
	}
	info, ok := debug.ReadBuildInfo()
	if !ok {
		return "(devel)"
	}
	return info.Main.Version
}