package main

import (
	"bufio"
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"time"
	"unicode/utf8"
)

var debugMode bool
var debugErrors []string
var logFile *os.File

func initLog() {
	if !debugMode {
		return
	}
	f, err := os.Create("diskmap_debug.log")
	if err != nil {
		fmt.Printf("  %s[debug] could not create log file: %v%s\n", red, err, reset)
		return
	}
	logFile = f
	logWrite("diskmap debug log — %s", time.Now().Format("2006-01-02 15:04:05"))
	logWrite("OS: %s  ARCH: %s", runtime.GOOS, runtime.GOARCH)
	logWrite(strings.Repeat("-", 60))
}

func logWrite(format string, args ...any) {
	msg := fmt.Sprintf(format, args...)
	if logFile != nil {
		fmt.Fprintln(logFile, msg)
	}
}

func dbg(format string, args ...any) {
	if !debugMode {
		return
	}
	msg := fmt.Sprintf(format, args...)
	fmt.Printf("  %s[dbg]%s %s\n", dim, reset, msg)
	logWrite("[dbg] " + msg)
}

func dbgErr(format string, args ...any) {
	if !debugMode {
		return
	}
	msg := fmt.Sprintf(format, args...)
	fmt.Printf("  %s[err]%s %s\n", red, reset, msg)
	logWrite("[err] " + msg)
	debugErrors = append(debugErrors, msg)
}

// ── ANSI colours ─────────────────────────────────────────────────────────────

const (
	reset  = "\033[0m"
	bold   = "\033[1m"
	dim    = "\033[2m"
	red    = "\033[31m"
	green  = "\033[32m"
	yellow = "\033[33m"
	cyan   = "\033[36m"
	white  = "\033[97m"
	orange = "\033[38;5;208m"
	rust   = "\033[38;5;166m"
)

type DirEntry struct {
	Name  string
	Path  string
	Size  int64
	IsDir bool
}

// ── Size helpers ──────────────────────────────────────────────────────────────

func humanSize(b int64) string {
	const (
		KB = 1 << 10
		MB = 1 << 20
		GB = 1 << 30
		TB = 1 << 40
	)
	switch {
	case b >= TB:
		return fmt.Sprintf("%6.1f TB", float64(b)/TB)
	case b >= GB:
		return fmt.Sprintf("%6.1f GB", float64(b)/GB)
	case b >= MB:
		return fmt.Sprintf("%6.1f MB", float64(b)/MB)
	case b >= KB:
		return fmt.Sprintf("%6.1f KB", float64(b)/KB)
	default:
		return fmt.Sprintf("%6d  B", b)
	}
}

func sizeColour(b int64) string {
	const GB = 1 << 30
	const MB = 1 << 20
	switch {
	case b >= 10*GB:
		return red
	case b >= 1*GB:
		return orange
	case b >= 100*MB:
		return yellow
	case b >= 10*MB:
		return cyan
	default:
		return dim + white
	}
}

// ── Scanning ──────────────────────────────────────────────────────────────────

func scanDir(path string) (total int64) {
	defer func() {
		if r := recover(); r != nil {
			dbgErr("PANIC recovered in scanDir(%s): %v", path, r)
		}
	}()
	dbg("entering %s", path)
	entries, err := os.ReadDir(path)
	if err != nil {
		dbgErr("ReadDir: %s => %v", path, err)
		return 0
	}
	for _, e := range entries {
		full := filepath.Join(path, e.Name())
		if e.IsDir() {
			total += scanDir(full)
		} else {
			info, err := e.Info()
			if err != nil {
				dbgErr("Info: %s => %v", full, err)
			} else {
				total += info.Size()
			}
		}
	}
	dbg("done %s => %s", path, humanSize(total))
	return total
}

func scanTopLevel(root string) ([]DirEntry, error) {
	dbg("scanTopLevel root=%s", root)
	entries, err := os.ReadDir(root)
	if err != nil {
		return nil, fmt.Errorf("cannot read %s: %w", root, err)
	}
	dbg("root has %d entries", len(entries))

	type workItem struct {
		name  string
		path  string
		isDir bool
	}

	items := make([]workItem, 0, len(entries))
	for _, e := range entries {
		name := e.Name()
		if strings.HasPrefix(name, ".") && root != "/" {
			continue
		}
		items = append(items, workItem{name, filepath.Join(root, name), e.IsDir()})
	}
	dbg("will scan %d items", len(items))

	sem := make(chan struct{}, 8)
	var wg sync.WaitGroup
	results := make([]DirEntry, len(items))
	var scanned atomic.Int64

	done := make(chan struct{})
	if !debugMode {
		go func() {
			frames := []string{"|", "/", "-", "\\"}
			i := 0
			for {
				select {
				case <-done:
					fmt.Printf("\r%s\r", strings.Repeat(" ", 60))
					return
				default:
					fmt.Printf("\r  %s%s%s scanning... %s%d%s items measured",
						cyan, frames[i%len(frames)], reset,
						yellow, scanned.Load(), reset)
					i++
					time.Sleep(80 * time.Millisecond)
				}
			}
		}()
	}

	for i, item := range items {
		wg.Add(1)
		i, item := i, item
		go func() {
			defer wg.Done()
			defer func() {
				if r := recover(); r != nil {
					dbgErr("PANIC in goroutine for %s: %v", item.path, r)
				}
			}()
			sem <- struct{}{}
			defer func() { <-sem }()

			var size int64
			if item.isDir {
				size = scanDir(item.path)
			} else {
				info, err := os.Stat(item.path)
				if err != nil {
					dbgErr("Stat: %s => %v", item.path, err)
				} else {
					size = info.Size()
				}
			}
			results[i] = DirEntry{item.name, item.path, size, item.isDir}
			scanned.Add(1)
		}()
	}

	wg.Wait()
	close(done)
	if !debugMode {
		time.Sleep(120 * time.Millisecond)
	}

	sort.Slice(results, func(i, j int) bool {
		return results[i].Size > results[j].Size
	})

	var nonEmpty []DirEntry
	for _, r := range results {
		if r.Name != "" {
			nonEmpty = append(nonEmpty, r)
		}
	}
	dbg("scanTopLevel done, %d results", len(nonEmpty))
	return nonEmpty, nil
}

// - tree (lwk duno how it does work)
func renderTree(path string, prefix string, depth int, maxDepth int) {
	if depth > maxDepth {
		return
	}
	dbg("renderTree depth=%d path=%s", depth, path)

	entries, err := os.ReadDir(path)
	if err != nil {
		dbgErr("renderTree ReadDir: %s => %v", path, err)
		return
	}

	type child struct {
		name  string
		path  string
		size  int64
		isDir bool
	}

	children := make([]child, 0, len(entries))
	for _, e := range entries {
		if strings.HasPrefix(e.Name(), ".") {
			continue
		}
		full := filepath.Join(path, e.Name())
		var sz int64
		if e.IsDir() {
			sz = scanDir(full)
		} else {
			info, _ := e.Info()
			if info != nil {
				sz = info.Size()
			}
		}
		children = append(children, child{e.Name(), full, sz, e.IsDir()})
	}

	sort.Slice(children, func(i, j int) bool {
		return children[i].size > children[j].size
	})

	const maxShow = 12
	shown := children
	hidden := 0
	if len(children) > maxShow {
		shown = children[:maxShow]
		hidden = len(children) - maxShow
	}

	for i, c := range shown {
		isLast := i == len(shown)-1 && hidden == 0
		connector := "├── "
		childPrefix := prefix + "│   "
		if isLast {
			connector = "└── "
			childPrefix = prefix + "    "
		}

		label := c.name
		if c.isDir {
			label = bold + cyan + c.name + "/" + reset
		}
		padLen := 30 - utf8.RuneCountInString(c.name)
		if padLen < 1 {
			padLen = 1
		}
		pad := strings.Repeat(" ", padLen)

		sc := sizeColour(c.size)
		fmt.Printf("%s%s%s%s%s%s\n",
			dim+prefix+connector+reset,
			label, pad,
			sc, humanSize(c.size), reset,
		)

		if c.isDir && depth < maxDepth {
			renderTree(c.path, childPrefix, depth+1, maxDepth)
		}
	}

	if hidden > 0 {
		fmt.Printf("%s%s%s... and %d more%s\n",
			dim, prefix+"└── ", yellow, hidden, reset)
	}
}

func printResults(results []DirEntry, root string) {
	var maxSize int64
	if len(results) > 0 {
		maxSize = results[0].Size
	}

	fmt.Printf("\n  %s%s%s\n", bold+yellow, root, reset)
	fmt.Printf("  %s%-4s  %-32s  %9s%s\n", bold+white, "#", "NAME", "SIZE", reset)
	fmt.Println("  " + strings.Repeat("-", 54))

	const barWidth = 18
	for i, r := range results {
		bar := ""
		if maxSize > 0 {
			filled := int(float64(r.Size) / float64(maxSize) * barWidth)
			bar = rust + strings.Repeat("#", filled) +
				dim + strings.Repeat(".", barWidth-filled) + reset
		}
		icon := "  "
		nameClr := white
		if r.IsDir {
			icon = "> "
			nameClr = bold + cyan
		}
		name := r.Name
		if len(name) > 30 {
			name = name[:27] + "..."
		}
		sc := sizeColour(r.Size)
		num := fmt.Sprintf("%s[%2d]%s", dim+green, i+1, reset)
		if !r.IsDir {
			num = "    "
		}
		fmt.Printf("  %s %s%s%-31s%s  %s%s%s  %s\n",
			num,
			icon, nameClr, name, reset,
			sc, humanSize(r.Size), reset,
			bar,
		)
	}
	fmt.Println("\n  " + strings.Repeat("-", 54))
}

// ── Drive detection ───────────────────────────────────────────────────────────

func detectDrives() []string {
	var drives []string
	if runtime.GOOS == "windows" {
		for c := 'A'; c <= 'Z'; c++ {
			p := string(c) + ":\\"
			if _, err := os.Stat(p); err == nil {
				drives = append(drives, p)
				dbg("found drive: %s", p)
			}
		}
		return drives
	}
	data, err := os.ReadFile("/proc/mounts")
	if err != nil {
		return []string{"/"}
	}
	seen := map[string]bool{}
	skip := map[string]bool{
		"tmpfs": true, "devtmpfs": true, "sysfs": true, "proc": true,
		"cgroup": true, "devpts": true, "securityfs": true, "pstore": true,
		"bpf": true, "autofs": true, "mqueue": true, "hugetlbfs": true,
		"debugfs": true, "tracefs": true, "fusectl": true, "configfs": true,
		"ramfs": true, "squashfs": true,
	}
	for _, line := range strings.Split(string(data), "\n") {
		fields := strings.Fields(line)
		if len(fields) < 3 {
			continue
		}
		dev, mount, fstype := fields[0], fields[1], fields[2]
		if skip[fstype] || !strings.HasPrefix(dev, "/") || seen[mount] {
			continue
		}
		seen[mount] = true
		drives = append(drives, mount)
		dbg("found mount: %s (%s)", mount, fstype)
	}
	if len(drives) == 0 {
		drives = []string{"/"}
	}
	sort.Strings(drives)
	return drives
}


var reader = bufio.NewReader(os.Stdin)

func readLine(prompt string) string {
	fmt.Print(prompt)
	line, _ := reader.ReadString('\n')
	return strings.TrimRight(line, "\r\n")
}

func readInt(prompt string) int {
	for {
		raw := strings.TrimSpace(readLine(prompt))
		var n int
		if _, err := fmt.Sscanf(raw, "%d", &n); err == nil {
			return n
		}
		fmt.Printf("  %splease enter a number%s\n", red, reset)
	}
}

// ── Banner ────────────────────────────────────────────────────────────────────

func printBanner() {
	fmt.Println()
	fmt.Printf("%s%s", rust, bold)
	fmt.Println("  ██████╗ ██╗███████╗██╗  ██╗███╗   ███╗ █████╗ ██████╗ ")
	fmt.Println("  ██╔══██╗██║██╔════╝██║ ██╔╝████╗ ████║██╔══██╗██╔══██╗")
	fmt.Println("  ██║  ██║██║███████╗█████╔╝ ██╔████╔██║███████║██████╔╝")
	fmt.Println("  ██║  ██║██║╚════██║██╔═██╗ ██║╚██╔╝██║██╔══██║██╔═══╝ ")
	fmt.Println("  ██████╔╝██║███████║██║  ██╗██║ ╚═╝ ██║██║  ██║██║     ")
	fmt.Println("  ╚═════╝ ╚═╝╚══════╝╚═╝  ╚═╝╚═╝     ╚═╝╚═╝  ╚═╝╚═╝     ")
	fmt.Printf("%s", reset)
	if debugMode {
		fmt.Printf("  %sdisk space analyzer%s  ·  %s[DEBUG MODE — log: diskmap_debug.log]%s\n\n",
			dim+white, reset, yellow+bold, reset)
	} else {
		fmt.Printf("  %sdisk space analyzer%s  ·  %sgithub.com/ekisde%s\n\n",
			dim+white, reset, dim, reset)
	}
}


func drillLoop(startRoot string) {
	var history []string
	current := startRoot

	for {
		fmt.Printf("\n  %s> Scanning%s %s%s%s ...\n",
			bold+cyan, reset, yellow, current, reset)

		start := time.Now()
		results, err := scanTopLevel(current)
		if err != nil {
			fmt.Printf("  %serror: %v%s\n\n", red, err, reset)
			if len(history) > 0 {
				current = history[len(history)-1]
				history = history[:len(history)-1]
				continue
			}
			return
		}

		elapsed := time.Since(start)
		printResults(results, current)
		fmt.Printf("  %s%d items  ·  %.2fs%s\n", dim+white, len(results), elapsed.Seconds(), reset)

		fmt.Println()
		if len(history) > 0 {
			fmt.Printf("  %s[num]%s drill in  %s[t]%s tree view  %s[b]%s back  %s[q]%s quit\n",
				green, reset, cyan, reset, yellow, reset, red, reset)
		} else {
			fmt.Printf("  %s[num]%s drill in  %s[t]%s tree view  %s[q]%s quit\n",
				green, reset, cyan, reset, red, reset)
		}

		raw := strings.TrimSpace(readLine("  > "))

		switch strings.ToLower(raw) {
		case "q", "quit", "exit":
			fmt.Printf("\n  %sbye%s\n\n", dim, reset)
			return

		case "b", "back":
			if len(history) == 0 {
				fmt.Printf("  %salready at the top%s\n", yellow, reset)
			} else {
				current = history[len(history)-1]
				history = history[:len(history)-1]
			}
			continue

		case "t", "tree":
			depthStr := strings.TrimSpace(readLine(
				fmt.Sprintf("  %sMax depth%s (default 2, max 5): ", bold+white, reset)))
			maxDepth := 2
			if depthStr != "" {
				fmt.Sscanf(depthStr, "%d", &maxDepth)
				if maxDepth < 1 {
					maxDepth = 1
				}
				if maxDepth > 5 {
					maxDepth = 5
				}
			}
			dbg("tree view maxDepth=%d root=%s", maxDepth, current)

			treeStart := time.Now()
			rootSize := scanDir(current)
			sc := sizeColour(rootSize)
			padLen := 32 - len(current)
			if padLen < 1 {
				padLen = 1
			}
			fmt.Printf("\n  %s%s%s%s  %s%s%s\n\n",
				bold+yellow, current, strings.Repeat(" ", padLen), reset,
				sc, humanSize(rootSize), reset)
			renderTree(current, "  ", 1, maxDepth)
			fmt.Printf("\n  %sscanned in %.2fs%s\n", dim+white, time.Since(treeStart).Seconds(), reset)
			continue
		}

		var choice int
		if _, err := fmt.Sscanf(raw, "%d", &choice); err != nil || choice < 1 || choice > len(results) {
			fmt.Printf("  %sinvalid — enter a number, t for tree, b to go back, q to quit%s\n", red, reset)
			continue
		}

		selected := results[choice-1]
		if !selected.IsDir {
			fmt.Printf("  %s%s is a file — pick a %s> %s(directory) entry%s\n",
				yellow, selected.Name, cyan, yellow, reset)
			continue
		}

		history = append(history, current)
		current = selected.Path
		dbg("drilling into: %s", current)
	}
}

// ── Main ──────────────────────────────────────────────────────────────────────

func main() {
	for _, arg := range os.Args[1:] {
		if arg == "-debug" || arg == "--debug" {
			debugMode = true
		}
	}

	enableWindowsANSI()
	initLog()
	printBanner()

	if debugMode {
		dbg("OS: %s  ARCH: %s", runtime.GOOS, runtime.GOARCH)
	}

	// Drives
	drives := detectDrives()
	fmt.Printf("  %sDetected drives / partitions:%s\n\n", bold+white, reset)
	for i, d := range drives {
		fmt.Printf("  %s[%d]%s  %s%s%s\n", cyan, i+1, reset, yellow, d, reset)
	}
	fmt.Println()

	var selectedRoot string
	if len(drives) == 1 {
		selectedRoot = drives[0]
		fmt.Printf("  %s-> Auto-selected:%s %s%s%s\n\n", dim, reset, yellow, selectedRoot, reset)
	} else {
		choice := readInt(fmt.Sprintf("  %sSelect drive [1-%d]: %s", bold+white, len(drives), reset))
		if choice < 1 || choice > len(drives) {
			fmt.Printf("\n  %sinvalid choice%s\n\n", red, reset)
			os.Exit(1)
		}
		selectedRoot = drives[choice-1]
		dbg("selected drive: %s", selectedRoot)
	}

	fmt.Printf("\n  %sStarting path%s\n", bold+white, reset)
	fmt.Printf("  %s(press Enter to start at %s, or type a subdirectory)%s\n", dim, selectedRoot, reset)
	rawPath := readLine("  > ")
	rawPath = strings.TrimSpace(rawPath)

	startRoot := selectedRoot
	if rawPath != "" {
		if !filepath.IsAbs(rawPath) {
			rawPath = filepath.Join(selectedRoot, rawPath)
		}
		info, err := os.Stat(rawPath)
		if err != nil {
			fmt.Printf("\n  %serror: %v%s\n\n", red, err, reset)
			os.Exit(1)
		}
		if !info.IsDir() {
			fmt.Printf("\n  %s%s is not a directory%s\n\n", red, rawPath, reset)
			os.Exit(1)
		}
		startRoot = rawPath
	}

	dbg("startRoot=%s", startRoot)
	logWrite("start root: %s", startRoot)

	drillLoop(startRoot)

	// if debug
	if debugMode {
		fmt.Println("  " + strings.Repeat("-", 48))
		if len(debugErrors) == 0 {
			fmt.Printf("  %s[debug] no errors — log saved to diskmap_debug.log%s\n\n", yellow, reset)
		} else {
			fmt.Printf("  %s[debug] %d error(s) — see diskmap_debug.log%s\n\n",
				red, len(debugErrors), reset)
		}
		logWrite(strings.Repeat("-", 60))
		logWrite("total errors: %d", len(debugErrors))
		if logFile != nil {
			logFile.Close()
		}
	}
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}