// dx — count files and directories, mwx'2016
//
// Go-Portierung des Perl-Originals (dx 1.0.2), Ausgabe eingefärbt nach der
// mwxcol-Palette (~/src/mwxcol). Ohne Terminal (Pipe, NO_COLOR, --color=never)
// bleibt die Ausgabe unverändert einfarbig.
package main

import (
	"bufio"
	"fmt"
	"io/fs"
	"os"
	"path/filepath"
	"runtime"
	"sort"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"unicode/utf8"

	"github.com/fatih/color"
)

const version = "2.1.0"

// ------------------------------------------------------------------ Palette

type rgb struct{ r, g, b int }

var (
	white  = rgb{0xf0, 0xf0, 0xea}
	yellow = rgb{0xf9, 0xe2, 0xaf}
	orange = rgb{0xfa, 0xb3, 0x87}
	red    = rgb{0xf3, 0x8b, 0xa8}
	green  = rgb{0x94, 0xe2, 0xd5}
	blue   = rgb{0xb4, 0xbe, 0xfe}
	pink   = rgb{0xf5, 0xb2, 0xf7}
	violet = rgb{0xcb, 0xa6, 0xf7}
	grey   = rgb{0xaa, 0xaa, 0xbb}
	dark   = rgb{0x77, 0x77, 0x88}
	darker = rgb{0x44, 0x44, 0x55}
)

func plain(c rgb) *color.Color { return color.RGB(c.r, c.g, c.b) }
func bold(c rgb) *color.Color  { return color.New(color.Bold).AddRGB(c.r, c.g, c.b) }

// Rollen wie im eza-Theme, damit dx und eza dieselbe Datei gleich einfärben.
var (
	cWhite   = plain(white)
	cYellow  = plain(yellow)
	cOrange  = plain(orange)
	cRed     = plain(red)
	cGreen   = plain(green)
	cBlue    = plain(blue)
	cPink    = plain(pink)
	cViolet  = plain(violet)
	cGrey    = plain(grey)
	cDark    = plain(dark)
	cDarker  = plain(darker)
	cDirB    = bold(blue)  // Verzeichnis
	cExecB   = bold(green) // ausführbar
	cWhiteB  = bold(white)
	cYellowB = bold(yellow)
	cOrangeB = bold(orange)
	cRedB    = bold(red)
	cGreenB  = bold(green)
	cDarkB   = bold(dark)
)

// ------------------------------------------------------------------- Zählen

type entry struct {
	name string // Anzeigename, Verzeichnisse mit "/"
	col  *color.Color
	n    int64
	size int64
	agg  bool // Sammelzeile aus -t, kein echter Eintrag
}

type result struct {
	entries  []entry
	totalN   int64
	totalS   int64
	unreadbl int64
}

func scan(dir string, onlyDirs bool) (result, error) {
	des, err := os.ReadDir(dir)
	if err != nil {
		return result{}, err
	}

	var res result
	var mu sync.Mutex
	var wg sync.WaitGroup
	sem := make(chan struct{}, runtime.NumCPU())

	for _, de := range des {
		name := de.Name()
		path := filepath.Join(dir, name)

		// Perls -d/-f folgen Symlinks; ein toter Link fällt durch beide Tests.
		st, err := os.Stat(path)
		if err != nil {
			atomic.AddInt64(&res.unreadbl, 1)
			continue
		}
		isLink := de.Type()&os.ModeSymlink != 0

		switch {
		case st.IsDir():
			wg.Go(func() {
				sem <- struct{}{}
				defer func() { <-sem }()

				// WalkDir folgt Symlinks nicht — wie `find <pfad> -type f`.
				n, size := walk(path, &res.unreadbl)

				col := cDirB
				if isLink {
					col = cPink
				}
				mu.Lock()
				res.entries = append(res.entries, entry{name: name + "/", col: col, n: n, size: size})
				res.totalN += n
				res.totalS += size
				mu.Unlock()
			})

		case st.Mode().IsRegular() && !onlyDirs:
			col := fileColor(name, st.Mode())
			if isLink {
				col = cPink
			}
			mu.Lock()
			res.entries = append(res.entries, entry{name: name, col: col, n: 1, size: st.Size()})
			res.totalN++
			res.totalS += st.Size()
			mu.Unlock()
		}
	}
	wg.Wait()

	return res, nil
}

func walk(root string, unreadable *int64) (n, size int64) {
	filepath.WalkDir(root, func(_ string, d fs.DirEntry, err error) error {
		if err != nil {
			atomic.AddInt64(unreadable, 1)
			return nil
		}
		if !d.Type().IsRegular() { // wie find -type f
			return nil
		}
		info, err := d.Info()
		if err != nil {
			atomic.AddInt64(unreadable, 1)
			return nil
		}
		n++
		size += info.Size()
		return nil
	})
	return
}

// ------------------------------------------------------------------ Ausgabe

// Schwellen wie im Perl-Original; unterhalb von 10 der jeweiligen Einheit
// jedoch mit einer Nachkommastelle — "1.9G" statt "2G".
func formSize(b int64) string {
	const k = 1024
	f := float64(b)
	switch {
	case b > k*k*k*k:
		return scaled(f/k/k/k/k, "T")
	case b > k*k*k:
		return scaled(f/k/k/k, "G")
	case b > k*k:
		return scaled(f/k/k, "M")
	case b > k:
		return scaled(f/k, "K")
	default:
		return fmt.Sprintf("%.0f", f) // ganze Bytes, nie gebrochen
	}
}

func scaled(v float64, unit string) string {
	if v < 9.95 {
		return fmt.Sprintf("%.1f%s", v, unit)
	}
	return fmt.Sprintf("%.0f%s", v, unit)
}

// Größenrampe wie im eza-Theme: byte/kilo grün, mega gelb, giga orange, huge rot.
func sizeColor(b int64, strong bool) *color.Color {
	switch {
	case b == 0:
		if strong {
			return cDarkB
		}
		return cDark
	case b > 1024*1024*1024*1024:
		if strong {
			return cRedB
		}
		return cRed
	case b > 1024*1024*1024:
		if strong {
			return cOrangeB
		}
		return cOrange
	case b > 1024*1024:
		if strong {
			return cYellowB
		}
		return cYellow
	default:
		if strong {
			return cGreenB
		}
		return cGreen
	}
}

func countColor(n int64) *color.Color {
	if n == 0 {
		return cDark
	}
	return cGrey
}

func pad(s string, w int) string {
	if n := w - utf8.RuneCountInString(s); n > 0 {
		return strings.Repeat(" ", n)
	}
	return ""
}

// Anteilsbalken in Achtelschritten, feste Breite barCells.
const barCells = 10

var eighths = []rune{' ', '▏', '▎', '▍', '▌', '▋', '▊', '▉'}

func bar(frac float64) string {
	if frac > 1 {
		frac = 1
	}
	if !(frac > 0) { // fängt auch NaN
		return strings.Repeat(" ", barCells)
	}
	units := int(frac*barCells*8 + 0.5)
	full, rem := units/8, units%8

	var b strings.Builder
	b.WriteString(strings.Repeat("█", full))
	used := full
	if rem > 0 {
		b.WriteRune(eighths[rem])
		used++
	}
	b.WriteString(strings.Repeat(" ", barCells-used))
	return b.String()
}

func percent(frac float64) string {
	p := frac * 100
	if p > 0 && p < 0.5 {
		return " <1%"
	}
	return fmt.Sprintf("%3.0f%%", p)
}

func printDir(w *bufio.Writer, dir string, res result, o opts) {
	metric := func(e entry) int64 {
		if o.byCount {
			return e.n
		}
		return e.size
	}
	sort.SliceStable(res.entries, func(i, j int) bool {
		a, b := res.entries[i], res.entries[j]
		if ka, kb := metric(a), metric(b); ka != kb {
			return ka < kb
		}
		return a.name < b.name // Perl ließ Gleichstände der Hash-Ordnung
	})

	// -t N: nur die N größten zeigen, der Rest wird zu einer Zeile.
	rows := res.entries
	var rest *entry
	if o.top > 0 && len(rows) > o.top {
		cut := len(rows) - o.top
		var agg entry
		for _, e := range rows[:cut] {
			agg.n += e.n
			agg.size += e.size
		}
		agg.name = fmt.Sprintf("… %d more", cut)
		agg.col, agg.agg = cDark, true
		rest, rows = &agg, rows[cut:]
	}
	if o.reverse {
		for i, j := 0, len(rows)-1; i < j; i, j = i+1, j-1 {
			rows[i], rows[j] = rows[j], rows[i]
		}
	}
	if rest != nil { // Restzeile bleibt am kleinen Ende
		if o.reverse {
			rows = append(rows, *rest)
		} else {
			rows = append([]entry{*rest}, rows...)
		}
	}

	total := res.totalS
	if o.byCount {
		total = res.totalN
	}

	ml := utf8.RuneCountInString(dir)
	for _, e := range rows {
		if n := utf8.RuneCountInString(e.name); n > ml {
			ml = n
		}
	}
	width := ml + 1

	fmt.Fprintf(w, "%s%s   %s %s\n",
		cDirB.Sprint(dir), pad(dir, width),
		cWhiteB.Sprintf("%8d", res.totalN),
		sizeColor(res.totalS, true).Sprintf("%8s", formSize(res.totalS)))

	for _, e := range rows {
		fmt.Fprintf(w, "  %s%s %s %s",
			e.col.Sprint(e.name), pad(e.name, width),
			countColor(e.n).Sprintf("%8d", e.n),
			sizeColor(e.size, false).Sprintf("%8s", formSize(e.size)))

		if o.bar {
			var frac float64
			if total > 0 {
				frac = float64(metric(e)) / float64(total)
			}
			// Balken trägt die Farbe der Spalte, deren Anteil er zeigt.
			col := sizeColor(e.size, false)
			if o.byCount {
				col = countColor(e.n)
			}
			if e.agg { // Sammelzeile bleibt zurückhaltend
				col = cDark
			}
			fmt.Fprintf(w, "  %s %s",
				col.Sprint(bar(frac)), cGrey.Sprint(percent(frac)))
		}
		fmt.Fprintln(w)
	}
}

// --------------------------------------------------------------------- CLI

type opts struct {
	byCount  bool // -n
	onlyDirs bool // -d
	reverse  bool // -r
	top      int  // -t N, 0 = alle
	bar      bool
}

func usage(w *os.File) {
	fmt.Fprintf(w, `dx %s, count files and directories, mwx'2016
usage: dx [-n] [-d] [-r] [-t N] [path]
  -n    sort by number of files (default by size)
  -d    count only directories
  -r    reverse order (largest first)
  -t N  show only the N largest, sum up the rest
  -h    this help
  --bar / --no-bar            share bar and percentage (default: when colorized)
  --color=auto|always|never   colorize output (default auto)
  --version
`, version)
}

func die(format string, a ...any) {
	fmt.Fprintf(os.Stderr, "dx: "+format+"\n", a...)
	os.Exit(2)
}

func main() {
	var o opts
	var barSet bool // sonst folgt der Balken der Farbentscheidung
	var dir string
	when := "auto"

	args := os.Args[1:]
	for i := 0; i < len(args); i++ {
		a := args[i]
		switch {
		case a == "--":
			if i+1 < len(args) {
				dir = args[i+1]
			}
			i = len(args)
		case a == "--help":
			usage(os.Stdout)
			return
		case a == "--version":
			fmt.Printf("dx %s\n", version)
			return
		case a == "--number":
			o.byCount = true
		case a == "--dirs":
			o.onlyDirs = true
		case a == "--reverse":
			o.reverse = true
		case a == "--bar":
			o.bar, barSet = true, true
		case a == "--no-bar":
			o.bar, barSet = false, true
		case a == "--no-color":
			when = "never"
		case strings.HasPrefix(a, "--top="):
			o.top = mustNum(strings.TrimPrefix(a, "--top="))
		case strings.HasPrefix(a, "--color="):
			when = strings.TrimPrefix(a, "--color=")
		case len(a) > 1 && a[0] == '-':
			flags := a[1:]
			for j := 0; j < len(flags); j++ { // Bündelung wie Getopt::Std: -dn
				switch flags[j] {
				case 'n':
					o.byCount = true
				case 'd':
					o.onlyDirs = true
				case 'r':
					o.reverse = true
				case 'h':
					usage(os.Stdout)
					return
				case 't': // Wert klebt am Flag (-t5) oder folgt (-t 5)
					val := flags[j+1:]
					if val == "" {
						if i+1 >= len(args) {
							die("option requires an argument -- t")
						}
						i++
						val = args[i]
					}
					o.top = mustNum(val)
					j = len(flags)
				default:
					fmt.Fprintf(os.Stderr, "dx: unknown option -- %c\n", flags[j])
					usage(os.Stderr)
					os.Exit(2)
				}
			}
		default:
			dir = a
		}
	}

	switch when {
	case "always":
		color.NoColor = false
	case "never":
		color.NoColor = true
	case "auto": // fatih/color entscheidet per isatty und NO_COLOR
	default:
		die("--color: expected auto|always|never, got %q", when)
	}
	if !barSet {
		o.bar = !color.NoColor
	}

	if dir == "" {
		dir = "."
	}

	res, err := scan(dir, o.onlyDirs)
	if err != nil {
		fmt.Fprintf(os.Stderr, "dx: %v\n", err)
		os.Exit(1)
	}

	w := bufio.NewWriter(os.Stdout)
	printDir(w, dir, res, o)
	w.Flush()

	if res.unreadbl > 0 {
		fmt.Fprintln(os.Stderr, cDark.Sprintf("dx: %d unreadable entries skipped", res.unreadbl))
	}
}

func mustNum(s string) int {
	n, err := strconv.Atoi(s)
	if err != nil || n < 1 {
		die("-t: expected a positive number, got %q", s)
	}
	return n
}