//go:build linux

package main

import (
	"bufio"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

type Process struct {
	PID       int
	Command   string
	CPUTime   uint64 // clock ticks
	RSSPages  uint64
	SwapBytes uint64
}

type Snapshot struct {
	Processes map[int]Process
	TotalCPU  uint64
}

func main() {
	if len(os.Args) < 2 {
		fmt.Printf("Usage: %s <keyword> [interval]\n", os.Args[0])
		fmt.Printf("\nExamples:\n")
		fmt.Printf("  %s firefox\n", os.Args[0])
		fmt.Printf("  %s firefox 1\n", os.Args[0])
		os.Exit(1)
	}

	keyword := strings.ToLower(os.Args[1])

	interval := time.Second

	if len(os.Args) >= 3 {
		seconds, err := strconv.ParseFloat(os.Args[2], 64)
		if err != nil || seconds <= 0 {
			fmt.Fprintf(os.Stderr, "invalid interval: %s\n", os.Args[2])
			os.Exit(1)
		}

		interval = time.Duration(seconds * float64(time.Second))
	}

	numCPU := uint64(os.Getuid()) // replaced below
	_ = numCPU

	cpuCount := uint64(getCPUCount())

	if cpuCount == 0 {
		cpuCount = 1
	}

	// First snapshot.
	first, err := takeSnapshot(keyword)
	if err != nil {
		fmt.Fprintf(os.Stderr, "snapshot error: %v\n", err)
		os.Exit(1)
	}

	time.Sleep(interval)

	// Second snapshot.
	second, err := takeSnapshot(keyword)
	if err != nil {
		fmt.Fprintf(os.Stderr, "snapshot error: %v\n", err)
		os.Exit(1)
	}

	result := calculate(first, second, cpuCount)

	printResult(keyword, result)
}

type Result struct {
	ProcessCount int
	CPUPercent   float64
	RAMBytes     uint64
	SwapBytes    uint64
}

func calculate(first, second Snapshot, cpuCount uint64) Result {
	var result Result

	// RAM and swap are taken from the second snapshot.
	for _, p := range second.Processes {
		result.ProcessCount++

		result.RAMBytes += pagesToBytes(p.RSSPages)
		result.SwapBytes += p.SwapBytes
	}

	/*
		CPU calculation:

		Process CPU time is measured in clock ticks.

		We calculate:

		    process_delta / total_cpu_delta * number_of_cpus * 100

		This allows CPU usage to exceed 100% when a process/group
		is using multiple cores.
	*/

	var processDelta uint64

	for pid, secondProcess := range second.Processes {
		firstProcess, exists := first.Processes[pid]
		if !exists {
			continue
		}

		if secondProcess.CPUTime >= firstProcess.CPUTime {
			processDelta += secondProcess.CPUTime - firstProcess.CPUTime
		}
	}

	totalCPUDelta := second.TotalCPU - first.TotalCPU

	if totalCPUDelta > 0 {
		result.CPUPercent =
			(float64(processDelta) / float64(totalCPUDelta)) *
				float64(cpuCount) *
				100
	}

	return result
}

func takeSnapshot(keyword string) (Snapshot, error) {
	snapshot := Snapshot{
		Processes: make(map[int]Process),
	}

	entries, err := os.ReadDir("/proc")
	if err != nil {
		return snapshot, err
	}

	for _, entry := range entries {
		if !entry.IsDir() {
			continue
		}

		pid, err := strconv.Atoi(entry.Name())
		if err != nil {
			// Not a PID directory.
			continue
		}

		process, err := readProcess(pid)
		if err != nil {
			// Process may have exited between ReadDir and reading
			// /proc/<pid>. That's normal.
			continue
		}

		commandLower := strings.ToLower(process.Command)

		if strings.Contains(commandLower, keyword) {
			snapshot.Processes[pid] = process
		}
	}

	snapshot.TotalCPU, _ = readTotalCPU()

	return snapshot, nil
}

func readProcess(pid int) (Process, error) {
	var process Process

	process.PID = pid

	// Read the complete command line.
	cmdline, err := os.ReadFile(
		filepath.Join("/proc", strconv.Itoa(pid), "cmdline"),
	)
	if err != nil {
		return process, err
	}

	/*
		/proc/<pid>/cmdline is separated by NUL bytes:

		    /usr/bin/firefox\0--some-option\0foo\0

		Convert those NULs into spaces so matching and displaying
		the command line is convenient.
	*/
	command := strings.ReplaceAll(string(cmdline), "\x00", " ")
	command = strings.TrimSpace(command)

	// Some kernel/system processes have an empty cmdline.
	// Fall back to /comm if necessary.
	if command == "" {
		comm, err := os.ReadFile(
			filepath.Join("/proc", strconv.Itoa(pid), "comm"),
		)
		if err != nil {
			return process, err
		}

		command = strings.TrimSpace(string(comm))
	}

	process.Command = command

	// CPU time from /proc/<pid>/stat.
	cpuTime, rssPages, err := readStat(pid)
	if err != nil {
		return process, err
	}

	process.CPUTime = cpuTime
	process.RSSPages = rssPages

	// Swap from /proc/<pid>/status.
	process.SwapBytes = readSwap(pid)

	return process, nil
}

func readStat(pid int) (cpuTime uint64, rssPages uint64, err error) {
	path := filepath.Join("/proc", strconv.Itoa(pid), "stat")

	data, err := os.ReadFile(path)
	if err != nil {
		return 0, 0, err
	}

	line := string(data)

	/*
		/proc/<pid>/stat contains:

		    pid (comm) state ppid ...

		The comm field itself can contain spaces and parentheses,
		so we cannot simply strings.Fields() the whole line.

		Find the LAST ')' which terminates comm.
	*/
	closeParen := strings.LastIndex(line, ")")
	if closeParen == -1 {
		return 0, 0, errors.New("invalid /proc/stat format")
	}

	rest := strings.TrimSpace(line[closeParen+1:])
	fields := strings.Fields(rest)

	/*
		After comm, fields start at:

		field 3  = state
		field 4  = ppid
		...
		field 14 = utime
		field 15 = stime
		...
		field 24 = rss

		Because fields[] starts at field 3:

		utime -> index 11
		stime -> index 12
		rss   -> index 21
	*/

	if len(fields) <= 21 {
		return 0, 0, errors.New("short /proc/stat entry")
	}

	utime, err := strconv.ParseUint(fields[11], 10, 64)
	if err != nil {
		return 0, 0, err
	}

	stime, err := strconv.ParseUint(fields[12], 10, 64)
	if err != nil {
		return 0, 0, err
	}

	rss, err := strconv.ParseInt(fields[21], 10, 64)
	if err != nil {
		return 0, 0, err
	}

	if rss < 0 {
		rss = 0
	}

	return utime + stime, uint64(rss), nil
}

func readSwap(pid int) uint64 {
	path := filepath.Join("/proc", strconv.Itoa(pid), "status")

	file, err := os.Open(path)
	if err != nil {
		return 0
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)

	for scanner.Scan() {
		line := scanner.Text()

		if !strings.HasPrefix(line, "VmSwap:") {
			continue
		}

		fields := strings.Fields(line)

		if len(fields) < 2 {
			return 0
		}

		/*
			VmSwap is normally reported in kB.
		*/
		kb, err := strconv.ParseUint(fields[1], 10, 64)
		if err != nil {
			return 0
		}

		return kb * 1024
	}

	return 0
}

func readTotalCPU() (uint64, error) {
	data, err := os.ReadFile("/proc/stat")
	if err != nil {
		return 0, err
	}

	scanner := bufio.NewScanner(strings.NewReader(string(data)))

	for scanner.Scan() {
		line := scanner.Text()

		if !strings.HasPrefix(line, "cpu ") {
			continue
		}

		fields := strings.Fields(line)

		if len(fields) < 2 {
			return 0, errors.New("invalid /proc/stat cpu line")
		}

		var total uint64

		for _, field := range fields[1:] {
			value, err := strconv.ParseUint(field, 10, 64)
			if err != nil {
				return 0, err
			}

			total += value
		}

		return total, nil
	}

	return 0, errors.New("cpu line not found")
}

func pagesToBytes(pages uint64) uint64 {
	pageSize := uint64(os.Getpagesize())
	return pages * pageSize
}

func getCPUCount() int {
	count := os.Getenv("NUMBER_OF_PROCESSORS")

	if count != "" {
		if n, err := strconv.Atoi(count); err == nil && n > 0 {
			return n
		}
	}

	// Linux: use /proc/cpuinfo.
	file, err := os.Open("/proc/cpuinfo")
	if err != nil {
		return 1
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)

	cpus := 0

	for scanner.Scan() {
		if strings.HasPrefix(scanner.Text(), "processor") {
			cpus++
		}
	}

	if cpus == 0 {
		return 1
	}

	return cpus
}

func printResult(keyword string, result Result) {
	fmt.Printf("\n")
	fmt.Printf("%s\n", keyword)
	fmt.Printf("  Processes : %d\n", result.ProcessCount)
	fmt.Printf("  CPU       : %.1f%%\n", result.CPUPercent)
	fmt.Printf("  RAM       : %s\n", formatBytes(result.RAMBytes))
	fmt.Printf("  Swap      : %s\n", formatBytes(result.SwapBytes))
	fmt.Printf("\n")
}

func formatBytes(bytes uint64) string {
	const (
		KB = 1024
		MB = 1024 * KB
		GB = 1024 * MB
	)

	switch {
	case bytes >= GB:
		return fmt.Sprintf("%.2f GB", float64(bytes)/float64(GB))

	case bytes >= MB:
		return fmt.Sprintf("%.0f MB", float64(bytes)/float64(MB))

	case bytes >= KB:
		return fmt.Sprintf("%.0f KB", float64(bytes)/float64(KB))

	default:
		return fmt.Sprintf("%d B", bytes)
	}
}