osmedeus/pkg/server/handlers/distributed.go

package handlers

import (
	"github.com/gofiber/fiber/v2"
	"github.com/j3ssie/osmedeus/v5/internal/distributed"
)

// ListWorkers returns all registered workers with their current status
// @Summary List all workers
// @Description Get a list of all registered workers in the distributed pool
// @Tags Distributed
// @Produce json
// @Success 200 {object} map[string]interface{} "List of workers"
// @Failure 500 {object} map[string]interface{} "Failed to list workers"
// @Security BearerAuth
// @Router /osm/api/workers [get]
func ListWorkers(master *distributed.Master) fiber.Handler {
	return func(c *fiber.Ctx) error {
		ctx := c.Context()
		workers, err := master.ListWorkers(ctx)
		if err != nil {
			return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{
				"error":   true,
				"message": err.Error(),
			})
		}

		data := make([]fiber.Map, 0, len(workers))
		for _, w := range workers {
			data = append(data, fiber.Map{
				"id":             w.ID,
				"hostname":       w.Hostname,
				"status":         w.Status,
				"current_task":   w.CurrentTaskID,
				"joined_at":      w.JoinedAt,
				"last_heartbeat": w.LastHeartbeat,
				"tasks_complete": w.TasksComplete,
				"tasks_failed":   w.TasksFailed,
			})
		}

		return c.JSON(fiber.Map{
			"data":  data,
			"count": len(data),
		})
	}
}

// GetWorker returns details for a specific worker
// @Summary Get worker details
// @Description Get details for a specific worker by ID
// @Tags Distributed
// @Produce json
// @Param id path string true "Worker ID"
// @Success 200 {object} map[string]interface{} "Worker details"
// @Failure 404 {object} map[string]interface{} "Worker not found"
// @Failure 500 {object} map[string]interface{} "Failed to get worker"
// @Security BearerAuth
// @Router /osm/api/workers/{id} [get]
func GetWorker(master *distributed.Master) fiber.Handler {
	return func(c *fiber.Ctx) error {
		workerID := c.Params("id")
		ctx := c.Context()

		workers, err := master.ListWorkers(ctx)
		if err != nil {
			return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{
				"error":   true,
				"message": err.Error(),
			})
		}

		for _, w := range workers {
			if w.ID == workerID {
				return c.JSON(fiber.Map{
					"id":             w.ID,
					"hostname":       w.Hostname,
					"status":         w.Status,
					"current_task":   w.CurrentTaskID,
					"joined_at":      w.JoinedAt,
					"last_heartbeat": w.LastHeartbeat,
					"tasks_complete": w.TasksComplete,
					"tasks_failed":   w.TasksFailed,
				})
			}
		}

		return c.Status(fiber.StatusNotFound).JSON(fiber.Map{
			"error":   true,
			"message": "Worker not found",
		})
	}
}

// ListTasks returns all tasks (running and completed)
// @Summary List all tasks
// @Description Get a list of all running and completed tasks
// @Tags Distributed
// @Produce json
// @Success 200 {object} map[string]interface{} "List of running and completed tasks"
// @Failure 500 {object} map[string]interface{} "Failed to list tasks"
// @Security BearerAuth
// @Router /osm/api/tasks [get]
func ListTasks(master *distributed.Master) fiber.Handler {
	return func(c *fiber.Ctx) error {
		ctx := c.Context()
		running, completed, err := master.ListTasks(ctx)
		if err != nil {
			return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{
				"error":   true,
				"message": err.Error(),
			})
		}

		runningData := make([]fiber.Map, 0, len(running))
		for _, t := range running {
			runningData = append(runningData, taskToMap(t))
		}

		completedData := make([]fiber.Map, 0)
		for _, r := range completed {
			completedData = append(completedData, taskResultToMap(r))
		}

		return c.JSON(fiber.Map{
			"running":   runningData,
			"completed": completedData,
		})
	}
}

// GetTask returns details for a specific task
// @Summary Get task details
// @Description Get details for a specific task by ID
// @Tags Distributed
// @Produce json
// @Param id path string true "Task ID"
// @Success 200 {object} map[string]interface{} "Task details"
// @Failure 404 {object} map[string]interface{} "Task not found"
// @Security BearerAuth
// @Router /osm/api/tasks/{id} [get]
func GetTask(master *distributed.Master) fiber.Handler {
	return func(c *fiber.Ctx) error {
		taskID := c.Params("id")
		ctx := c.Context()

		task, result, err := master.GetTaskStatus(ctx, taskID)
		if err != nil {
			return c.Status(fiber.StatusNotFound).JSON(fiber.Map{
				"error":   true,
				"message": err.Error(),
			})
		}

		if task != nil {
			return c.JSON(taskToMap(task))
		}

		if result != nil {
			return c.JSON(taskResultToMap(result))
		}

		return c.Status(fiber.StatusNotFound).JSON(fiber.Map{
			"error":   true,
			"message": "Task not found",
		})
	}
}

// SubmitTaskRequest represents a task submission request
type SubmitTaskRequest struct {
	WorkflowName string                 `json:"workflow_name"`
	WorkflowKind string                 `json:"workflow_kind"`
	Target       string                 `json:"target"`
	Params       map[string]interface{} `json:"params"`
}

// SubmitTask submits a new task to the distributed queue
// @Summary Submit a new task
// @Description Submit a new task to the distributed worker queue
// @Tags Distributed
// @Accept json
// @Produce json
// @Param task body SubmitTaskRequest true "Task configuration"
// @Success 202 {object} map[string]interface{} "Task submitted"
// @Failure 400 {object} map[string]interface{} "Invalid request"
// @Failure 500 {object} map[string]interface{} "Failed to submit task"
// @Security BearerAuth
// @Router /osm/api/tasks [post]
func SubmitTask(master *distributed.Master) fiber.Handler {
	return func(c *fiber.Ctx) error {
		var req SubmitTaskRequest
		if err := c.BodyParser(&req); err != nil {
			return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{
				"error":   true,
				"message": "Invalid request body",
			})
		}

		if req.WorkflowName == "" {
			return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{
				"error":   true,
				"message": "workflow_name is required",
			})
		}

		if req.Target == "" {
			return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{
				"error":   true,
				"message": "target is required",
			})
		}

		task := &distributed.Task{
			WorkflowName: req.WorkflowName,
			WorkflowKind: req.WorkflowKind,
			Target:       req.Target,
			Params:       req.Params,
		}

		ctx := c.Context()
		if err := master.SubmitTask(ctx, task); err != nil {
			return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{
				"error":   true,
				"message": err.Error(),
			})
		}

		return c.Status(fiber.StatusAccepted).JSON(fiber.Map{
			"message": "Task submitted",
			"task_id": task.ID,
		})
	}
}

// taskToMap converts a Task to a fiber.Map
func taskToMap(t *distributed.Task) fiber.Map {
	m := fiber.Map{
		"id":            t.ID,
		"scan_id":       t.ScanID,
		"workflow_name": t.WorkflowName,
		"workflow_kind": t.WorkflowKind,
		"target":        t.Target,
		"status":        t.Status,
		"worker_id":     t.WorkerID,
		"created_at":    t.CreatedAt,
	}
	if t.StartedAt != nil {
		m["started_at"] = t.StartedAt
	}
	return m
}

// taskResultToMap converts a TaskResult to a fiber.Map
func taskResultToMap(r *distributed.TaskResult) fiber.Map {
	return fiber.Map{
		"task_id":      r.TaskID,
		"status":       r.Status,
		"output":       r.Output,
		"error":        r.Error,
		"exports":      r.Exports,
		"completed_at": r.CompletedAt,
	}
}