package main

import (
	"bufio"
	"flag"
	"fmt"
	_ "image/png"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"
	"unicode/utf8"

	"github.com/BurntSushi/toml"
)

const (
	Reset = "\033[0m"
	Blue  = "\033[1;34m"
)

var dist string

var ansiEscapeRe = regexp.MustCompile(`\x1b\[[0-9;?]*[ -/]*[@-~]`)

type Config struct {
	Ascii        bool     `toml:"ascii"`
	Layout       []string `toml:"layout"`
	CustomLogo   string   `toml:"custom_logo"`
	LogoPosition string   `toml:"logo_position"`
}

func appendMissingConfigKeys(pathConf string, md toml.MetaData) error {
	var missing []string

	if !md.IsDefined("ascii") {
		missing = append(missing, "# Set this to \"true\" if you like ascii more than images\nascii = false")
	}

	if !md.IsDefined("layout") {
		missing = append(missing, `# Stats layout: reorder or remove items to customize your fetch
# Available: "dist", "host", "cpu", "kernel", "ram", "gpu", "de/wm", "pkgs", "shell", "terminal", "uptime"
layout = [
	"host",
    "dist",
    "cpu",
    "kernel",
    "ram",
    "gpu",
    "de/wm",
    "pkgs",
    "shell",
    "terminal",
    "uptime",
]`)
	}

	if !md.IsDefined("custom_logo") {
		missing = append(missing, "# Absolute path to a custom image or .txt file. Leave empty to use auto-detection.\ncustom_logo = \"\"")
	}

	if !md.IsDefined("logo_position") {
		missing = append(missing, "# Position of the logo block relative to the stats block: \"left\" or \"right\".\nlogo_position = \"left\"")
	}

	if len(missing) == 0 {
		return nil
	}

	f, err := os.OpenFile(pathConf, os.O_APPEND|os.O_WRONLY, 0644)
	if err != nil {
		return err
	}
	defer func(f *os.File) {
		_ = f.Close()
	}(f)

	block := "\n# Added automatically after update\n\n" + strings.Join(missing, "\n\n") + "\n"
	_, err = f.WriteString(block)
	return err
}

func getConfigPath() (string, error) {
	configDir, err := os.UserConfigDir()
	if err != nil {
		return "", err
	}
	directoryName := "gogofetch"
	fullPath := filepath.Join(configDir, directoryName, "config.toml")

	return fullPath, nil
}

func initConfig() (Config, string) {
	var conf Config
	pathConf, err := getConfigPath()
	if err != nil {
		return conf, ""
	}

	err = os.MkdirAll(filepath.Dir(pathConf), 0755)
	if err != nil {
		return Config{}, ""
	}

	if _, err := os.Stat(pathConf); os.IsNotExist(err) {
		defaultConfig := []byte(`
# Set this to "true" if you like ascii more than images
ascii = false

# Stats layout: reorder or remove items to customize your fetch
# Available: "dist", "host", "cpu", "kernel", "ram", "gpu", "de/wm", "pkgs", "shell", "terminal", "uptime"
layout = [
   "host",
    "dist",
    "cpu",
    "kernel",
    "ram",
    "gpu",
    "de/wm",
    "pkgs",
    "shell",
    "terminal",
    "uptime",
]

# Absolute path to a custom image or .txt file. Leave empty to use auto-detection.
custom_logo = ""

# Position of the logo block relative to the stats block: "left" or "right".
logo_position = "left"
`)
		err := os.WriteFile(pathConf, defaultConfig, 0644)
		if err != nil {
			return Config{}, ""
		}
	}

	md, err := toml.DecodeFile(pathConf, &conf)
	if err != nil {
		fmt.Printf("Error loading config at %s: %v\n", pathConf, err)
		return conf, pathConf
	}

	if err := appendMissingConfigKeys(pathConf, md); err != nil {
		fmt.Printf("Warning: Could not update config at %s: %v\n", pathConf, err)
	} else {
		if _, err := toml.DecodeFile(pathConf, &conf); err != nil {
			fmt.Printf("Error reloading updated config at %s: %v\n", pathConf, err)
		}
	}

	return conf, pathConf
}

func getDist() string {
	f, _ := os.Open("/etc/os-release")
	defer func(f *os.File) {
		err := f.Close()
		if err != nil {
			fmt.Println("Error closing file:", err)
		}
	}(f)
	s := bufio.NewScanner(f)
	for s.Scan() {
		t := s.Text()
		if strings.HasPrefix(t, "PRETTY_NAME") {
			distro := strings.TrimPrefix(t, "PRETTY_NAME=")
			dist = strings.Trim(distro, "\"")
			dist = strings.TrimSpace(dist)
			return dist
		}
	}
	return "no /etc/os-release file found"
}

func getRam() string {
	f, err := os.Open("/proc/meminfo")
	if err != nil {
		return "error"
	}

	var total, available int
	s := bufio.NewScanner(f)
	for s.Scan() {
		line := s.Text()
		if strings.HasPrefix(line, "MemTotal:") {
			_, err2 := fmt.Sscanf(line, "MemTotal: %d kB", &total)
			if err2 != nil {
				return ""
			}
		}
		if strings.HasPrefix(line, "MemAvailable:") {
			_, err := fmt.Sscanf(line, "MemAvailable: %d kB", &available)
			if err != nil {
				return ""
			}
		}
	}
	totalMB := total / 1024
	availableMB := available / 1024
	usedMB := totalMB - availableMB

	ram := fmt.Sprintf("%d / %d MiB (%d MiB available)", usedMB, totalMB, availableMB)
	ram = strings.TrimSpace(ram)
	return ram

}

func getCpu() string {
	f, err := os.Open("/proc/cpuinfo")
	if err != nil {
		return "error"
	}

	var cpu string
	s := bufio.NewScanner(f)
	for s.Scan() {
		line := s.Text()
		if strings.HasPrefix(line, "model name") {
			parts := strings.SplitN(line, ":", 2)
			if len(parts) > 1 {
				cpu := strings.TrimSpace(parts[1])
				return cpu
			}
		}
	}

	return strings.TrimSpace(cpu)
}

func getGpu() string {
	out, err := exec.Command("sh", "-c", `lspci -mm | awk -F'"' '$2=="VGA compatible controller" || $2=="3D controller" || $2=="Display controller"'`).Output()
	if err != nil || len(out) == 0 {
		return "Unknown GPU"
	}

	lines := strings.Split(strings.TrimSpace(string(out)), "\n")
	var gpus []string

	for _, line := range lines {
		re := regexp.MustCompile(`"([^"]+)"`)
		matches := re.FindAllStringSubmatch(line, -1)

		if len(matches) >= 3 {
			vendor := matches[1][1]
			device := matches[2][1]

			cleaned := cleanGpuString(vendor + " " + device)
			gpus = append(gpus, cleaned)
		}
	}

	if len(gpus) > 0 {
		return strings.Join(gpus, " / ")
	}
	return "GPU not found"
}

func cleanGpuString(raw string) string {
	prettyNameRe := regexp.MustCompile(`\[(.*?)]`)
	match := prettyNameRe.FindStringSubmatch(raw)

	res := raw
	if len(match) > 1 {
		vendor := ""
		if strings.Contains(strings.ToLower(raw), "nvidia") {
			vendor = "NVIDIA "
		}
		if strings.Contains(strings.ToLower(raw), "intel") {
			vendor = "Intel "
		}
		if strings.Contains(strings.ToLower(raw), "amd") {
			vendor = "AMD "
		}

		res = vendor + match[1]
	}

	re := regexp.MustCompile(`\(.*?\)|\[.*?]`)
	res = re.ReplaceAllString(res, "")

	fluff := []string{"Corporation", "Controller", "VGA compatible", "3D", "Graphics", "Device"}
	for _, word := range fluff {
		res = strings.ReplaceAll(res, word, "")
	}

	return strings.Join(strings.Fields(res), " ")
}

func getKernel() string {
	out, err := exec.Command("sh", "-c", "uname -r").Output()
	if err != nil {
		return "Error getting kernel, how did we get there?"
	}
	return strings.TrimSpace(string(out))
}

func getHostname() string {
	hostname, err := exec.Command("sh", "-c", "uname -n").Output()
	if err != nil {
		return "Error getting hostname, awful"
	}
	username, errUser := exec.Command("sh", "-c", "whoami").Output()
	if errUser != nil {
		return "Error getting username, nobody here but us chickens!"
	}
	fullHostname := fmt.Sprintf("%s@%s", strings.TrimSpace(string(username)), strings.TrimSpace(string(hostname)))
	return fullHostname
}

func checkNix() bool {
	_, err := exec.Command("sh", "-c", "command -v nix-env >/dev/null 2>&1").Output()
	if err != nil {
		return false
	}
	return true
}

func addNix() string {
	out, _ := exec.Command("sh", "-c", "nix-env -q | wc -l").Output()
	nixPkgs := strings.TrimSpace(string(out)) + " (nix)"
	return nixPkgs
}

func getPkgs() string {
	var pkgs string
	haveNix := checkNix()
	if strings.HasPrefix(dist, "Arch") {
		out, _ := exec.Command("sh", "-c", "pacman -Qq | wc -l").Output()
		pkgs = strings.TrimSpace(string(out)) + " (pacman)"
		if haveNix {
			pkgs = strings.TrimSpace(pkgs + ", " + addNix())
		}
		return pkgs
	}
	if strings.HasPrefix(dist, "Debian") {
		out, _ := exec.Command("sh", "-c", "dpkg -l | grep ^ii | wc -l").Output()
		pkgs = strings.TrimSpace(string(out)) + " (apt)"
		if haveNix {
			pkgs = strings.TrimSpace(pkgs + ", " + addNix())
		}
		return pkgs
	}
	if strings.HasPrefix(dist, "NixOS") {
		entries, err := os.ReadDir("/run/current-system/sw/bin")
		if err != nil {
			return "can't get nix packages"
		}
		return strconv.Itoa(len(entries)) + " (nix)"
	}
	return "unknown"
}

func getLogo(dist string) string {
	switch {
	case strings.HasPrefix(dist, "Arch"):
		return "logos/arch.png"
	case strings.HasPrefix(dist, "NixOS"):
		return "logos/nixos.png"
	case strings.HasPrefix(dist, "Debian"):
		return "logos/debian.png"
	default:
		return "logos/linux.png"
	}
}

func getAscii(dist string) string {
	switch {
	case strings.HasPrefix(dist, "Arch"):
		return "logos/arch.txt"
	case strings.HasPrefix(dist, "NixOS"):
		return "logos/nixos.txt"
	case strings.HasPrefix(dist, "Debian"):
		return "logos/debian.txt"
	default:
		return "logos/linux.txt"
	}
}

func printAsciiLogo(path string) string {
	content, err := os.ReadFile(path)
	if err != nil {
		return fmt.Sprintf("Error: Could not read ASCII file at %s\n", path)
	}
	return string(content)
}

func renderLogoChafa(path string, width uint) []string {
	cols := width / 8
	if cols == 0 {
		cols = 1
	}
	sizeArg := fmt.Sprintf("--size=%dx%d", cols, cols/2)
	cmd := exec.Command("chafa", path, "--format", "symbols", sizeArg, "--symbols=half")
	out, err := cmd.Output()
	if err != nil {
		return []string{fmt.Sprintf("Chafa error: %v", err)}
	}
	return splitLines(string(out))
}

func normalizeLogoPosition(pos string) string {
	value := strings.ToLower(strings.TrimSpace(pos))
	if value == "left" || value == "right" {
		return value
	}
	return "right"
}

func splitLines(content string) []string {
	normalized := strings.ReplaceAll(content, "\r\n", "\n")
	normalized = strings.TrimRight(normalized, "\n")
	if normalized == "" {
		return []string{}
	}
	return strings.Split(normalized, "\n")
}

func maxLineLen(lines []string) int {
	maxLen := 0
	for _, line := range lines {
		lineLen := visibleWidth(line)
		if lineLen > maxLen {
			maxLen = lineLen
		}
	}
	return maxLen
}

func visibleWidth(line string) int {
	stripped := ansiEscapeRe.ReplaceAllString(line, "")
	return utf8.RuneCountInString(stripped)
}

func padRightVisible(line string, targetWidth int) string {
	pad := targetWidth - visibleWidth(line)
	if pad <= 0 {
		return line
	}
	return line + strings.Repeat(" ", pad)
}

func renderSideBySide(statsLines []string, logoLines []string, position string) {
	const gap = "   "
	statsWidth := maxLineLen(statsLines)
	logoWidth := maxLineLen(logoLines)

	lineCount := len(statsLines)
	if len(logoLines) > lineCount {
		lineCount = len(logoLines)
	}

	for i := 0; i < lineCount; i++ {
		stats := ""
		logo := ""

		if i < len(statsLines) {
			stats = statsLines[i]
		}
		if i < len(logoLines) {
			logo = logoLines[i]
		}

		if position == "left" {
			fmt.Printf("%s%s%s\n", padRightVisible(logo, logoWidth), gap, stats)
			continue
		}

		fmt.Printf("%s%s%s\n", padRightVisible(stats, statsWidth), gap, logo)
	}
}

func getAbsLogoPath(relativePath string) string {
	if _, err := os.Stat(relativePath); err == nil {
		return relativePath
	}

	exePath, err := os.Executable()
	if err == nil {
		exeDir := filepath.Dir(exePath)

		nixPath := filepath.Join(exeDir, "..", "share", "gogofetch", relativePath)

		if _, err := os.Stat(nixPath); err == nil {
			return nixPath
		}
	}

	systemPath := filepath.Join("/usr/share/gogofetch", relativePath)
	if _, err := os.Stat(systemPath); err == nil {
		return systemPath
	}

	return relativePath
}

func getShell() string {
	out, err := exec.Command("ps", "-p", strconv.Itoa(os.Getppid()), "-o", "comm=").Output()
	if err != nil {
		return path.Base(os.Getenv("SHELL"))
	}
	return strings.TrimSpace(string(out))
}

func getTerminal() string {
	return os.Getenv("TERM")
}

func getDE() string {
	return os.Getenv("XDG_CURRENT_DESKTOP")
}

func getUptime() string {
	data, err := os.ReadFile("/proc/uptime")
	if err != nil {
		return "0"
	}
	uptimeStr := strings.Fields(string(data))[0]
	seconds, err := strconv.ParseFloat(uptimeStr, 64)
	if err != nil {
		_, err := fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		if err != nil {
			return "error printing uptime"
		}
		return ""
	}
	d := time.Duration(seconds) * time.Second
	formatted := d.Round(time.Second).String()
	spaced := strings.NewReplacer("h", "h ", "m", "m ", "s", "s ").Replace(formatted)
	return spaced
}

func main() {
	conf, _ := initConfig()

	if len(conf.Layout) == 0 {
		conf.Layout = []string{"host","dist", "cpu", "ram", "shell", "uptime"}
		fmt.Println("Warning: No layout specified in config, using default")
	}

	dist := getDist()
	var defaultLogo string

	if conf.CustomLogo != "" {
		defaultLogo = conf.CustomLogo
	} else if conf.Ascii {
		defaultLogo = getAscii(dist)
	} else {
		defaultLogo = getLogo(dist)
	}

	width := flag.Int("w", 190, "Logo width in lines")
	logoPath := flag.String("l", getAbsLogoPath(defaultLogo), "Path to the logo")
	logoPos := flag.String("p", normalizeLogoPosition(conf.LogoPosition), "Logo position: left or right")
	flag.Parse()
	position := normalizeLogoPosition(*logoPos)

	type Info struct {
		Label string
		Value string
	}
	var stats []Info

	for _, item := range conf.Layout {
		switch strings.ToLower(item) {
		case "host":
				stats = append(stats, Info{"host", getHostname()})
		case "dist":
			stats = append(stats, Info{"OS", dist})
		case "cpu":
			stats = append(stats, Info{"cpu", getCpu()})
		case "kernel", "krnl":
			stats = append(stats, Info{"krnl", getKernel()})
		case "ram":
			stats = append(stats, Info{"ram", getRam()})
		case "gpu":
			stats = append(stats, Info{"gpu", getGpu()})
		case "de/wm", "de", "wm":
			stats = append(stats, Info{"de/wm", getDE()})
		case "pkgs":
			stats = append(stats, Info{"pkgs", getPkgs()})
		case "shell":
			stats = append(stats, Info{"shell", getShell()})
		case "terminal", "term":
			stats = append(stats, Info{"term", getTerminal()})
		case "uptime":
			stats = append(stats, Info{"uptime", getUptime()})
		}
	}

	var statsLines []string
	for _, info := range stats {
		statsLines = append(statsLines, fmt.Sprintf("-> %s%s%s: %s", Blue, info.Label, Reset, info.Value))
	}

	var logoLines []string
	if _, err := os.Stat(*logoPath); err == nil {
		if conf.Ascii {
			logoLines = splitLines(printAsciiLogo(*logoPath))
		} else {
			logoLines = renderLogoChafa(*logoPath, uint(*width))
		}
	} else {
		logoLines = []string{fmt.Sprintf("Logo not found at: %s", *logoPath)}
	}

	renderSideBySide(statsLines, logoLines, position)
}