package platform

import (
	"bufio"
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"
)

// Estimated wall-clock durations for each SAT/validate test, derived from real
// production logs in _benchmark/_v8/.
//
// Rule: whenever the commands, timeout parameters, or number of sub-jobs inside
// the corresponding Run*Pack function change, re-measure the wall-clock duration
// from actual task logs and update the matching constant here.
//
// Sources:
//   - SATEstimatedCPUValidateSec:                 xFusion v8.6 — 62 s
//   - SATEstimatedMemoryValidateSec:               xFusion v8.6 — 68 s
//   - SATEstimatedNvidiaGPUValidateSec:            xFusion v8.6/v8.22 — 77–87 s/GPU (measured per-GPU; re-measure after switch to all-GPU simultaneous)
//   - SATEstimatedNvidiaGPUStressSec:              xFusion v8.6/v8.22 — 444–448 s/GPU (measured per-GPU; re-measure after switch to all-GPU simultaneous)
//   - SATEstimatedNvidiaTargetedStressSec:         xFusion v8.6/v8.22 — 347–348 s/GPU (measured per-GPU; re-measure after switch to all-GPU simultaneous)
//   - SATEstimatedNvidiaTargetedPowerSec:          MSI v8.22 / xFusion v8.6 — 346–351 s/GPU (measured per-GPU; re-measure after switch to all-GPU simultaneous)
//   - SATEstimatedNvidiaPulseTestSec:              xFusion v8.6 — 4 926 s / 8 GPU (all simultaneous)
//   - SATEstimatedNvidiaInterconnectSec:           xFusion v8.6/v8.22 — 210–384 s / 8 GPU (all simultaneous)
//   - SATEstimatedNvidiaBandwidthSec:              xFusion v8.6/v8.22 — 2 664–2 688 s / 8 GPU (all simultaneous);
//     on multi-socket systems now runs as up to 3 passes (per-socket + all-GPU) — re-measure and bump this once
//     real multi-socket task logs exist, current value only covers the single-pass/single-socket case.
const (
	// CPU stress: stress-ng 60 s + lscpu/sensors overhead.
	SATEstimatedCPUValidateSec = 65
	// CPU stress: stress-ng 1800 s (stress mode default).
	SATEstimatedCPUStressSec = 1800

	// RAM: memtester 256 MB / 1 pass.
	SATEstimatedMemoryValidateSec = 70
	// RAM: memtester 512 MB / 1 pass (extrapolated from validate timing, linear with size).
	SATEstimatedMemoryStressSec = 140

	// TPM capabilities, PCR values, and existing self-test result queries.
	SATEstimatedTPMValidateSec = 5

	// NVIDIA dcgmi diag Level 2 (medium), all GPUs simultaneously.
	SATEstimatedNvidiaGPUValidateSec = 85
	// NVIDIA dcgmi diag Level 3 (targeted stress), all GPUs simultaneously.
	SATEstimatedNvidiaGPUStressSec = 450

	// NVIDIA dcgmi targeted_stress 300 s + overhead, all GPUs simultaneously.
	SATEstimatedNvidiaTargetedStressSec = 350
	// NVIDIA dcgmi targeted_power 300 s + overhead, all GPUs simultaneously.
	SATEstimatedNvidiaTargetedPowerSec = 350

	// NVIDIA dcgmi pulse_test, all GPUs simultaneously (not per-GPU).
	SATEstimatedNvidiaPulseTestSec = 5000

	// NCCL all_reduce_perf, all GPUs simultaneously.
	SATEstimatedNvidiaInterconnectSec = 300
	// nvbandwidth, all GPUs simultaneously. Tool runs all built-in tests
	// without a user-configurable time limit; duration is determined by nvbandwidth itself.
	SATEstimatedNvidiaBandwidthSec = 2700
)

// satJobBoundaryHook, if set, is called with a SAT job's log file name right
// after that job's output has been written to disk — a natural point for an
// external blackbox sync to pick up newly-completed data promptly instead of
// waiting out its own adaptive schedule. Nil by default (no-op); set once
// via SetJobBoundaryHook by whichever process wires up blackbox.
var satJobBoundaryHook func(jobName string)

// SetJobBoundaryHook installs the callback invoked after each SAT job
// finishes and its log file has been written. Pass nil to clear it.
func SetJobBoundaryHook(hook func(jobName string)) {
	satJobBoundaryHook = hook
}

// satSyncBracketHook, if set, is called synchronously immediately before and
// after a satJob marked syncBracket runs — i.e. around the actual load
// command of a diagnostic (nvbandwidth, memtester, stress-ng, dcgmi diag...),
// not the cheap discovery/inventory steps around it. phase is "before" or
// "after". Unlike satJobBoundaryHook (fire-and-forget, fires after every
// job), this is meant to block until an external blackbox sync has actually
// reached removable media, so that a crash during the load itself still
// leaves durable evidence that the load started (and, on the far side, that
// it finished). Nil by default. An error is logged, never fails the job —
// a stuck blackbox target must not block the diagnostic the operator asked
// for.
var satSyncBracketHook func(jobName, phase string) error

// SetSyncBracketHook installs the callback invoked synchronously before and
// after a syncBracket-marked SAT job. Pass nil to clear it.
func SetSyncBracketHook(hook func(jobName, phase string) error) {
	satSyncBracketHook = hook
}

func runSyncBracketHook(job satJob, phase string, logFunc func(string)) {
	if !job.syncBracket || satSyncBracketHook == nil {
		return
	}
	if err := satSyncBracketHook(job.name, phase); err != nil && logFunc != nil {
		logFunc(fmt.Sprintf("%s: blackbox sync wait (%s) did not complete cleanly: %v", job.name, phase, err))
	}
}

var (
	satExecCommand  = exec.Command
	satLookPath     = exec.LookPath
	satGlob         = filepath.Glob
	satStat         = os.Stat
	satFreeMemBytes = freeMemBytes

	rocmSMIExecutableGlobs = []string{
		"/opt/rocm/bin/rocm-smi",
		"/opt/rocm-*/bin/rocm-smi",
	}
	rocmSMIScriptGlobs = []string{
		"/opt/rocm/libexec/rocm_smi/rocm_smi.py",
		"/opt/rocm-*/libexec/rocm_smi/rocm_smi.py",
	}
	rvsExecutableGlobs = []string{
		"/opt/rocm/bin/rvs",
		"/opt/rocm-*/bin/rvs",
	}
	dcgmProfTesterCandidates = []string{
		"dcgmproftester",
		"dcgmproftester13",
		"dcgmproftester12",
		"dcgmproftester11",
	}
)

// streamExecOutput runs cmd and streams each output line to logFunc (if non-nil).
// If livePath is non-empty, each line is also appended to that file as it
// arrives — so the command's own output already exists on disk (and is thus
// pickable up by a concurrent blackbox sync) while it's still running,
// instead of only appearing once the whole pack finishes writing the final
// job file. Best-effort: a failure to open/write livePath never fails the job.
// Returns combined stdout+stderr as a byte slice.
func streamExecOutput(cmd *exec.Cmd, logFunc func(string), livePath string) ([]byte, error) {
	pr, pw := io.Pipe()
	cmd.Stdout = pw
	cmd.Stderr = pw

	var liveFile *os.File
	if livePath != "" {
		if f, err := os.OpenFile(livePath, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0644); err == nil {
			liveFile = f
		}
	}

	var buf bytes.Buffer
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		if liveFile != nil {
			defer liveFile.Close()
		}
		scanner := bufio.NewScanner(pr)
		for scanner.Scan() {
			line := scanner.Text()
			buf.WriteString(line + "\n")
			if liveFile != nil {
				_, _ = liveFile.WriteString(line + "\n")
			}
			if logFunc != nil {
				logFunc(line)
			}
		}
	}()

	err := cmd.Start()
	if err != nil {
		_ = pw.Close()
		wg.Wait()
		return nil, err
	}
	waitErr := cmd.Wait()
	_ = pw.Close()
	wg.Wait()
	return buf.Bytes(), waitErr
}

// satJob describes one command and the checks needed to turn its process and
// output results into a SAT verdict.
type satJob struct {
	name string
	cmd  []string
	env  []string // extra env vars (appended to os.Environ)
	// validate checks successful command output against the tool's documented
	// result format. It may inspect artifacts from earlier jobs in runDir.
	// FAILED means the tool proved a test failure; UNSUPPORTED means the pinned
	// output contract could not be recognized safely.
	validate   func(runDir string, out []byte) (status, detail string)
	collectGPU bool  // collect GPU metrics via nvidia-smi while this job runs
	gpuIndices []int // GPU indices to collect metrics for (empty = all)
	// informational marks a preflight/metadata job (e.g. dcgmi discovery) whose
	// failure shouldn't flip the pack's overall status — the diagnostic jobs
	// that follow it are the actual test of GPU health.
	informational bool
	// retries is the number of extra attempts (with a short backoff) if the
	// job's first run fails. Used for jobs racing nv-hostengine startup.
	retries int
	// syncBracket marks a job as the actual load step of a pack (as opposed
	// to the cheap inventory/discovery steps around it) — see
	// satSyncBracketHook. Set this on the command that can hang or crash the
	// host, not on nvidia-smi/dcgmi discovery calls.
	syncBracket bool
}

type satStats struct {
	OK            int
	Failed        int
	Unsupported   int
	Informational int
}

func withNvidiaPersistenceMode(jobs ...satJob) []satJob {
	out := make([]satJob, 0, len(jobs)+1)
	out = append(out, satJob{
		name: "00-nvidia-smi-persistence-mode.log",
		cmd:  []string{"nvidia-smi", "-pm", "1"},
	})
	out = append(out, jobs...)
	return out
}

func nvidiaSATJobs() []satJob {
	return withNvidiaPersistenceMode(
		satJob{name: "01-nvidia-smi-q.log", cmd: []string{"nvidia-smi", "-q"}},
		satJob{name: "02-dmidecode-baseboard.log", cmd: []string{"dmidecode", "-t", "baseboard"}},
		satJob{name: "03-dmidecode-system.log", cmd: []string{"dmidecode", "-t", "system"}},
		// nvidia-bug-report.sh appends .gz to --output-file whenever gzip is
		// available, so the artifact actually lands at nvidia-bug-report.log.gz.
		satJob{name: "04-nvidia-bug-report.log.gz", cmd: []string{"nvidia-bug-report.sh", "--output-file", "{{run_dir}}/nvidia-bug-report.log"}},
		satJob{name: "05-bee-gpu-burn.log", cmd: []string{"bee-gpu-burn", "--seconds", "5", "--size-mb", "64"}, syncBracket: true},
	)
}

func nvidiaDCGMJobs(diagLevel int, gpuIndices []int) []satJob {
	if diagLevel < 1 || diagLevel > 4 {
		diagLevel = 3
	}
	diagArgs := append([]string{"dcgmi", "diag", "-r", strconv.Itoa(diagLevel)}, nvidiaDCGMDiagDebugArgs("dcgmi-diag")...)
	if len(gpuIndices) > 0 {
		ids := make([]string, len(gpuIndices))
		for i, idx := range gpuIndices {
			ids[i] = strconv.Itoa(idx)
		}
		diagArgs = append(diagArgs, "-i", strings.Join(ids, ","))
	}
	return withNvidiaPersistenceMode(
		satJob{name: "01-nvidia-smi-q.log", cmd: []string{"nvidia-smi", "-q"}},
		satJob{name: "02-dmidecode-baseboard.log", cmd: []string{"dmidecode", "-t", "baseboard"}},
		satJob{name: "03-dmidecode-system.log", cmd: []string{"dmidecode", "-t", "system"}},
		satJob{name: "04-dcgmi-discovery.log", cmd: []string{"dcgmi", "discovery", "-l"}, informational: true, retries: 2},
		satJob{name: "05-dcgmi-diag.log", cmd: diagArgs, gpuIndices: gpuIndices, syncBracket: true},
	)
}

// nvidiaDCGMDiagDebugArgs returns flags that make dcgmi diag emit its full
// internal nvvs debug log (e.g. the real reason behind a terse "Detected
// driver major version 0" failure) into the SAT run dir instead of the
// default /var/log/nvidia-dcgm/nvvs.log, which may not be captured otherwise.
func nvidiaDCGMDiagDebugArgs(logPrefix string) []string {
	return []string{"-v", "-d", "DEBUG", "--debugLogFile", "{{run_dir}}/" + logPrefix + "-debug.log"}
}

func nvidiaDCGMNamedDiagCommand(name string, durationSec int, gpuIndices []int) []string {
	args := append([]string{"dcgmi", "diag", "-r", name}, nvidiaDCGMDiagDebugArgs("dcgmi-"+strings.ReplaceAll(name, "_", "-"))...)
	if durationSec > 0 {
		args = append(args, "-p", fmt.Sprintf("%s.test_duration=%d", name, durationSec))
	}
	if len(gpuIndices) > 0 {
		args = append(args, "-i", joinIndexList(gpuIndices))
	}
	return args
}

func normalizeNvidiaBurnDuration(durationSec int) int {
	if durationSec <= 0 {
		return 300
	}
	return durationSec
}

func nvidiaVisibleDevicesEnv(gpuIndices []int) []string {
	if len(gpuIndices) == 0 {
		return nil
	}
	return []string{
		"CUDA_DEVICE_ORDER=PCI_BUS_ID",
		"CUDA_VISIBLE_DEVICES=" + joinIndexList(gpuIndices),
	}
}

func runAcceptancePackCtx(ctx context.Context, baseDir, prefix string, jobs []satJob, logFunc func(string)) (string, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	if baseDir == "" {
		baseDir = "/var/log/bee-sat"
	}
	ts := time.Now().UTC().Format("20060102-150405")
	runDir := filepath.Join(baseDir, prefix+"-"+ts)
	if err := os.MkdirAll(runDir, 0755); err != nil {
		return "", err
	}
	verboseLog := filepath.Join(runDir, "verbose.log")

	var summary strings.Builder
	stats := satStats{}
	nvidiaPack := isNvidiaAcceptancePack(prefix)
	perGPU := map[int]*nvidiaGPUStatusFile{}
	selectedGPUIndices := map[int]struct{}{}
	fmt.Fprintf(&summary, "run_at_utc=%s\n", time.Now().UTC().Format(time.RFC3339))
	for _, job := range jobs {
		if ctx.Err() != nil {
			break
		}
		for _, idx := range job.gpuIndices {
			selectedGPUIndices[idx] = struct{}{}
			status := perGPU[idx]
			if status == nil {
				status = &nvidiaGPUStatusFile{Index: idx}
				perGPU[idx] = status
			}
			status.Selected = true
		}
		cmd := make([]string, 0, len(job.cmd))
		for _, arg := range job.cmd {
			cmd = append(cmd, strings.ReplaceAll(arg, "{{run_dir}}", runDir))
		}

		var out []byte
		var err error
		jobDetail := ""

		if nvidiaPack && nvidiaJobNeedsHealthCheck(job) {
			if msg, healthErr := checkNvidiaJobHealth(job.gpuIndices); healthErr != nil {
				if logFunc != nil {
					logFunc(msg)
				}
				out = []byte(msg + "\n")
				err = healthErr
				jobDetail = msg
			}
		}

		if err == nil {
			runSyncBracketHook(job, "before", logFunc)
			for attempt := 0; ; attempt++ {
				if job.collectGPU {
					out, err = runSATCommandWithMetrics(ctx, verboseLog, job.name, cmd, job.env, job.gpuIndices, runDir, logFunc)
				} else {
					out, err = runSATCommandCtx(ctx, verboseLog, job.name, cmd, job.env, logFunc, filepath.Join(runDir, job.name))
				}
				if err == nil || attempt >= job.retries || ctx.Err() != nil {
					break
				}
				if logFunc != nil {
					logFunc(fmt.Sprintf("%s: retrying after failure (attempt %d/%d)", job.name, attempt+1, job.retries))
				}
				time.Sleep(2 * time.Second)
			}
		}

		if nvidiaPack && nvidiaJobNeedsHealthCheck(job) {
			if msg, healthErr := checkNvidiaJobHealth(job.gpuIndices); healthErr != nil {
				if logFunc != nil {
					logFunc(msg)
				}
				if len(out) > 0 && !bytes.HasSuffix(out, []byte("\n")) {
					out = append(out, '\n')
				}
				out = append(out, []byte(msg+"\n")...)
				if err == nil {
					err = healthErr
				}
				jobDetail = msg
			}
		}

		status, rc := classifySATResult(job.name, out, err)
		validationDetail := singleLineSATDetail(jobDetail)
		validated := false
		if status == "OK" && job.validate != nil {
			status, validationDetail = validateSATJobOutput(job, runDir, out)
			validated = true
		}
		if validated {
			if validationDetail != "" {
				if len(out) > 0 && !bytes.HasSuffix(out, []byte("\n")) {
					out = append(out, '\n')
				}
				out = append(out, []byte(fmt.Sprintf("[bee-validator] %s: %s\n", status, validationDetail))...)
			}
		}

		if writeErr := os.WriteFile(filepath.Join(runDir, job.name), out, 0644); writeErr != nil {
			return "", writeErr
		}
		if satJobBoundaryHook != nil {
			satJobBoundaryHook(job.name)
		}
		runSyncBracketHook(job, "after", logFunc)
		if ctx.Err() != nil {
			return "", ctx.Err()
		}
		if job.informational && status != "OK" {
			stats.Informational++
		} else {
			stats.Add(status)
		}
		if nvidiaPack && len(job.gpuIndices) > 0 && nvidiaJobNeedsHealthCheck(job) {
			for _, idx := range job.gpuIndices {
				updateNvidiaGPUStatus(perGPU, idx, status, job.name, string(out))
			}
		}
		key := strings.TrimSuffix(strings.TrimPrefix(job.name, "0"), ".log")
		fmt.Fprintf(&summary, "%s_rc=%d\n", key, rc)
		fmt.Fprintf(&summary, "%s_status=%s\n", key, status)
		if validationDetail != "" {
			fmt.Fprintf(&summary, "%s_detail=%s\n", key, singleLineSATDetail(validationDetail))
		}
	}
	writeSATStats(&summary, stats)
	if err := os.WriteFile(filepath.Join(runDir, "summary.txt"), []byte(summary.String()), 0644); err != nil {
		return "", err
	}
	if nvidiaPack {
		if err := writeNvidiaGPUStatusFiles(runDir, stats.Overall(), perGPU, selectedGPUIndices); err != nil {
			return "", err
		}
	}

	return runDir, nil
}

func isNvidiaAcceptancePack(prefix string) bool {
	return strings.HasPrefix(prefix, "gpu-nvidia") || prefix == "nccl-tests"
}

func validateSATJobOutput(job satJob, runDir string, out []byte) (string, string) {
	status, detail := job.validate(runDir, out)
	status = strings.ToUpper(strings.TrimSpace(status))
	switch status {
	case "OK", "FAILED", "UNSUPPORTED", "PARTIAL":
		return status, strings.TrimSpace(detail)
	default:
		return "UNSUPPORTED", fmt.Sprintf("validator returned invalid status %q", status)
	}
}

func singleLineSATDetail(detail string) string {
	return strings.Join(strings.Fields(detail), " ")
}

func updateNvidiaGPUStatus(perGPU map[int]*nvidiaGPUStatusFile, idx int, status, jobName, detail string) {
	entry := perGPU[idx]
	if entry == nil {
		entry = &nvidiaGPUStatusFile{Index: idx}
		perGPU[idx] = entry
	}
	if nvidiaSATStatusSeverity(status) >= nvidiaSATStatusSeverity(entry.RunStatus) {
		entry.RunStatus = status
		entry.FailingJob = jobName
		entry.Reason = firstLine(detail)
	}
}

func writeNvidiaGPUStatusFiles(runDir, overall string, perGPU map[int]*nvidiaGPUStatusFile, selected map[int]struct{}) error {
	health, err := readNvidiaGPUHealth()
	if err == nil {
		for _, gpu := range health {
			entry := perGPU[gpu.Index]
			if entry == nil {
				entry = &nvidiaGPUStatusFile{Index: gpu.Index}
				perGPU[gpu.Index] = entry
			}
			entry.Name = gpu.Name
			entry.Observed = true
			entry.HealthRaw = gpu.RawLine
			if gpu.NeedsReset {
				entry.Health = "RESET_REQUIRED"
				if entry.RunStatus == "" || nvidiaSATStatusSeverity("FAILED") >= nvidiaSATStatusSeverity(entry.RunStatus) {
					entry.RunStatus = "FAILED"
					if strings.TrimSpace(entry.Reason) == "" {
						entry.Reason = "GPU requires reset"
					}
				}
			} else {
				entry.Health = "OK"
			}
		}
	}
	for idx := range selected {
		entry := perGPU[idx]
		if entry == nil {
			entry = &nvidiaGPUStatusFile{Index: idx}
			perGPU[idx] = entry
		}
		entry.Selected = true
	}
	var indices []int
	for idx := range perGPU {
		indices = append(indices, idx)
	}
	sort.Ints(indices)
	for _, idx := range indices {
		entry := perGPU[idx]
		if entry.RunStatus == "" {
			entry.RunStatus = overall
		}
		if entry.Health == "" {
			entry.Health = "UNKNOWN"
		}
		if entry.Name == "" {
			entry.Name = "Unknown GPU"
		}
		var body strings.Builder
		fmt.Fprintf(&body, "gpu_index=%d\n", entry.Index)
		fmt.Fprintf(&body, "gpu_name=%s\n", entry.Name)
		fmt.Fprintf(&body, "selected=%t\n", entry.Selected)
		fmt.Fprintf(&body, "observed=%t\n", entry.Observed)
		fmt.Fprintf(&body, "run_status=%s\n", entry.RunStatus)
		fmt.Fprintf(&body, "health_status=%s\n", entry.Health)
		if strings.TrimSpace(entry.FailingJob) != "" {
			fmt.Fprintf(&body, "failing_job=%s\n", entry.FailingJob)
		}
		if strings.TrimSpace(entry.Reason) != "" {
			fmt.Fprintf(&body, "reason=%s\n", entry.Reason)
		}
		if strings.TrimSpace(entry.HealthRaw) != "" {
			fmt.Fprintf(&body, "health_raw=%s\n", entry.HealthRaw)
		}
		if err := os.WriteFile(filepath.Join(runDir, fmt.Sprintf("gpu-%d-status.txt", idx)), []byte(body.String()), 0644); err != nil {
			return err
		}
	}
	return nil
}

// runSATCommandCtx runs cmd and returns its combined output. livePath is
// variadic purely so existing callers are unaffected: pass a path (job's
// output file) to also stream output to disk live as it runs, so a crash
// mid-command leaves whatever had printed so far instead of nothing at all.
func runSATCommandCtx(ctx context.Context, verboseLog, name string, cmd []string, env []string, logFunc func(string), livePath ...string) ([]byte, error) {
	start := time.Now().UTC()
	resolvedCmd, err := resolveSATCommand(cmd)
	appendSATVerboseLog(verboseLog,
		fmt.Sprintf("[%s] start %s", start.Format(time.RFC3339), name),
		"cmd: "+strings.Join(resolvedCmd, " "),
	)
	if logFunc != nil {
		logFunc(fmt.Sprintf("=== %s ===", name))
	}
	if err != nil {
		appendSATVerboseLog(verboseLog,
			fmt.Sprintf("[%s] finish %s", time.Now().UTC().Format(time.RFC3339), name),
			"rc: 1",
			fmt.Sprintf("duration_ms: %d", time.Since(start).Milliseconds()),
			"",
		)
		return []byte(err.Error() + "\n"), err
	}

	c := exec.CommandContext(ctx, resolvedCmd[0], resolvedCmd[1:]...)
	c.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	c.Cancel = func() error {
		if c.Process != nil {
			_ = syscall.Kill(-c.Process.Pid, syscall.SIGKILL)
		}
		return nil
	}
	if len(env) > 0 {
		c.Env = append(os.Environ(), env...)
	}
	var live string
	if len(livePath) > 0 {
		live = livePath[0]
	}
	out, err := streamExecOutput(c, logFunc, live)

	rc := 0
	if err != nil {
		rc = 1
	}
	appendSATVerboseLog(verboseLog,
		fmt.Sprintf("[%s] finish %s", time.Now().UTC().Format(time.RFC3339), name),
		fmt.Sprintf("rc: %d", rc),
		fmt.Sprintf("duration_ms: %d", time.Since(start).Milliseconds()),
		"",
	)
	return out, err
}

// smartctlSelfTestPollInterval/Timeout bound how long we poll the drive after
// launching `smartctl -t short`, which SMART/ATA specs put at ~2 minutes.
const (
	smartctlSelfTestPollInterval = 5 * time.Second
	smartctlSelfTestTimeout      = 4 * time.Minute
)

// waitForSmartctlSelfTest polls `smartctl -a` until the short self-test
// started on devPath finishes (or the timeout/context elapses) and returns
// the final output, which reflects the actual test result rather than the
// "Testing has begun" launch acknowledgement.
func waitForSmartctlSelfTest(ctx context.Context, verboseLog, devPath string, logFunc func(string)) []byte {
	deadline := time.Now().Add(smartctlSelfTestTimeout)
	var last []byte
	for {
		out, _ := runSATCommandCtx(ctx, verboseLog, "smartctl-self-test-status", []string{"smartctl", "-a", devPath}, nil, nil)
		last = out
		if ctx.Err() != nil {
			return last
		}
		lower := bytes.ToLower(out)
		if !bytes.Contains(lower, []byte("self-test routine in progress")) &&
			!bytes.Contains(lower, []byte("% of test remaining")) {
			return last
		}
		if time.Now().After(deadline) {
			return last
		}
		select {
		case <-ctx.Done():
			return last
		case <-time.After(smartctlSelfTestPollInterval):
		}
	}
}

func listStorageDevices() ([]string, error) {
	out, err := satExecCommand("lsblk", "-dn", "-o", "NAME,TYPE,TRAN").Output()
	if err != nil {
		return nil, err
	}
	return parseStorageDevices(string(out)), nil
}

// storageSATCommands returns the commands to run for a single storage device.
// extended=false (Check): read-only SMART/NVMe data collection, no self-test.
// extended=true (Load): data collection + short self-test.
func storageSATCommands(devPath string, extended bool) []satJob {
	if strings.Contains(filepath.Base(devPath), "nvme") {
		jobs := []satJob{
			{name: "nvme-id-ctrl", cmd: []string{"nvme", "id-ctrl", devPath, "-o", "json"}},
			{name: "nvme-smart-log", cmd: []string{"nvme", "smart-log", devPath, "-o", "json"}},
		}
		if extended {
			jobs = append(jobs, satJob{name: "nvme-device-self-test", cmd: []string{"nvme", "device-self-test", devPath, "-s", "1", "--wait"}, syncBracket: true})
		}
		return jobs
	}
	jobs := []satJob{
		{name: "smartctl-health", cmd: []string{"smartctl", "-H", "-A", "-i", devPath}},
	}
	if extended {
		jobs = append(jobs, satJob{name: "smartctl-self-test-short", cmd: []string{"smartctl", "-t", "short", devPath}, syncBracket: true})
	}
	return jobs
}

func (s *satStats) Add(status string) {
	switch status {
	case "OK":
		s.OK++
	case "UNSUPPORTED", "PARTIAL":
		s.Unsupported++
	default:
		s.Failed++
	}
}

func (s satStats) Overall() string {
	if s.Failed > 0 {
		return "FAILED"
	}
	if s.Unsupported > 0 {
		return "PARTIAL"
	}
	return "OK"
}

func writeSATStats(summary *strings.Builder, stats satStats) {
	fmt.Fprintf(summary, "overall_status=%s\n", stats.Overall())
	fmt.Fprintf(summary, "job_ok=%d\n", stats.OK)
	fmt.Fprintf(summary, "job_failed=%d\n", stats.Failed)
	fmt.Fprintf(summary, "job_unsupported=%d\n", stats.Unsupported)
	fmt.Fprintf(summary, "job_informational_failed=%d\n", stats.Informational)
}

func classifySATResult(name string, out []byte, err error) (string, int) {
	rc := 0
	if err != nil {
		rc = 1
	}
	if err == nil {
		return "OK", rc
	}

	text := strings.ToLower(string(out))
	// No output at all means the tool failed to start (mlock limit, binary missing,
	// etc.) — we cannot say anything about hardware health → UNSUPPORTED.
	if len(strings.TrimSpace(text)) == 0 {
		return "UNSUPPORTED", rc
	}
	if strings.Contains(text, "unsupported") ||
		strings.Contains(text, "not supported") ||
		strings.Contains(text, "not found in path") ||
		strings.Contains(text, "invalid opcode") ||
		strings.Contains(text, "unknown command") ||
		strings.Contains(text, "not implemented") ||
		strings.Contains(text, "not available") ||
		strings.Contains(text, "cuda_error_system_not_ready") ||
		strings.Contains(text, "no such device") ||
		// nvidia-smi on a machine with no NVIDIA GPU
		strings.Contains(text, "couldn't communicate with the nvidia driver") ||
		strings.Contains(text, "no nvidia gpu") ||
		// Some NVMe firmwares start self-test but never expose progress to nvme-cli
		// while waiting, so the CLI stops polling without proving device failure.
		(strings.Contains(name, "self-test") &&
			strings.Contains(text, "no progress for") &&
			strings.Contains(text, "stop waiting")) ||
		(strings.Contains(name, "self-test") && strings.Contains(text, "aborted")) {
		return "UNSUPPORTED", rc
	}
	return "FAILED", rc
}

func hasSMARTOverallHealth(out []byte) bool {
	m := smartHealthRE.FindStringSubmatch(string(out))
	return len(m) > 1 && strings.TrimSpace(m[1]) != ""
}

func runROCmSMI(args ...string) ([]byte, error) {
	cmd, err := resolveROCmSMICommand(args...)
	if err != nil {
		return nil, err
	}
	return satExecCommand(cmd[0], cmd[1:]...).CombinedOutput()
}

func resolveSATCommand(cmd []string) ([]string, error) {
	if len(cmd) == 0 {
		return nil, errors.New("empty SAT command")
	}
	switch cmd[0] {
	case "rocm-smi":
		return resolveROCmSMICommand(cmd[1:]...)
	case "rvs":
		return resolveRVSCommand(cmd[1:]...)
	}
	path, err := satLookPath(cmd[0])
	if err != nil {
		return nil, fmt.Errorf("%s not found in PATH: %w", cmd[0], err)
	}
	return append([]string{path}, cmd[1:]...), nil
}

func resolveRVSCommand(args ...string) ([]string, error) {
	if path, err := satLookPath("rvs"); err == nil {
		return append([]string{path}, args...), nil
	}
	for _, path := range expandExistingPaths(rvsExecutableGlobs) {
		return append([]string{path}, args...), nil
	}
	return nil, errors.New("rvs not found in PATH or under /opt/rocm")
}

func resolveROCmSMICommand(args ...string) ([]string, error) {
	if path, err := satLookPath("rocm-smi"); err == nil {
		return append([]string{path}, args...), nil
	}

	for _, path := range rocmSMIExecutableCandidates() {
		return append([]string{path}, args...), nil
	}

	pythonPath, pyErr := satLookPath("python3")
	if pyErr == nil {
		for _, script := range rocmSMIScriptCandidates() {
			cmd := []string{pythonPath, script}
			cmd = append(cmd, args...)
			return cmd, nil
		}
	}

	return nil, errors.New("rocm-smi not found in PATH or under /opt/rocm")
}

func resolveDCGMProfTesterCommand(args ...string) ([]string, error) {
	for _, candidate := range dcgmProfTesterCandidates {
		if path, err := satLookPath(candidate); err == nil {
			return append([]string{path}, args...), nil
		}
	}
	return nil, errors.New("dcgmproftester not found in PATH")
}

func ensureAMDRuntimeReady() error {
	if _, err := os.Stat("/dev/kfd"); err == nil {
		return nil
	}
	if raw, err := os.ReadFile("/sys/module/amdgpu/initstate"); err == nil {
		state := strings.TrimSpace(string(raw))
		if strings.EqualFold(state, "live") {
			return nil
		}
		return fmt.Errorf("AMD driver is present but not initialized: amdgpu initstate=%q", state)
	}
	return errors.New("AMD GPUs are present but the runtime is not initialized: /dev/kfd is missing and amdgpu is not loaded")
}

func rocmSMIExecutableCandidates() []string {
	return expandExistingPaths(rocmSMIExecutableGlobs)
}

func rocmSMIScriptCandidates() []string {
	return expandExistingPaths(rocmSMIScriptGlobs)
}

func expandExistingPaths(patterns []string) []string {
	seen := make(map[string]struct{})
	var paths []string
	for _, pattern := range patterns {
		matches, err := satGlob(pattern)
		if err != nil {
			continue
		}
		sort.Strings(matches)
		for _, match := range matches {
			if _, err := satStat(match); err != nil {
				continue
			}
			if _, ok := seen[match]; ok {
				continue
			}
			seen[match] = struct{}{}
			paths = append(paths, match)
		}
	}
	return paths
}

func parseStorageDevices(raw string) []string {
	var devices []string
	for _, line := range strings.Split(strings.TrimSpace(raw), "\n") {
		fields := strings.Fields(strings.TrimSpace(line))
		if len(fields) < 2 || fields[1] != "disk" {
			continue
		}
		if len(fields) >= 3 && strings.EqualFold(fields[2], "usb") {
			continue
		}
		devices = append(devices, "/dev/"+fields[0])
	}
	return devices
}

// runSATCommandWithMetrics runs a command while collecting GPU metrics in the background.
// On completion it writes gpu-metrics.csv and gpu-metrics.html into runDir.
func runSATCommandWithMetrics(ctx context.Context, verboseLog, name string, cmd []string, env []string, gpuIndices []int, runDir string, logFunc func(string)) ([]byte, error) {
	stopCh := make(chan struct{})
	doneCh := make(chan struct{})
	var metricRows []GPUMetricRow
	start := time.Now()

	go func() {
		defer close(doneCh)
		ticker := time.NewTicker(time.Second)
		defer ticker.Stop()
		for {
			select {
			case <-stopCh:
				return
			case <-ticker.C:
				samples, err := sampleGPUMetrics(gpuIndices)
				if err != nil {
					continue
				}
				elapsed := time.Since(start).Seconds()
				for i := range samples {
					samples[i].ElapsedSec = elapsed
				}
				metricRows = append(metricRows, samples...)
			}
		}
	}()

	out, err := runSATCommandCtx(ctx, verboseLog, name, cmd, env, logFunc, filepath.Join(runDir, name))

	close(stopCh)
	<-doneCh

	if len(metricRows) > 0 {
		_ = WriteGPUMetricsCSV(filepath.Join(runDir, "gpu-metrics.csv"), metricRows)
		_ = WriteGPUMetricsHTML(filepath.Join(runDir, "gpu-metrics.html"), metricRows)
	}

	return out, err
}

func appendSATVerboseLog(path string, lines ...string) {
	if path == "" {
		return
	}
	f, err := os.OpenFile(path, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644)
	if err != nil {
		return
	}
	defer f.Close()
	for _, line := range lines {
		_, _ = io.WriteString(f, line+"\n")
	}
}

func envInt(name string, fallback int) int {
	raw := strings.TrimSpace(os.Getenv(name))
	if raw == "" {
		return fallback
	}
	value, err := strconv.Atoi(raw)
	if err != nil || value <= 0 {
		return fallback
	}
	return value
}