/*
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 ncol.go
*/

package main

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

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

Nice COLumns realigns and styles data tables using ANSI color sequences. In
particular, all auto-detected numbers are styled so they're easier to read
at a glance. Input tables can be either lines of space-separated values or
tab-separated values, and are auto-detected using the first non-empty line.

When not given filepaths to read data from, this tool reads from standard
input by default.

The only option is the help option, using any of "-h", "--h", "-help", or
"--help".
`

// altDigitStyle is used to make 4+ digit-runs easier to read
const altDigitStyle = "\x1b[38;2;168;168;168m"

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, sb *strings.Builder) bool {
	if len(s) == 0 {
		return false
	}

	digits := 0

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

		if b := s[0]; '0' <= b && b <= '9' {
			sb.WriteByte(b)
			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, sb *strings.Builder) bool {
	if len(s) == 0 {
		return false
	}

	sb.Reset()

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

	s = skipLeadingANSI(s)
	if len(s) == 0 {
		return false
	}
	if b := s[0]; b == '.' {
		sb.WriteByte(b)
		return isDigits(s[1:], sb)
	}

	digits := 0

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

		b := s[0]
		sb.WriteByte(b)

		if b == '.' {
			return isDigits(s[1:], sb)
		}

		if !('0' <= b && b <= '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

	Numeric []int

	Sums []float64

	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++
		})
	}

	var sb strings.Builder

	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)
			t.Numeric = append(t.Numeric, 0)
			t.Sums = append(t.Sums, 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, &sb) {
			dd := countDotDecimals(s)
			if t.MaxDotDecimals[i] < dd {
				t.MaxDotDecimals[i] = dd
			}

			t.Numeric[i]++
			f, _ := strconv.ParseFloat(sb.String(), 64)
			t.Sums[i] += f
		}
	})
}

func realign(w *bufio.Writer, t table) {
	var sb strings.Builder
	sums := make([]string, 0, t.Columns)

	for i := 0; i < t.Columns; i++ {
		s := `-`
		width := 1

		if t.Numeric[i] > 0 {
			decs := t.MaxDotDecimals[i]
			if decs > 0 {
				decs--
			}

			var buf [64]byte
			s = string(strconv.AppendFloat(buf[:0], t.Sums[i], 'f', decs, 64))
			width = len(s)
		}

		if t.MaxWidth[i] < width {
			t.MaxWidth[i] = width
		}
		sums = append(sums, s)
	}

	for _, line := range t.Rows {
		due := 0

		writeRowTiles(w, line, t, &sb)

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

			if isNumeric(s, &sb) {
				dd := countDotDecimals(s)
				rpad := t.MaxDotDecimals[i] - dd
				width := countWidth(s)
				lpad := t.MaxWidth[i] - (width + rpad) + due
				writeSpaces(w, lpad)
				f, _ := strconv.ParseFloat(sb.String(), 64)
				writeNumericItem(w, s, numericStyle(f))
				due = rpad
				return
			}

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

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

	if t.Columns == 0 {
		return
	}

	due := 0
	writeSpaces(w, t.Columns)

	for i, s := range sums {
		due += 2

		if t.Numeric[i] == 0 {
			writeSpaces(w, due)
			w.WriteString(s)
			due = t.MaxWidth[i] - countWidth(s)
			continue
		}

		lpad := t.MaxWidth[i] - len(s) + due
		writeSpaces(w, lpad)
		writeNumericItem(w, s, numericStyle(t.Sums[i]))
		due = 0
	}

	w.WriteByte('\n')
}

// 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])
}

func writeRowTiles(w *bufio.Writer, row string, t table, sb *strings.Builder) {
	end := 0
	t.LoopItems(row, t.Columns, func(i int, s string) {
		writeTile(w, s, sb)
		end = i
	})

	if end < len(t.MaxWidth)-1 {
		w.WriteString("\x1b[0m")
	}
	for i := end + 1; i < len(t.MaxWidth); i++ {
		w.WriteString("×")
	}
	w.WriteString("\x1b[0m")
}

func writeTile(w *bufio.Writer, s string, sb *strings.Builder) {
	if len(s) == 0 {
		w.WriteString("\x1b[0m○")
		return
	}

	if isNumeric(s, sb) {
		f, _ := strconv.ParseFloat(sb.String(), 64)
		w.WriteString(numericStyle(f))
		w.WriteString("■")
		return
	}

	if s[0] == ' ' || s[len(s)-1] == ' ' {
		w.WriteString("\x1b[38;2;196;160;0m■")
		return
	}

	w.WriteString("\x1b[38;2;128;128;128m■")
}

func numericStyle(f float64) string {
	if f > 0 {
		if float64(int64(f)) == f {
			return "\x1b[38;2;0;135;0m"
		}
		return "\x1b[38;2;0;155;95m"
	}
	if f < 0 {
		if float64(int64(f)) == f {
			return "\x1b[38;2;204;0;0m"
		}
		return "\x1b[38;2;215;95;95m"
	}
	if f == 0 {
		return "\x1b[38;2;0;95;215m"
	}
	return "\x1b[38;2;128;128;128m"
}

func writeNumericItem(w *bufio.Writer, s string, startStyle string) {
	w.WriteString(startStyle)
	if len(s) > 0 && (s[0] == '-' || s[0] == '+') {
		w.WriteByte(s[0])
		s = s[1:]
	}

	dot := strings.IndexByte(s, '.')
	if dot < 0 {
		restyleDigits(w, s, altDigitStyle)
		w.WriteString("\x1b[0m")
		return
	}

	if len(s[:dot]) > 3 {
		restyleDigits(w, s[:dot], altDigitStyle)
		w.WriteString("\x1b[0m")
		w.WriteString(startStyle)
		w.WriteByte('.')
	} else {
		w.WriteString(s[:dot])
		w.WriteByte('.')
	}

	rest := s[dot+1:]
	restyleDigits(w, rest, altDigitStyle)
	if len(rest) < 4 {
		w.WriteString("\x1b[0m")
	}
}

// restyleDigits renders a run of digits as alternating styled/unstyled runs
// of 3 digits, which greatly improves readability, and is the only purpose
// of this app; string is assumed to be all decimal digits
func restyleDigits(w *bufio.Writer, digits string, altStyle string) {
	if len(digits) < 4 {
		// digit sequence is short, so emit it as is
		w.WriteString(digits)
		return
	}

	// separate leading 0..2 digits which don't align with the 3-digit groups
	i := len(digits) % 3
	// emit leading digits unstyled, if there are any
	w.WriteString(digits[:i])
	// the rest is guaranteed to have a length which is a multiple of 3
	digits = digits[i:]

	// start by styling, unless there were no leading digits
	style := i != 0

	for len(digits) > 0 {
		if style {
			w.WriteString(altStyle)
			w.WriteString(digits[:3])
			w.WriteString("\x1b[0m")
		} else {
			w.WriteString(digits[:3])
		}

		// advance to the next triple: the start of this func is supposed
		// to guarantee this step always works
		digits = digits[3:]

		// alternate between styled and unstyled 3-digit groups
		style = !style
	}
}