package platform

import (
	"fmt"
	"os/exec"
	"strconv"
	"strings"
	"time"
)

const nvidiaRecoverHelper = "/usr/local/bin/bee-nvidia-recover"

func runNvidiaRecover(args ...string) (string, error) {
	helperArgs := append([]string{nvidiaRecoverHelper}, args...)
	if _, err := exec.LookPath("systemd-run"); err == nil {
		unit := fmt.Sprintf("bee-nvidia-recover-%d", time.Now().UnixNano())
		cmdArgs := []string{
			"systemd-run",
			"--quiet",
			"--pipe",
			"--wait",
			"--collect",
			"--service-type=oneshot",
			"--unit", unit,
		}
		cmdArgs = append(cmdArgs, helperArgs...)
		raw, err := exec.Command("sudo", cmdArgs...).CombinedOutput()
		return string(raw), err
	}
	raw, err := exec.Command("sudo", helperArgs...).CombinedOutput()
	return string(raw), err
}

func resetNvidiaGPU(index int) (string, error) {
	if index < 0 {
		return "", fmt.Errorf("gpu index must be >= 0")
	}
	out, err := runNvidiaRecover("reset-gpu", strconv.Itoa(index))
	if strings.TrimSpace(out) == "" && err == nil {
		out = "GPU reset completed.\n"
	}
	return out, err
}

func restartNvidiaDrivers() (string, error) {
	out, err := runNvidiaRecover("restart-drivers")
	if strings.TrimSpace(out) == "" && err == nil {
		out = "NVIDIA drivers restarted.\n"
	}
	return out, err
}