package main

import (
	"log"
	"os"
	"os/exec"
	"time"

	"github.com/NVIDIA/go-nvml/pkg/nvml"
	"github.com/shirou/gopsutil/v3/cpu"
	"github.com/shirou/gopsutil/v3/disk"
	"github.com/shirou/gopsutil/v3/net"
)

const (
	checkInterval     = 10 * time.Second
	monitoringPeriod = 5 * time.Minute
	cpuThreshold     = 20.0    // percentage
	gpuThreshold     = 20.0    // percentage
	diskThreshold    = 5 * 1024 * 1024    // 5 MB/s
	networkThreshold = 1 * 1024 * 1024    // 1 MB/s
)

type ResourceUsage struct {
	timestamp  time.Time
	cpuUsage   float64
	gpuUsage   float64
	diskIO     uint64
	networkIO  uint64
}

func main() {
	// Check if running as root
	if os.Geteuid() != 0 {
		log.Fatal("This program must be run as root")
	}

	// Initialize NVML for GPU monitoring
	ret := nvml.Init()
	if ret != nvml.SUCCESS {
		log.Printf("Warning: Could not initialize NVML: %v", ret)
	}
	defer nvml.Shutdown()

	usageHistory := make([]ResourceUsage, 0)
	ticker := time.NewTicker(checkInterval)
	defer ticker.Stop()

	log.Printf("Starting idle monitoring. System will suspend when:\n")
	log.Printf("- CPU usage < %.1f%%\n", cpuThreshold)
	log.Printf("- GPU usage < %.1f%%\n", gpuThreshold)
	log.Printf("- Disk I/O < %.1f MB/s\n", float64(diskThreshold)/(1024*1024))
	log.Printf("- Network I/O < %.1f MB/s\n", float64(networkThreshold)/(1024*1024))
	log.Printf("Over the last %v\n", monitoringPeriod)

	for range ticker.C {
		usage := getCurrentUsage()
		usageHistory = append(usageHistory, usage)

		// Remove entries older than monitoring period
		cutoff := time.Now().Add(-monitoringPeriod)
		for i, u := range usageHistory {
			if u.timestamp.After(cutoff) {
				usageHistory = usageHistory[i:]
				break
			}
		}

		if len(usageHistory) > 0 && isSystemIdle(usageHistory) {
			log.Println("System has been idle for the monitoring period. Suspending...")
			if err := suspendSystem(); err != nil {
				log.Printf("Failed to suspend system: %v", err)
			}
		}
	}
}

func getCurrentUsage() ResourceUsage {
	usage := ResourceUsage{
		timestamp: time.Now(),
	}

	// Get CPU usage
	if cpuPercent, err := cpu.Percent(0, false); err == nil && len(cpuPercent) > 0 {
		usage.cpuUsage = cpuPercent[0]
	}

	// Get GPU usage
	count, ret := nvml.DeviceGetCount()
	if ret == nvml.SUCCESS && count > 0 {
		device, ret := nvml.DeviceGetHandleByIndex(0)
		if ret == nvml.SUCCESS {
			utilization, ret := device.GetUtilizationRates()
			if ret == nvml.SUCCESS {
				usage.gpuUsage = float64(utilization.Gpu)
			}
		}
	}

	// Get disk I/O
	if diskStats, err := disk.IOCounters(); err == nil {
		var totalIO uint64
		for _, stat := range diskStats {
			totalIO += stat.ReadBytes + stat.WriteBytes
		}
		usage.diskIO = totalIO
	}

	// Get network I/O
	if netStats, err := net.IOCounters(false); err == nil && len(netStats) > 0 {
		usage.networkIO = netStats[0].BytesSent + netStats[0].BytesRecv
	}

	return usage
}

func isSystemIdle(history []ResourceUsage) bool {
	if len(history) < 2 {
		return false
	}

	var avgCPU, avgGPU float64
	samples := len(history)

	for _, usage := range history {
		avgCPU += usage.cpuUsage
		avgGPU += usage.gpuUsage
	}

	// Calculate I/O rates using first and last samples
	duration := history[samples-1].timestamp.Sub(history[0].timestamp).Seconds()
	diskIORate := float64(history[samples-1].diskIO-history[0].diskIO) / duration
	netIORate := float64(history[samples-1].networkIO-history[0].networkIO) / duration

	avgCPU /= float64(samples)
	avgGPU /= float64(samples)

	return avgCPU < cpuThreshold &&
		avgGPU < gpuThreshold &&
		diskIORate < float64(diskThreshold) &&
		netIORate < float64(networkThreshold)
}

func suspendSystem() error {
	cmd := exec.Command("systemctl", "suspend")
	return cmd.Run()
}