bee/audit/internal/webui/tasks_http_persistence.go

package webui

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"

	"bee/audit/internal/platform"
)

func (h *handler) handleAPITasksList(w http.ResponseWriter, _ *http.Request) {
	tasks := globalQueue.snapshot()
	writeJSON(w, tasks)
}

func (h *handler) handleAPITasksCancel(w http.ResponseWriter, r *http.Request) {
	id := r.PathValue("id")
	t, ok := globalQueue.findByID(id)
	if !ok {
		writeError(w, http.StatusNotFound, "task not found")
		return
	}
	globalQueue.mu.Lock()
	defer globalQueue.mu.Unlock()
	switch t.Status {
	case TaskPending:
		t.Status = TaskCancelled
		now := time.Now()
		t.DoneAt = &now
		globalQueue.persistLocked()
		taskSerialEvent(t, "finished with status="+t.Status)
		writeJSON(w, map[string]string{"status": "cancelled"})
	case TaskRunning:
		if t.job == nil || !t.job.abort() {
			writeError(w, http.StatusConflict, "task is not cancellable")
			return
		}
		writeJSON(w, map[string]string{"status": "aborting"})
	default:
		writeError(w, http.StatusConflict, "task is not running or pending")
	}
}

func (h *handler) handleAPITasksPriority(w http.ResponseWriter, r *http.Request) {
	id := r.PathValue("id")
	t, ok := globalQueue.findByID(id)
	if !ok {
		writeError(w, http.StatusNotFound, "task not found")
		return
	}
	var req struct {
		Delta int `json:"delta"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeError(w, http.StatusBadRequest, "invalid body")
		return
	}
	globalQueue.mu.Lock()
	defer globalQueue.mu.Unlock()
	if t.Status != TaskPending {
		writeError(w, http.StatusConflict, "only pending tasks can be reprioritised")
		return
	}
	t.Priority += req.Delta
	globalQueue.persistLocked()
	writeJSON(w, map[string]int{"priority": t.Priority})
}

func (h *handler) handleAPITasksCancelAll(w http.ResponseWriter, _ *http.Request) {
	globalQueue.mu.Lock()
	now := time.Now()
	n := 0
	for _, t := range globalQueue.tasks {
		switch t.Status {
		case TaskPending:
			t.Status = TaskCancelled
			t.DoneAt = &now
			taskSerialEvent(t, "finished with status="+t.Status)
			n++
		case TaskRunning:
			if t.job != nil {
				t.job.abort()
			}
			n++
		}
	}
	globalQueue.persistLocked()
	globalQueue.mu.Unlock()
	writeJSON(w, map[string]int{"cancelled": n})
}

func (h *handler) handleAPITasksKillWorkers(w http.ResponseWriter, _ *http.Request) {
	// Cancel all queued/running tasks in the queue first.
	globalQueue.mu.Lock()
	now := time.Now()
	cancelled := 0
	for _, t := range globalQueue.tasks {
		switch t.Status {
		case TaskPending:
			t.Status = TaskCancelled
			t.DoneAt = &now
			taskSerialEvent(t, "finished with status="+t.Status)
			cancelled++
		case TaskRunning:
			if t.job != nil {
				t.job.abort()
			}
			if taskMayLeaveOrphanWorkers(t.Target) {
				platform.KillTestWorkers()
			}
			t.Status = TaskCancelled
			t.DoneAt = &now
			taskSerialEvent(t, "finished with status="+t.Status)
			cancelled++
		}
	}
	globalQueue.persistLocked()
	globalQueue.mu.Unlock()

	// Kill orphaned test worker processes at the OS level.
	killed := platform.KillTestWorkers()
	writeJSON(w, map[string]any{
		"cancelled": cancelled,
		"killed":    len(killed),
		"processes": killed,
	})
}

func (h *handler) handleAPITasksStream(w http.ResponseWriter, r *http.Request) {
	id := r.PathValue("id")
	src, ok := globalQueue.taskStreamSource(id)
	if !ok {
		http.Error(w, "task not found", http.StatusNotFound)
		return
	}
	if src.job != nil {
		streamJob(w, r, src.job)
		return
	}
	if src.status == TaskDone || src.status == TaskFailed || src.status == TaskCancelled {
		j := newTaskJobState(src.logPath)
		j.finish(src.errMsg)
		streamJob(w, r, j)
		return
	}
	if !sseStart(w) {
		return
	}
	sseWrite(w, "", "Task is queued. Waiting for worker...")
	ticker := time.NewTicker(200 * time.Millisecond)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			src, ok = globalQueue.taskStreamSource(id)
			if !ok {
				sseWrite(w, "done", "task not found")
				return
			}
			if src.job != nil {
				streamSubscribedJob(w, r, src.job)
				return
			}
			if src.status == TaskDone || src.status == TaskFailed || src.status == TaskCancelled {
				j := newTaskJobState(src.logPath)
				j.finish(src.errMsg)
				streamSubscribedJob(w, r, j)
				return
			}
		case <-r.Context().Done():
			return
		}
	}
}

func (q *taskQueue) assignTaskLogPathLocked(t *Task) {
	if q.logsDir == "" || t.ID == "" {
		return
	}
	q.ensureTaskArtifactPathsLocked(t)
}

func (q *taskQueue) loadLocked() {
	if q.statePath == "" {
		return
	}
	data, err := os.ReadFile(q.statePath)
	if err != nil || len(data) == 0 {
		return
	}
	var persisted []persistedTask
	if err := json.Unmarshal(data, &persisted); err != nil {
		return
	}
	for _, pt := range persisted {
		t := &Task{
			ID:             pt.ID,
			Name:           pt.Name,
			Target:         pt.Target,
			Priority:       pt.Priority,
			Status:         pt.Status,
			CreatedAt:      pt.CreatedAt,
			StartedAt:      pt.StartedAt,
			DoneAt:         pt.DoneAt,
			ErrMsg:         pt.ErrMsg,
			LogPath:        pt.LogPath,
			ArtifactsDir:   pt.ArtifactsDir,
			ReportJSONPath: pt.ReportJSONPath,
			ReportHTMLPath: pt.ReportHTMLPath,
			params:         pt.Params,
		}
		q.assignTaskLogPathLocked(t)
		if t.Status == TaskRunning {
			state, ok := readTaskRunnerState(t)
			switch {
			case ok && state.Status == TaskRunning && processAlive(state.PID):
				t.runnerPID = state.PID
				t.job = newTaskJobState(t.LogPath)
			case ok && state.Status != TaskRunning:
				t.runnerPID = state.PID
				t.Status = state.Status
				t.ErrMsg = state.Error
				now := state.UpdatedAt
				if now.IsZero() {
					now = time.Now()
				}
				t.DoneAt = &now
			default:
				if taskMayLeaveOrphanWorkers(t.Target) {
					_ = platform.KillTestWorkers()
				}
				now := time.Now()
				t.Status = TaskFailed
				t.DoneAt = &now
				t.ErrMsg = "interrupted by bee-web restart"
			}
		} else if t.Status == TaskPending {
			t.StartedAt = nil
			t.DoneAt = nil
			t.ErrMsg = ""
		}
		q.tasks = append(q.tasks, t)
	}
	q.prune()
	q.persistLocked()
}

func (q *taskQueue) persistLocked() {
	if q.statePath == "" {
		return
	}
	state := make([]persistedTask, 0, len(q.tasks))
	for _, t := range q.tasks {
		state = append(state, persistedTask{
			ID:             t.ID,
			Name:           t.Name,
			Target:         t.Target,
			Priority:       t.Priority,
			Status:         t.Status,
			CreatedAt:      t.CreatedAt,
			StartedAt:      t.StartedAt,
			DoneAt:         t.DoneAt,
			ErrMsg:         t.ErrMsg,
			LogPath:        t.LogPath,
			ArtifactsDir:   t.ArtifactsDir,
			ReportJSONPath: t.ReportJSONPath,
			ReportHTMLPath: t.ReportHTMLPath,
			Params:         t.params,
		})
	}
	data, err := json.MarshalIndent(state, "", "  ")
	if err != nil {
		return
	}
	tmp := q.statePath + ".tmp"
	if err := os.WriteFile(tmp, data, 0644); err != nil {
		return
	}
	_ = os.Rename(tmp, q.statePath)
}

func taskElapsedSec(t *Task, now time.Time) int {
	if t == nil || t.StartedAt == nil || t.StartedAt.IsZero() {
		return 0
	}
	start := *t.StartedAt
	if !t.CreatedAt.IsZero() && start.Before(t.CreatedAt) {
		start = t.CreatedAt
	}
	end := now
	if t.DoneAt != nil && !t.DoneAt.IsZero() {
		end = *t.DoneAt
	}
	if end.Before(start) {
		return 0
	}
	return int(end.Sub(start).Round(time.Second) / time.Second)
}

func taskFolderStatus(status string) string {
	status = strings.TrimSpace(strings.ToLower(status))
	switch status {
	case TaskRunning, TaskDone, TaskFailed, TaskCancelled:
		return status
	default:
		return TaskPending
	}
}

func sanitizeTaskFolderPart(s string) string {
	s = strings.TrimSpace(strings.ToLower(s))
	if s == "" {
		return "task"
	}
	var b strings.Builder
	lastDash := false
	for _, r := range s {
		isAlnum := (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9')
		if isAlnum {
			b.WriteRune(r)
			lastDash = false
			continue
		}
		if !lastDash {
			b.WriteByte('-')
			lastDash = true
		}
	}
	out := strings.Trim(b.String(), "-")
	if out == "" {
		return "task"
	}
	return out
}

func taskArtifactsDir(root string, t *Task, status string) string {
	if strings.TrimSpace(root) == "" || t == nil {
		return ""
	}
	prefix := taskFolderNumberPrefix(t.ID)
	return filepath.Join(root, fmt.Sprintf("%s_%s_%s", prefix, sanitizeTaskFolderPart(t.Name), taskFolderStatus(status)))
}

func taskFolderNumberPrefix(taskID string) string {
	taskID = strings.TrimSpace(taskID)
	if strings.HasPrefix(taskID, "TASK-") && len(taskID) >= len("TASK-000") {
		num := strings.TrimSpace(strings.TrimPrefix(taskID, "TASK-"))
		if len(num) == 3 {
			allDigits := true
			for _, r := range num {
				if r < '0' || r > '9' {
					allDigits = false
					break
				}
			}
			if allDigits {
				return num
			}
		}
	}
	fallback := sanitizeTaskFolderPart(taskID)
	if fallback == "" {
		return "000"
	}
	return fallback
}

func ensureTaskReportPaths(t *Task) {
	if t == nil || strings.TrimSpace(t.ArtifactsDir) == "" {
		return
	}
	if t.LogPath == "" || filepath.Base(t.LogPath) == "task.log" {
		t.LogPath = filepath.Join(t.ArtifactsDir, "task.log")
	}
	t.ReportJSONPath = filepath.Join(t.ArtifactsDir, "report.json")
	t.ReportHTMLPath = filepath.Join(t.ArtifactsDir, "report.html")
}

func (q *taskQueue) ensureTaskArtifactPathsLocked(t *Task) {
	if t == nil || strings.TrimSpace(q.logsDir) == "" || strings.TrimSpace(t.ID) == "" {
		return
	}
	if strings.TrimSpace(t.ArtifactsDir) == "" {
		t.ArtifactsDir = taskArtifactsDir(q.logsDir, t, t.Status)
	}
	if t.ArtifactsDir != "" {
		_ = os.MkdirAll(t.ArtifactsDir, 0755)
	}
	ensureTaskReportPaths(t)
}

func (q *taskQueue) finalizeTaskArtifactPathsLocked(t *Task) {
	if t == nil || strings.TrimSpace(q.logsDir) == "" || strings.TrimSpace(t.ID) == "" {
		return
	}
	q.ensureTaskArtifactPathsLocked(t)
	dstDir := taskArtifactsDir(q.logsDir, t, t.Status)
	if dstDir == "" {
		return
	}
	if t.ArtifactsDir != "" && t.ArtifactsDir != dstDir {
		if _, err := os.Stat(dstDir); err != nil {
			_ = os.Rename(t.ArtifactsDir, dstDir)
		}
		t.ArtifactsDir = dstDir
	}
	ensureTaskReportPaths(t)
}