bee/audit/internal/webui/api.go

package webui

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os/exec"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"sync/atomic"
	"syscall"
	"time"

	"bee/audit/internal/app"
	"bee/audit/internal/platform"
)

var ansiEscapeRE = regexp.MustCompile(`\x1b\[[0-9;]*[a-zA-Z]|\x1b[()][A-Z0-9]|\x1b[DABC]`)
var apiListNvidiaGPUs = func(a *app.App) ([]platform.NvidiaGPU, error) {
	if a == nil {
		return nil, fmt.Errorf("app not configured")
	}
	return a.ListNvidiaGPUs()
}
var apiListNvidiaGPUStatuses = func(a *app.App) ([]platform.NvidiaGPUStatus, error) {
	if a == nil {
		return nil, fmt.Errorf("app not configured")
	}
	return a.ListNvidiaGPUStatuses()
}

const (
	taskPriorityBenchmark      = 10
	taskPriorityBurn           = 20
	taskPriorityValidateStress = 30
	taskPriorityValidate       = 40
	taskPriorityAudit          = 50
	taskPriorityInstallToRAM   = 60
	taskPriorityInstall        = 70
)

// ── Job ID counter ────────────────────────────────────────────────────────────

var jobCounter atomic.Uint64

func newJobID(_ string) string {
	start := int((jobCounter.Add(1) - 1) % 1000)
	globalQueue.mu.Lock()
	defer globalQueue.mu.Unlock()
	for offset := 0; offset < 1000; offset++ {
		n := (start + offset) % 1000
		id := fmt.Sprintf("TASK-%03d", n)
		if !taskIDInUseLocked(id) {
			return id
		}
	}
	return fmt.Sprintf("TASK-%03d", start)
}

func taskIDInUseLocked(id string) bool {
	for _, t := range globalQueue.tasks {
		if t != nil && t.ID == id {
			return true
		}
	}
	return false
}

type taskRunResponse struct {
	TaskID    string   `json:"task_id,omitempty"`
	JobID     string   `json:"job_id,omitempty"`
	TaskIDs   []string `json:"task_ids,omitempty"`
	JobIDs    []string `json:"job_ids,omitempty"`
	TaskCount int      `json:"task_count,omitempty"`
}

type nvidiaTaskSelection struct {
	GPUIndices []int
	Label      string
}

func writeTaskRunResponse(w http.ResponseWriter, tasks []*Task) {
	if len(tasks) == 0 {
		writeJSON(w, taskRunResponse{})
		return
	}
	ids := make([]string, 0, len(tasks))
	for _, t := range tasks {
		if t == nil || strings.TrimSpace(t.ID) == "" {
			continue
		}
		ids = append(ids, t.ID)
	}
	resp := taskRunResponse{TaskCount: len(ids)}
	if len(ids) > 0 {
		resp.TaskID = ids[0]
		resp.JobID = ids[0]
		resp.TaskIDs = ids
		resp.JobIDs = ids
	}
	writeJSON(w, resp)
}

func shouldSplitHomogeneousNvidiaTarget(target string) bool {
	switch strings.TrimSpace(target) {
	case "nvidia", "nvidia-targeted-stress", "nvidia-bench-perf", "nvidia-bench-power", "nvidia-compute",
		"nvidia-targeted-power", "nvidia-pulse", "nvidia-interconnect",
		"nvidia-bandwidth", "nvidia-stress":
		return true
	default:
		return false
	}
}

func defaultTaskPriority(target string, params taskParams) int {
	switch strings.TrimSpace(target) {
	case "install":
		return taskPriorityInstall
	case "install-to-ram":
		return taskPriorityInstallToRAM
	case "nvme-format":
		return taskPriorityInstall
	case "audit":
		return taskPriorityAudit
	case "nvidia-bench-perf", "nvidia-bench-power", "nvidia-bench-autotune":
		return taskPriorityBenchmark
	case "nvidia-stress", "amd-stress", "memory-stress", "sat-stress", "platform-stress", "nvidia-compute", "scenario":
		return taskPriorityBurn
	case "nvidia", "nvidia-targeted-stress", "nvidia-targeted-power", "nvidia-pulse",
		"nvidia-interconnect", "nvidia-bandwidth", "memory", "storage", "cpu",
		"amd", "amd-mem", "amd-bandwidth", "nvidia-config", "pcie-link", "nvidia-pcie-bandwidth":
		if params.StressMode {
			return taskPriorityValidateStress
		}
		return taskPriorityValidate
	default:
		return 0
	}
}

func expandHomogeneousNvidiaSelections(gpus []platform.NvidiaGPU, include, exclude []int) ([]nvidiaTaskSelection, error) {
	if len(gpus) == 0 {
		return nil, fmt.Errorf("no NVIDIA GPUs detected")
	}
	indexed := make(map[int]platform.NvidiaGPU, len(gpus))
	allIndices := make([]int, 0, len(gpus))
	for _, gpu := range gpus {
		indexed[gpu.Index] = gpu
		allIndices = append(allIndices, gpu.Index)
	}
	sort.Ints(allIndices)

	selected := allIndices
	if len(include) > 0 {
		selected = make([]int, 0, len(include))
		seen := make(map[int]struct{}, len(include))
		for _, idx := range include {
			if _, ok := indexed[idx]; !ok {
				continue
			}
			if _, dup := seen[idx]; dup {
				continue
			}
			seen[idx] = struct{}{}
			selected = append(selected, idx)
		}
		sort.Ints(selected)
	}
	if len(exclude) > 0 {
		skip := make(map[int]struct{}, len(exclude))
		for _, idx := range exclude {
			skip[idx] = struct{}{}
		}
		filtered := selected[:0]
		for _, idx := range selected {
			if _, ok := skip[idx]; ok {
				continue
			}
			filtered = append(filtered, idx)
		}
		selected = filtered
	}
	if len(selected) == 0 {
		return nil, fmt.Errorf("no NVIDIA GPUs selected")
	}

	modelGroups := make(map[string][]platform.NvidiaGPU)
	modelOrder := make([]string, 0)
	for _, idx := range selected {
		gpu := indexed[idx]
		model := strings.TrimSpace(gpu.Name)
		if model == "" {
			model = fmt.Sprintf("GPU %d", gpu.Index)
		}
		if _, ok := modelGroups[model]; !ok {
			modelOrder = append(modelOrder, model)
		}
		modelGroups[model] = append(modelGroups[model], gpu)
	}
	sort.Slice(modelOrder, func(i, j int) bool {
		left := modelGroups[modelOrder[i]]
		right := modelGroups[modelOrder[j]]
		if len(left) == 0 || len(right) == 0 {
			return modelOrder[i] < modelOrder[j]
		}
		return left[0].Index < right[0].Index
	})

	var groups []nvidiaTaskSelection
	var singles []nvidiaTaskSelection
	for _, model := range modelOrder {
		group := modelGroups[model]
		sort.Slice(group, func(i, j int) bool { return group[i].Index < group[j].Index })
		indices := make([]int, 0, len(group))
		for _, gpu := range group {
			indices = append(indices, gpu.Index)
		}
		if len(indices) >= 2 {
			groups = append(groups, nvidiaTaskSelection{
				GPUIndices: indices,
				Label:      fmt.Sprintf("%s; GPUs %s", model, joinTaskIndices(indices)),
			})
			continue
		}
		gpu := group[0]
		singles = append(singles, nvidiaTaskSelection{
			GPUIndices: []int{gpu.Index},
			Label:      fmt.Sprintf("GPU %d — %s", gpu.Index, model),
		})
	}
	return append(groups, singles...), nil
}

func joinTaskIndices(indices []int) string {
	parts := make([]string, 0, len(indices))
	for _, idx := range indices {
		parts = append(parts, fmt.Sprintf("%d", idx))
	}
	return strings.Join(parts, ",")
}

func formatGPUIndexList(indices []int) string {
	parts := make([]string, len(indices))
	for i, idx := range indices {
		parts[i] = strconv.Itoa(idx)
	}
	return strings.Join(parts, ",")
}

func formatSplitTaskName(baseName, selectionLabel string) string {
	baseName = strings.TrimSpace(baseName)
	selectionLabel = strings.TrimSpace(selectionLabel)
	if baseName == "" {
		return selectionLabel
	}
	if selectionLabel == "" {
		return baseName
	}
	return baseName + " (" + selectionLabel + ")"
}

func buildNvidiaTaskSet(target string, priority int, createdAt time.Time, params taskParams, baseName string, appRef *app.App, idPrefix string) ([]*Task, error) {
	if !shouldSplitHomogeneousNvidiaTarget(target) || params.ParallelGPUs || params.StaggerGPUStart {
		// Parallel mode, ramp-up mode (or non-splittable target): one task for all selected GPUs.
		if (params.ParallelGPUs || params.StaggerGPUStart) && shouldSplitHomogeneousNvidiaTarget(target) {
			// Resolve the selected GPU indices so ExcludeGPUIndices is applied.
			gpus, err := apiListNvidiaGPUs(appRef)
			if err != nil {
				return nil, err
			}
			resolved, err := expandSelectedGPUIndices(gpus, params.GPUIndices, params.ExcludeGPUIndices)
			if err != nil {
				return nil, err
			}
			params.GPUIndices = resolved
			params.ExcludeGPUIndices = nil
		}
		t := &Task{
			ID:        newJobID(idPrefix),
			Name:      baseName,
			Target:    target,
			Priority:  priority,
			Status:    TaskPending,
			CreatedAt: createdAt,
			params:    params,
		}
		return []*Task{t}, nil
	}
	gpus, err := apiListNvidiaGPUs(appRef)
	if err != nil {
		return nil, err
	}
	selections, err := expandHomogeneousNvidiaSelections(gpus, params.GPUIndices, params.ExcludeGPUIndices)
	if err != nil {
		return nil, err
	}
	tasks := make([]*Task, 0, len(selections))
	for _, selection := range selections {
		taskParamsCopy := params
		taskParamsCopy.GPUIndices = append([]int(nil), selection.GPUIndices...)
		taskParamsCopy.ExcludeGPUIndices = nil
		displayName := formatSplitTaskName(baseName, selection.Label)
		taskParamsCopy.DisplayName = displayName
		tasks = append(tasks, &Task{
			ID:        newJobID(idPrefix),
			Name:      displayName,
			Target:    target,
			Priority:  priority,
			Status:    TaskPending,
			CreatedAt: createdAt,
			params:    taskParamsCopy,
		})
	}
	return tasks, nil
}

// expandSelectedGPUIndices returns the sorted list of selected GPU indices after
// applying include/exclude filters, without splitting by model.
func expandSelectedGPUIndices(gpus []platform.NvidiaGPU, include, exclude []int) ([]int, error) {
	indexed := make(map[int]struct{}, len(gpus))
	allIndices := make([]int, 0, len(gpus))
	for _, gpu := range gpus {
		indexed[gpu.Index] = struct{}{}
		allIndices = append(allIndices, gpu.Index)
	}
	sort.Ints(allIndices)

	selected := allIndices
	if len(include) > 0 {
		selected = make([]int, 0, len(include))
		seen := make(map[int]struct{}, len(include))
		for _, idx := range include {
			if _, ok := indexed[idx]; !ok {
				continue
			}
			if _, dup := seen[idx]; dup {
				continue
			}
			seen[idx] = struct{}{}
			selected = append(selected, idx)
		}
		sort.Ints(selected)
	}
	if len(exclude) > 0 {
		skip := make(map[int]struct{}, len(exclude))
		for _, idx := range exclude {
			skip[idx] = struct{}{}
		}
		filtered := selected[:0]
		for _, idx := range selected {
			if _, ok := skip[idx]; ok {
				continue
			}
			filtered = append(filtered, idx)
		}
		selected = filtered
	}
	if len(selected) == 0 {
		return nil, fmt.Errorf("no NVIDIA GPUs selected")
	}
	return selected, nil
}

// ── SSE helpers ───────────────────────────────────────────────────────────────

func sseWrite(w http.ResponseWriter, event, data string) bool {
	f, ok := w.(http.Flusher)
	if !ok {
		return false
	}
	if event != "" {
		fmt.Fprintf(w, "event: %s\n", event)
	}
	fmt.Fprintf(w, "data: %s\n\n", data)
	f.Flush()
	return true
}

func sseStart(w http.ResponseWriter) bool {
	_, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "streaming not supported", http.StatusInternalServerError)
		return false
	}
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("Access-Control-Allow-Origin", "*")
	return true
}

// streamJob streams lines from a jobState to a SSE response.
func streamJob(w http.ResponseWriter, r *http.Request, j *jobState) {
	if !sseStart(w) {
		return
	}
	streamSubscribedJob(w, r, j)
}

func streamSubscribedJob(w http.ResponseWriter, r *http.Request, j *jobState) {
	existing, ch := j.subscribe()
	for _, line := range existing {
		sseWrite(w, "", line)
	}
	if ch == nil {
		// Job already finished
		sseWrite(w, "done", j.err)
		return
	}
	for {
		select {
		case line, ok := <-ch:
			if !ok {
				sseWrite(w, "done", j.err)
				return
			}
			sseWrite(w, "", line)
		case <-r.Context().Done():
			return
		}
	}
}

// streamCmdJob runs an exec.Cmd and streams stdout+stderr lines into j.
func streamCmdJob(j *jobState, cmd *exec.Cmd) error {
	pr, pw := io.Pipe()
	cmd.Stdout = pw
	cmd.Stderr = pw

	if err := cmd.Start(); err != nil {
		_ = pw.Close()
		_ = pr.Close()
		return err
	}
	// Lower the CPU scheduling priority of stress/audit subprocesses to nice+10
	// so the X server and kernel interrupt handling remain responsive under load
	// (prevents KVM/IPMI graphical console from freezing during GPU stress tests).
	if cmd.Process != nil {
		_ = syscall.Setpriority(syscall.PRIO_PROCESS, cmd.Process.Pid, 10)
	}

	scanDone := make(chan error, 1)
	go func() {
		defer func() {
			if rec := recover(); rec != nil {
				scanDone <- fmt.Errorf("stream scanner panic: %v", rec)
			}
		}()
		scanner := bufio.NewScanner(pr)
		scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
		for scanner.Scan() {
			// Split on \r to handle progress-bar style output (e.g. \r overwrites)
			// and strip ANSI escape codes so logs are readable in the browser.
			parts := strings.Split(scanner.Text(), "\r")
			for _, part := range parts {
				line := ansiEscapeRE.ReplaceAllString(part, "")
				if line != "" {
					j.append(line)
				}
			}
		}
		if err := scanner.Err(); err != nil && !errors.Is(err, io.ErrClosedPipe) {
			scanDone <- err
			return
		}
		scanDone <- nil
	}()

	err := cmd.Wait()
	_ = pw.Close()
	scanErr := <-scanDone
	_ = pr.Close()
	if err != nil {
		return err
	}
	return scanErr
}

// ── Audit ─────────────────────────────────────────────────────────────────────