/*
The MIT License (MIT)

Copyright © 2025 pacman64

Permission is hereby granted, free of charge, to any person obtaining a copy of
this software and associated documentation files (the “Software”), to deal
in the Software without restriction, including without limitation the rights to
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
of the Software, and to permit persons to whom the Software is furnished to do
so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/

/*
To compile a smaller-sized command-line app, you can use the `go` command as
follows:

go build -ldflags "-s -w" -trimpath realign.go
*/

package main

import (
	"bufio"
	"errors"
	"io"
	"os"
	"strings"
	"unicode/utf8"
)

const info = `
realign [options...] [filenames...]

Realign all detected columns, right-aligning any detected numbers in any
column. ANSI style-codes are also kept as given.

The only option available is to show this help message, using any of
"-h", "--h", "-help", or "--help", without the quotes.
`

func main() {
	if len(os.Args) > 1 {
		switch os.Args[1] {
		case `-h`, `--h`, `-help`, `--help`:
			os.Stderr.WriteString(info[1:])
			return
		}
	}

	args := os.Args[1:]
	if len(args) > 0 && args[0] == `--` {
		args = args[1:]
	}

	if err := run(args); err != nil {
		os.Stderr.WriteString(err.Error())
		os.Stderr.WriteString("\n")
		os.Exit(1)
	}
}

func run(paths []string) error {
	var res table

	for _, p := range paths {
		if err := handleFile(&res, p); err != nil {
			return err
		}
	}

	if len(paths) == 0 {
		if err := handleReader(&res, os.Stdin); err != nil {
			return err
		}
	}

	bw := bufio.NewWriter(os.Stdout)
	defer bw.Flush()
	realign(bw, res)
	return nil
}

func handleFile(res *table, path string) error {
	f, err := os.Open(path)
	if err != nil {
		// on windows, file-not-found error messages may mention `CreateFile`,
		// even when trying to open files in read-only mode
		return errors.New(`can't open file named ` + path)
	}
	defer f.Close()
	return handleReader(res, f)
}

func handleReader(res *table, r io.Reader) error {
	const gb = 1024 * 1024 * 1024
	sc := bufio.NewScanner(r)
	sc.Buffer(nil, 8*gb)

	for sc.Scan() {
		res.update(sc.Text())
	}
	return sc.Err()
}

// loopItemsSSV loops over a line's items, allocation-free style; when given
// empty strings, the callback func is never called
func loopItemsSSV(s string, max int, f func(i int, s string)) {
	for len(s) > 0 && s[len(s)-1] == ' ' {
		s = s[:len(s)-1]
	}

	for i := 0; true; i++ {
		for len(s) > 0 && s[0] == ' ' {
			s = s[1:]
		}

		if len(s) == 0 {
			return
		}

		if i+1 == max {
			f(i, s)
			return
		}

		if j := strings.IndexByte(s, ' '); j >= 0 {
			f(i, s[:j])
			s = s[j+1:]
			continue
		}

		f(i, s)
		return
	}
}

// loopItemsTSV loops over a line's tab-separated items, allocation-free style;
// when given empty strings, the callback func is never called
func loopItemsTSV(s string, max int, f func(i int, s string)) {
	if len(s) == 0 {
		return
	}

	for i := 0; true; i++ {
		if i+1 == max {
			f(i, s)
			return
		}

		if j := strings.IndexByte(s, '\t'); j >= 0 {
			f(i, s[:j])
			s = s[j+1:]
			continue
		}

		f(i, s)
		return
	}
}

func skipSingleLeadingANSI(s string) string {
	if len(s) < 3 || s[0] != '\x1b' || s[1] != '[' {
		return s
	}

	s = s[2:]

	for len(s) > 0 {
		b := s[0]
		s = s[1:]
		if ('A' <= b && b <= 'Z') || ('a' <= b && b <= 'z') {
			break
		}
	}

	return s
}

func skipLeadingANSI(s string) string {
	prev := len(s)

	for len(s) > 0 {
		s = skipSingleLeadingANSI(s)
		if len(s) == prev {
			return s
		}
		prev = len(s)
	}

	return s
}

// isNumeric checks if a string is a valid/useable float64, which excludes
// NaNs and the infinities
// func isNumeric(s string) bool {
// 	f, err := strconv.ParseFloat(s, 64)
// 	return err == nil && !math.IsNaN(f) && !math.IsInf(f, 0)
// }

func isDigits(s string) bool {
	if len(s) == 0 {
		return false
	}

	digits := 0

	for {
		s = skipLeadingANSI(s)
		if len(s) == 0 {
			break
		}

		if '0' <= s[0] && s[0] <= '9' {
			s = s[1:]
			digits++
		} else {
			return false
		}
	}

	s = skipLeadingANSI(s)
	return len(s) == 0 && digits > 0
}

// isNumeric checks if a string is valid/useable as a number
func isNumeric(s string) bool {
	if len(s) == 0 {
		return false
	}

	s = skipLeadingANSI(s)
	if len(s) > 0 && (s[0] == '+' || s[0] == '-') {
		s = s[1:]
	}

	s = skipLeadingANSI(s)
	if len(s) == 0 {
		return false
	}
	if s[0] == '.' {
		return isDigits(s[1:])
	}

	digits := 0

	for {
		s = skipLeadingANSI(s)
		if len(s) == 0 {
			break
		}

		if s[0] == '.' {
			return isDigits(s[1:])
		}

		if !('0' <= s[0] && s[0] <= '9') {
			return false
		}

		digits++
		s = s[1:]
	}

	s = skipLeadingANSI(s)
	return len(s) == 0 && digits > 0
}

// // countDecimals counts decimal digits from the string given, assuming it
// // represents a valid/useable float64, when parsed
// func countDecimals(s string) int {
// 	if dot := strings.IndexByte(s, '.'); dot >= 0 {
// 		return len(s) - dot - 1
// 	}
// 	return 0
// }

// countDecimals counts decimal digits from the string given, assuming it
// represents a valid/useable float64, when parsed
func countDecimals(s string) int {
	dot := strings.IndexByte(s, '.')
	if dot < 0 {
		return 0
	}

	decs := 0
	s = s[dot+1:]

	for len(s) > 0 {
		s = skipLeadingANSI(s)
		if len(s) == 0 {
			break
		}
		if '0' <= s[0] && s[0] <= '9' {
			decs++
		}
		s = s[1:]
	}

	return decs
}

// countDotDecimals is like func countDecimals, but this one also includes
// the dot, when any decimals are present, else the count stays at 0
func countDotDecimals(s string) int {
	decs := countDecimals(s)
	if decs > 0 {
		return decs + 1
	}
	return decs
}

// func countWidth(s string) int{
// 	return utf8.RuneCountInString(s)
// }

func countWidth(s string) int {
	c := 0
	for len(s) > 0 {
		i := strings.Index(s, "\x1b[")
		if i < 0 {
			c += utf8.RuneCountInString(s)
			return c
		}

		c += utf8.RuneCountInString(s[:i])
		s = s[i+2:]

		for len(s) > 0 {
			b := s[0]
			s = s[1:]
			if ('A' <= b && b <= 'Z') || ('a' <= b && b <= 'z') {
				break
			}
		}
	}

	return c
}

// table has all summary info gathered from the data, along with the row
// themselves, stored as lines/strings
type table struct {
	Columns int

	Rows []string

	MaxWidth []int

	MaxDotDecimals []int

	LoopItems func(line string, items int, f func(i int, s string))
}

func (t *table) update(line string) {
	if len(line) == 0 {
		return
	}

	t.Rows = append(t.Rows, line)

	if t.LoopItems == nil {
		if strings.ContainsRune(line, '\t') {
			t.LoopItems = loopItemsTSV
		} else {
			t.LoopItems = loopItemsSSV
		}
		t.LoopItems(line, t.Columns, func(i int, s string) {
			t.Columns++
		})
	}

	t.LoopItems(line, t.Columns, func(i int, s string) {
		// ensure column-info-slices have enough room
		if i >= len(t.MaxWidth) {
			t.MaxWidth = append(t.MaxWidth, 0)
			t.MaxDotDecimals = append(t.MaxDotDecimals, 0)
		}

		// keep track of widest rune-counts for each column
		w := countWidth(s)
		if t.MaxWidth[i] < w {
			t.MaxWidth[i] = w
		}

		// update stats for numeric items
		if isNumeric(s) {
			dd := countDotDecimals(s)
			if t.MaxDotDecimals[i] < dd {
				t.MaxDotDecimals[i] = dd
			}
		}
	})
}

func realign(w *bufio.Writer, t table) {
	for _, line := range t.Rows {
		due := 0

		t.LoopItems(line, t.Columns, func(i int, s string) {
			if i > 0 {
				due += 2
			}

			if isNumeric(s) {
				dd := countDotDecimals(s)
				rpad := t.MaxDotDecimals[i] - dd
				width := countWidth(s)
				lpad := t.MaxWidth[i] - (width + rpad) + due
				writeSpaces(w, lpad)
				w.WriteString(s)
				due = rpad
				return
			}

			writeSpaces(w, due)
			w.WriteString(s)
			due = t.MaxWidth[i] - countWidth(s)
		})

		if w.WriteByte('\n') != nil {
			return
		}
	}
}

// writeSpaces does what it says, minimizing calls to write-like funcs
func writeSpaces(w *bufio.Writer, n int) {
	const spaces = `                                `
	if n < 1 {
		return
	}

	for n >= len(spaces) {
		w.WriteString(spaces)
		n -= len(spaces)
	}
	w.WriteString(spaces[:n])
}