/*
The MIT License (MIT)

Copyright © 2024 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.
*/

/*
Single-file source-code for nt: this version has no http(s) support. Even
the unit-tests from the original nt are omitted.

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

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

package main

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

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


Nice Tables realigns and styles TSV (tab-separated values) data using ANSI
sequences. When not given filepaths to read data from, this tool reads from
standard input.

If you're using Linux or MacOS, you may find this cmd-line shortcut useful:

# View Nice Table(s) / Very Nice Table(s); uses my tools "nt" and "nn"
vnt() {
    awk 1 "$@" | nl -b a -w 1 -v 0 | nt | nn |
        awk '(NR - 1) % 5 == 1 && NR > 1 { print "" } 1' | less -JMKiCRS
}
`

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

	if err := run(os.Args[1:]); err != nil {
		os.Stderr.WriteString("\x1b[31m")
		os.Stderr.WriteString(err.Error())
		os.Stderr.WriteString("\x1b[0m\n")
		os.Exit(1)
	}
}

func run(paths []string) error {
	bw := bufio.NewWriter(os.Stdout)
	defer bw.Flush()

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

	if len(paths) == 0 {
		return niceTable(bw, os.Stdin)
	}
	return nil
}

func handleFile(w *bufio.Writer, 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 niceTable(w, f)
}

// loopItems lets you loop over a line's items, allocation-free style;
// when given empty strings, the callback func is never called
func loopItems(s string, sep byte, f func(i int, s string)) {
	if len(s) == 0 {
		return
	}

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

		f(i, s)
		return
	}
}

// tryNumeric tries to parse a strings as a valid/useable float64, which
// excludes NaNs and the infinities, returning a boolean instead of an
// error
func tryNumeric(s string) (f float64, ok bool) {
	f, err := strconv.ParseFloat(s, 64)
	if err == nil && !math.IsNaN(f) && !math.IsInf(f, 0) {
		return f, true
	}
	return f, false
}

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

func numericStyle(f float64) string {
	if f > 0 {
		if float64(int64(f)) == f {
			return "\x1b[38;5;29m"
		}
		return "\x1b[38;5;22m"
	}
	if f < 0 {
		if float64(int64(f)) == f {
			return "\x1b[38;5;1m"
		}
		return "\x1b[38;5;167m"
	}
	if f == 0 {
		// return "\x1b[38;5;33m"
		return "\x1b[38;5;26m"
	}
	return "\x1b[38;5;244m"
}

const (
	columnGap       = `  `
	tilesBackground = "\x1b[48;5;255m"
)

// writeSpaces does what it says, minimizing calls to write funcs
func writeSpaces(w *bufio.Writer, n int) {
	const spaces = `                                `
	for n >= len(spaces) {
		w.WriteString(spaces)
		n -= len(spaces)
	}

	if n > 0 {
		w.WriteString(spaces[:n])
	}
}

func niceTable(w *bufio.Writer, r io.Reader) error {
	const gb = 1024 * 1024 * 1024
	sc := bufio.NewScanner(r)
	sc.Buffer(nil, 8*gb)

	var res table
	for sc.Scan() {
		res.update(sc.Text())
	}
	if err := sc.Err(); err != nil {
		return err
	}

	if len(res.Rows) == 0 {
		return nil
	}

	totalsRow := makeTotalsRow(res)
	loopItems(totalsRow, '\t', func(i int, s string) {
		// keep track of widest rune-counts for each column
		w := utf8.RuneCountInString(s)
		if res.MaxWidth[i] < w {
			res.MaxWidth[i] = w
		}
	})

	for _, row := range res.Rows {
		writeRowTiles(w, row, res)
		writeRowItems(w, row, res)
		if err := w.WriteByte('\n'); err != nil {
			return nil
		}
	}

	writeSpaces(w, len(res.MaxWidth))
	writeRowItems(w, totalsRow, res)
	w.WriteByte('\n')
	return nil
}

func makeTotalsRow(res table) string {
	var sb strings.Builder
	for i, t := range res.Sums {
		if i > 0 {
			sb.WriteByte('\t')
		}
		if res.Numeric[i] > 0 {
			var buf [32]byte
			decs := res.MaxDecimals[i]
			s := strconv.AppendFloat(buf[:0], t, 'f', decs, 64)
			sb.Write(s)
		} else {
			sb.WriteString(`-`)
		}
	}
	return sb.String()
}

// table has all summary info gathered from TSV data, along with the row
// themselves, stored as lines/strings
type table struct {
	Rows []string

	MaxWidth []int

	Numeric []int

	MaxDecimals []int

	Sums []float64
}

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

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

	loopItems(line, '\t', func(i int, s string) {
		// ensure column-info-slices have enough room
		if i >= len(t.MaxWidth) {
			t.MaxWidth = append(t.MaxWidth, 0)
			t.Numeric = append(t.Numeric, 0)
			t.MaxDecimals = append(t.MaxDecimals, 0)
			t.Sums = append(t.Sums, 0.0)
		}

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

		// update stats for numeric items
		if f, ok := tryNumeric(s); ok {
			decs := countDecimals(s)
			if t.MaxDecimals[i] < decs {
				t.MaxDecimals[i] = decs
			}
			t.Numeric[i]++
			t.Sums[i] += f
		}
	})
}

func writeRowTiles(w *bufio.Writer, row string, t table) {
	w.WriteString(tilesBackground)

	end := 0
	loopItems(row, '\t', func(i int, s string) {
		writeTile(w, s)
		end = i
	})

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

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

	if f, ok := tryNumeric(s); ok {
		w.WriteString(numericStyle(f))
		w.WriteString("■")
	} else {
		w.WriteString("\x1b[38;5;244m■")
	}
}

func writeRowItems(w *bufio.Writer, row string, t table) {
	loopItems(row, '\t', func(i int, s string) {
		w.WriteString(columnGap)

		if f, ok := tryNumeric(s); ok {
			trail := 0
			decs := countDecimals(s)
			if decs > 0 {
				trail = t.MaxDecimals[i] - decs
			} else if t.MaxDecimals[i] > 0 {
				trail = t.MaxDecimals[i] + 1
			}

			n := utf8.RuneCountInString(s)
			lead := t.MaxWidth[i] - n - trail
			writeSpaces(w, lead)
			writeNumericItem(w, s, numericStyle(f))
			if i < len(t.MaxWidth)-1 {
				writeSpaces(w, trail)
			}
			return
		}

		w.WriteString(s)
		if i < len(t.MaxWidth)-1 {
			n := utf8.RuneCountInString(s)
			writeSpaces(w, t.MaxWidth[i]-n)
		}
	})
}

// func writeNumericItem(w *bufio.Writer, s string, startStyle string) {
// 	w.WriteString(startStyle)
// 	w.WriteString(s)
// 	w.WriteString("\x1b[0m")
// }

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