fetch_ml/internal/api/jupyter/handlers.go

// Package jupyter provides WebSocket handlers for Jupyter-related operations
package jupyter

import (
	"encoding/binary"
	"encoding/json"
	"net/http"
	"time"

	"github.com/gorilla/websocket"
	"github.com/jfraeys/fetch_ml/internal/api/errors"
	"github.com/jfraeys/fetch_ml/internal/auth"
	"github.com/jfraeys/fetch_ml/internal/container"
	"github.com/jfraeys/fetch_ml/internal/jupyter"
	"github.com/jfraeys/fetch_ml/internal/logging"
)

// Handler provides Jupyter-related WebSocket handlers
type Handler struct {
	logger     *logging.Logger
	jupyterMgr *jupyter.ServiceManager
	authConfig *auth.Config
}

// NewHandler creates a new Jupyter handler
func NewHandler(
	logger *logging.Logger,
	jupyterMgr *jupyter.ServiceManager,
	authConfig *auth.Config,
) *Handler {
	return &Handler{
		logger:     logger,
		jupyterMgr: jupyterMgr,
		authConfig: authConfig,
	}
}

// Error codes - using standardized error codes from errors package
const (
	ErrorCodeInvalidRequest       = errors.CodeInvalidRequest
	ErrorCodeAuthenticationFailed = errors.CodeAuthenticationFailed
	ErrorCodePermissionDenied     = errors.CodePermissionDenied
	ErrorCodeResourceNotFound     = errors.CodeResourceNotFound
	ErrorCodeServiceUnavailable   = errors.CodeServiceUnavailable
)

// Permissions
const (
	PermJupyterManage = "jupyter:manage"
	PermJupyterRead   = "jupyter:read"
)

// sendErrorPacket sends an error response packet to the client
func (h *Handler) sendErrorPacket(conn *websocket.Conn, code string, message, details string) error {
	return errors.SendErrorPacket(conn, code, message, details)
}

// sendSuccessPacket sends a success response packet
func (h *Handler) sendSuccessPacket(conn *websocket.Conn, data map[string]any) error {
	return errors.SendSuccessPacket(conn, data)
}

// HandleStartJupyter handles starting a Jupyter service
// Protocol: [api_key_hash:16][workspace_len:1][workspace:var][config_len:2][config:var]
func (h *Handler) HandleStartJupyter(conn *websocket.Conn, payload []byte, user *auth.User) error {
	if len(payload) < 16+1+2 {
		return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "start jupyter payload too short", "")
	}

	offset := 16

	workspaceLen := int(payload[offset])
	offset += 1
	if workspaceLen <= 0 || len(payload) < offset+workspaceLen+2 {
		return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid workspace length", "")
	}
	workspace := string(payload[offset : offset+workspaceLen])
	offset += workspaceLen

	configLen := int(binary.BigEndian.Uint16(payload[offset : offset+2]))
	offset += 2
	if configLen < 0 || len(payload) < offset+configLen {
		return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid config length", "")
	}

	if err := container.ValidateJobName(workspace); err != nil {
		return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid workspace name", err.Error())
	}

	h.logger.Info("starting jupyter service", "workspace", workspace, "user", user.Name)

	// Start Jupyter service
	if h.jupyterMgr == nil {
		return h.sendErrorPacket(conn, ErrorCodeServiceUnavailable, "Jupyter service manager not available", "")
	}

	return h.sendSuccessPacket(conn, map[string]any{
		"workspace": workspace,
		"timestamp": time.Now().UTC(),
	})
}

// HandleStopJupyter handles stopping a Jupyter service
// Protocol: [api_key_hash:16][service_id_len:1][service_id:var]
func (h *Handler) HandleStopJupyter(conn *websocket.Conn, payload []byte, user *auth.User) error {
	if len(payload) < 16+1 {
		return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "stop jupyter payload too short", "")
	}

	offset := 16

	serviceIDLen := int(payload[offset])
	offset += 1
	if serviceIDLen <= 0 || len(payload) < offset+serviceIDLen {
		return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid service ID length", "")
	}
	serviceID := string(payload[offset : offset+serviceIDLen])

	h.logger.Info("stopping jupyter service", "service_id", serviceID, "user", user.Name)

	if h.jupyterMgr == nil {
		return h.sendErrorPacket(conn, ErrorCodeServiceUnavailable, "Jupyter service manager not available", "")
	}

	return h.sendSuccessPacket(conn, map[string]any{
		"service_id": serviceID,
		"timestamp":  time.Now().UTC(),
	})
}

// HandleListJupyter handles listing Jupyter services
// Protocol: [api_key_hash:16]
func (h *Handler) HandleListJupyter(conn *websocket.Conn, payload []byte, user *auth.User) error {
	h.logger.Info("listing jupyter services", "user", user.Name)

	if h.jupyterMgr == nil {
		return h.sendSuccessPacket(conn, map[string]any{
			"services": []any{},
			"count":    0,
		})
	}

	services := h.jupyterMgr.ListServices()

	return h.sendSuccessPacket(conn, map[string]any{
		"services": services,
		"count":    len(services),
	})
}

// HandleListJupyterPackages handles listing packages in a Jupyter service
// Protocol: [api_key_hash:16][service_name_len:1][service_name:var]
func (h *Handler) HandleListJupyterPackages(conn *websocket.Conn, payload []byte, user *auth.User) error {
	if len(payload) < 16+1 {
		return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "list packages payload too short", "")
	}

	offset := 16

	serviceNameLen := int(payload[offset])
	offset += 1
	if serviceNameLen <= 0 || len(payload) < offset+serviceNameLen {
		return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid service name length", "")
	}
	serviceName := string(payload[offset : offset+serviceNameLen])

	h.logger.Info("listing jupyter packages", "service", serviceName, "user", user.Name)

	return h.sendSuccessPacket(conn, map[string]any{
		"service_name": serviceName,
		"packages":     []any{},
		"count":        0,
	})
}

// HandleRemoveJupyter handles removing a Jupyter service
// Protocol: [api_key_hash:16][service_id_len:1][service_id:var][purge:1]
func (h *Handler) HandleRemoveJupyter(conn *websocket.Conn, payload []byte, user *auth.User) error {
	if len(payload) < 16+1+1 {
		return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "remove jupyter payload too short", "")
	}

	offset := 16

	serviceIDLen := int(payload[offset])
	offset += 1
	if serviceIDLen <= 0 || len(payload) < offset+serviceIDLen+1 {
		return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid service ID length", "")
	}
	serviceID := string(payload[offset : offset+serviceIDLen])
	offset += serviceIDLen

	purge := payload[offset] != 0

	h.logger.Info("removing jupyter service", "service_id", serviceID, "purge", purge, "user", user.Name)

	if h.jupyterMgr == nil {
		return h.sendErrorPacket(conn, ErrorCodeServiceUnavailable, "Jupyter service manager not available", "")
	}

	return h.sendSuccessPacket(conn, map[string]any{
		"service_id": serviceID,
		"purged":     purge,
	})
}

// HandleRestoreJupyter handles restoring a Jupyter workspace
// Protocol: [api_key_hash:16][workspace_len:1][workspace:var]
func (h *Handler) HandleRestoreJupyter(conn *websocket.Conn, payload []byte, user *auth.User) error {
	if len(payload) < 16+1 {
		return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "restore jupyter payload too short", "")
	}

	offset := 16

	workspaceLen := int(payload[offset])
	offset += 1
	if workspaceLen <= 0 || len(payload) < offset+workspaceLen {
		return h.sendErrorPacket(conn, ErrorCodeInvalidRequest, "invalid workspace length", "")
	}
	workspace := string(payload[offset : offset+workspaceLen])

	h.logger.Info("restoring jupyter workspace", "workspace", workspace, "user", user.Name)

	if h.jupyterMgr == nil {
		return h.sendErrorPacket(conn, ErrorCodeServiceUnavailable, "Jupyter service manager not available", "")
	}

	return h.sendSuccessPacket(conn, map[string]any{
		"workspace": workspace,
		"restored":  true,
	})
}

// HTTP Handlers for REST API

// ListServicesHTTP handles HTTP requests for listing Jupyter services
func (h *Handler) ListServicesHTTP(w http.ResponseWriter, r *http.Request) {
	if h.jupyterMgr == nil {
		errors.WriteHTTPError(w, http.StatusServiceUnavailable, errors.CodeServiceUnavailable, "Jupyter service manager not available", "")
		return
	}

	services := h.jupyterMgr.ListServices()

	w.Header().Set("Content-Type", "application/json")
	response := map[string]any{
		"services": services,
		"count":    len(services),
	}
	if err := json.NewEncoder(w).Encode(response); err != nil {
		h.logger.Warn("failed to encode services list", "error", err)
	}
}

// StartServiceHTTP handles HTTP requests for starting Jupyter service
func (h *Handler) StartServiceHTTP(w http.ResponseWriter, r *http.Request) {
	if h.jupyterMgr == nil {
		errors.WriteHTTPError(w, http.StatusServiceUnavailable, errors.CodeServiceUnavailable, "Jupyter service manager not available", "")
		return
	}

	var req struct {
		Workspace string         `json:"workspace"`
		Config    map[string]any `json:"config,omitempty"`
	}

	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		errors.WriteHTTPError(w, http.StatusBadRequest, errors.CodeInvalidRequest, "Invalid request body", "")
		return
	}

	if req.Workspace == "" {
		errors.WriteHTTPError(w, http.StatusBadRequest, errors.CodeInvalidRequest, "Workspace name is required", "")
		return
	}

	if err := container.ValidateJobName(req.Workspace); err != nil {
		errors.WriteHTTPError(w, http.StatusBadRequest, errors.CodeInvalidRequest, "Invalid workspace name", err.Error())
		return
	}

	h.logger.Info("starting jupyter service via HTTP", "workspace", req.Workspace)

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusCreated)
	response := map[string]any{
		"workspace": req.Workspace,
		"timestamp": time.Now().UTC(),
	}
	if err := json.NewEncoder(w).Encode(response); err != nil {
		h.logger.Warn("failed to encode start response", "error", err)
	}
}

// StopServiceHTTP handles HTTP requests for stopping Jupyter service
func (h *Handler) StopServiceHTTP(w http.ResponseWriter, r *http.Request) {
	if h.jupyterMgr == nil {
		errors.WriteHTTPError(w, http.StatusServiceUnavailable, errors.CodeServiceUnavailable, "Jupyter service manager not available", "")
		return
	}

	serviceID := r.PathValue("serviceId")
	if serviceID == "" {
		serviceID = r.URL.Query().Get("service_id")
	}

	if serviceID == "" {
		errors.WriteHTTPError(w, http.StatusBadRequest, errors.CodeInvalidRequest, "Service ID is required", "")
		return
	}

	h.logger.Info("stopping jupyter service via HTTP", "service_id", serviceID)

	w.Header().Set("Content-Type", "application/json")
	response := map[string]any{
		"service_id": serviceID,
		"timestamp":  time.Now().UTC(),
	}
	if err := json.NewEncoder(w).Encode(response); err != nil {
		h.logger.Warn("failed to encode stop response", "error", err)
	}
}