package webui

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"log/slog"
	"net/http"
	"sort"
	"time"

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

// gpuReadyWait bounds how long /api/sat/run-all waits for the GPU driver
// stack to come up before it plans the GPU tests. The per-GPU GSP firmware
// boot on a multi-GPU box lags the device nodes by tens of seconds.
// Overridable from tests.
var (
	gpuReadyWait         = 75 * time.Second
	gpuReadyPollInterval = 3 * time.Second
	apiRuntimeHealthNow  = func(a *app.App) (schema.RuntimeHealth, error) {
		return a.RuntimeHealthNow()
	}
)

type satRunAllRequest struct {
	StressMode bool `json:"stress_mode"`
	// AMDTargets is the operator's AMD check selection (intent). It is still
	// gated on an AMD GPU actually being present.
	AMDTargets []string `json:"amd_targets"`
	// NvidiaGPUIndices optionally narrows the NVIDIA tests to a subset; empty
	// means "every GPU the backend enumerates".
	NvidiaGPUIndices []int `json:"nvidia_gpu_indices"`
}

type satRunAllResponse struct {
	TaskIDs   []string `json:"task_ids"`
	TaskCount int      `json:"task_count"`
	// Notes carries anything the backend decided to skip or override, so the
	// page can show it without reasoning about hardware itself.
	Notes []string `json:"notes,omitempty"`
}

// handleAPISATRunAll plans and enqueues the full validate/check task set
// server-side. Hardware presence and readiness are decided here, never in the
// browser: the page sends only operator intent (stress toggle, AMD checkbox
// selection, an optional GPU subset).
func (h *handler) handleAPISATRunAll(w http.ResponseWriter, r *http.Request) {
	if h.opts.App == nil {
		writeError(w, http.StatusServiceUnavailable, "app not configured")
		return
	}
	var req satRunAllRequest
	if r.Body != nil {
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil && !errors.Is(err, io.EOF) {
			writeError(w, http.StatusBadRequest, "invalid request body")
			return
		}
	}
	specs, notes := h.planSATRunAll(r.Context(), req)

	var ids []string
	for _, spec := range specs {
		tasks, err := h.enqueueSATTarget(spec.target, spec.params)
		if err != nil {
			writeError(w, http.StatusInternalServerError, err.Error())
			return
		}
		for _, t := range tasks {
			if t != nil {
				ids = append(ids, t.ID)
			}
		}
	}
	slog.Info("sat run-all planned", "tasks", len(ids), "stress", req.StressMode, "notes", len(notes))
	writeJSON(w, satRunAllResponse{TaskIDs: ids, TaskCount: len(ids), Notes: notes})
}

type satRunAllSpec struct {
	target string
	params taskParams
}

func (h *handler) planSATRunAll(ctx context.Context, req satRunAllRequest) ([]satRunAllSpec, []string) {
	var specs []satRunAllSpec
	var notes []string
	skip := func(msg string) {
		notes = append(notes, msg)
		slog.Warn("sat run-all: check skipped", "reason", msg)
	}

	cpuDur := 60
	if req.StressMode {
		cpuDur = 1800
	}
	specs = append(specs,
		satRunAllSpec{target: "cpu", params: taskParams{Duration: cpuDur, StressMode: req.StressMode}},
		satRunAllSpec{target: "memory", params: taskParams{StressMode: req.StressMode}},
		satRunAllSpec{target: "storage", params: taskParams{StressMode: req.StressMode}},
		satRunAllSpec{target: "pcie-link", params: taskParams{}},
	)

	if h.opts.App.TPMPresent() {
		specs = append(specs, satRunAllSpec{target: "tpm", params: taskParams{}})
	} else {
		skip("TPM: no TPM device on this host; check skipped")
	}

	gp := h.opts.App.DetectGPUPresence()

	if gp.Nvidia || gp.NvidiaInitializing {
		// nvidia-config only collects inventory and NVLink state; safe to run
		// even while the compute stack is still coming up.
		specs = append(specs, satRunAllSpec{target: "nvidia-config", params: taskParams{}})

		health, gpus, ready := h.waitForNvidiaReady(ctx)
		switch {
		case health.NvidiaGSPMode == "gsp-stuck":
			skip("NVIDIA: GSP firmware init is stuck; reboot with GSP=off. GPU compute/interconnect/bandwidth tests skipped")
		case !ready:
			skip("NVIDIA: nvidia-smi did not enumerate a GPU after " + gpuReadyWait.String() +
				". GPU compute/interconnect/bandwidth tests skipped; see the GPU Config check")
		default:
			indices := make([]int, 0, len(gpus))
			for _, g := range gpus {
				indices = append(indices, g.Index)
			}
			if len(req.NvidiaGPUIndices) > 0 {
				indices = intersectSortedInts(indices, req.NvidiaGPUIndices)
			}
			if len(indices) == 0 {
				skip("NVIDIA: driver ready but no GPU to test (enumeration empty, or the requested subset matched nothing); GPU tests skipped")
				break
			}
			if !health.CUDAReady {
				notes = append(notes, "NVIDIA: CUDA runtime not confirmed ready; GPU tests queued anyway")
			}
			gpuTargets := []string{"nvidia", "nvidia-interconnect", "nvidia-bandwidth", "nvidia-pcie-bandwidth"}
			if req.StressMode {
				// Stress tier adds the targeted dcgmi diag load tests.
				gpuTargets = append(gpuTargets, "nvidia-targeted-stress", "nvidia-targeted-power", "nvidia-pulse")
			}
			for _, target := range gpuTargets {
				specs = append(specs, satRunAllSpec{
					target: target,
					params: taskParams{GPUIndices: append([]int(nil), indices...), StressMode: req.StressMode},
				})
			}
		}
	}

	if gp.AMD {
		for _, target := range req.AMDTargets {
			switch target {
			case "amd", "amd-mem", "amd-bandwidth":
				specs = append(specs, satRunAllSpec{target: target, params: taskParams{StressMode: req.StressMode}})
			}
		}
	}

	return specs, notes
}

// waitForNvidiaReady waits until the same nvidia-smi query used by
// ListNvidiaGPUs returns at least one GPU. A loaded kernel module alone is
// not evidence that NVIDIA user-space commands can address a GPU yet.
func (h *handler) waitForNvidiaReady(ctx context.Context) (schema.RuntimeHealth, []platform.NvidiaGPU, bool) {
	deadline := time.Now().Add(gpuReadyWait)
	var last schema.RuntimeHealth
	for {
		health, err := apiRuntimeHealthNow(h.opts.App)
		if err == nil {
			last = health
			if health.NvidiaGSPMode == "gsp-stuck" {
				return health, nil, false
			}
		}
		gpus, listErr := apiListNvidiaGPUs(h.opts.App)
		if listErr == nil && len(gpus) > 0 {
			return last, gpus, true
		}
		if time.Now().After(deadline) || ctx.Err() != nil {
			return last, nil, false
		}
		select {
		case <-ctx.Done():
			return last, nil, false
		case <-time.After(gpuReadyPollInterval):
		}
	}
}

// intersectSortedInts returns the ascending-sorted values present in both a
// and b.
func intersectSortedInts(a, b []int) []int {
	set := make(map[int]struct{}, len(b))
	for _, v := range b {
		set[v] = struct{}{}
	}
	out := make([]int, 0, len(a))
	for _, v := range a {
		if _, ok := set[v]; ok {
			out = append(out, v)
		}
	}
	sort.Ints(out)
	return out
}