package main

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

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

var dist string

var logoMap = map[string]string{
	"Arch Linux": "logos/arch.png",
	"Debian":     "logos/debian.png",
	"NixOS":      "logos/nixos.png",
	"default":    "logos/linux.png",
}

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

	cpu = fmt.Sprintf(cpu)
	cpu = strings.TrimSpace(cpu)
	return cpu
}

func getGpu() string {
	out, err := exec.Command("sh", "-c", "lspci | grep -E 'VGA|3D'").Output()
	if err != nil {
		return "Unknown GPU"
	}
	line := string(out)
	parts := strings.Split(line, ": ")
	if len(parts) > 1 {
		return cleanGpuString(parts[1])
	}
	return "GPU not found"
}

func cleanGpuString(raw string) string {
	re := regexp.MustCompile(`\(.*\)`)
	res := re.ReplaceAllString(raw, "")
	res = strings.ReplaceAll(res, "Corporation", "")
	res = strings.ReplaceAll(res, "[", "")
	res = strings.ReplaceAll(res, "]", "")
	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(string(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(string(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 getAscii() string {
//	var ascii string
//	if strings.HasPrefix(dist, "Arch") {
//		content := "                    _       _      _\n     /\\            | |     | |    (_)\n    /  \\   _ __ ___| |__   | |     _ _ __  _   ___  __\n   / /\\ \\ | '__/ __| '_ \\  | |    | | '_ \\| | | \\ \\/ /\n  / ____ \\| | | (__| | | | | |____| | | | | |_| |>  <\n /_/    \\_\\_|  \\___|_| |_| |______|_|_| |_|\\__,_/_/\\_\\\n"
//		ascii = fmt.Sprintf(content)
//		return ascii
//	}
//	if strings.HasPrefix(dist, "Debian") {
//		content := "  _____           _       _\n |  __ \\         | |     (_)\n | |  | |   ___  | |__    _    __ _   _ __\n | |  | |  / _ \\ | '_ \\  | |  / _` | | '_ \\\n | |__| | |  __/ | |_) | | | | (_| | | | | |\n |_____/   \\___| |_.__/  |_|  \\__,_| |_| |_|\n"
//		ascii = fmt.Sprintf(content)
//		return ascii
//	}
//	if strings.HasPrefix(dist, "NixOS") {
//		content := " _   _ _       ____   _____\n | \\ | (_)     / __ \\ / ____|\n |  \\| |___  _| |  | | (___\n | . ` | \\ \\/ / |  | |\\___ \\\n | |\\  | |>  <| |__| |____) |\n |_| \\_|_/_/\\_\\\\____/|_____/\n"
//		ascii = fmt.Sprintf(content)
//		return ascii
//	}
//	return "unknown"
//}

func printLogo(path string, width uint) {
	cols := width / 8
	sizeArg := fmt.Sprintf("--size=%dx%d", cols, cols/2)
	cmd := exec.Command("chafa", path, sizeArg, "--symbols=half")
	cmd.Stdout = os.Stdout
	if err := cmd.Run(); err != nil {
		fmt.Println("Chafa error:", err)
	}
}

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 {
	out, err := exec.Command("sh", "-c", "uptime -p").Output()
	if err != nil {
		return "Unknown"
	}
	uptime := strings.TrimPrefix(string(out), "up ")
	return strings.TrimSpace(uptime)
}

func main() {
	type Info struct {
		Label string
		Value string
	}
	dist := getDist()
	defaultLogo, ok := logoMap[dist]
	if !ok {
		defaultLogo = logoMap["default"]
	}

	absoluteDefaultLogo := getAbsLogoPath(defaultLogo)

	stats := []Info{
		{"dist", dist},
		{"host", getHostname()},
		{"cpu", getCpu()},
		{"krnl", getKernel()},
		{"ram", getRam()},
		{"gpu", getGpu()},
		{"de/wm", getDE()},
		{"pkgs", getPkgs()},
		{"shell", getShell()},
		{"term", getTerminal()},
		{"uptime", getUptime()},
	}

	logoPath := flag.String("l", absoluteDefaultLogo, "Absolute or relative path to the logo image")
	width := flag.Int("w", 200, "Logo width in lines")
	flag.Parse()
	if _, err := os.Stat(*logoPath); os.IsNotExist(err) {
		fmt.Printf("Image '%s' not found\n", *logoPath)
		return
	}

	printLogo(*logoPath, uint(*width))

	for _, s := range stats {
		fmt.Printf("~ %s%s%s: %s\n", Blue, s.Label, Reset, s.Value)
	}

}