fetch_ml/internal/api/helpers/response_helpers.go

// Package helpers provides shared utilities for WebSocket handlers.
package helpers

import (
	"encoding/json"
	"fmt"
	"strings"

	"github.com/jfraeys/fetch_ml/internal/queue"
)

// ErrorCode represents WebSocket error codes
type ErrorCode byte

// TaskErrorMapper maps task errors to error codes
type TaskErrorMapper struct{}

// NewTaskErrorMapper creates a new task error mapper
func NewTaskErrorMapper() *TaskErrorMapper {
	return &TaskErrorMapper{}
}

// MapError maps a task error to an error code based on status and error message
func (m *TaskErrorMapper) MapError(t *queue.Task, defaultCode ErrorCode) ErrorCode {
	if t == nil {
		return defaultCode
	}
	status := strings.ToLower(strings.TrimSpace(t.Status))
	errStr := strings.ToLower(strings.TrimSpace(t.Error))

	if status == "cancelled" {
		return 0x24 // ErrorCodeJobCancelled
	}
	if strings.Contains(errStr, "out of memory") || strings.Contains(errStr, "oom") {
		return 0x30 // ErrorCodeOutOfMemory
	}
	if strings.Contains(errStr, "no space left") || strings.Contains(errStr, "disk full") {
		return 0x31 // ErrorCodeDiskFull
	}
	if strings.Contains(errStr, "rate limit") || strings.Contains(errStr, "too many requests") || strings.Contains(errStr, "throttle") {
		return 0x33 // ErrorCodeServiceUnavailable
	}
	if strings.Contains(errStr, "timed out") || strings.Contains(errStr, "timeout") || strings.Contains(errStr, "deadline") {
		return 0x14 // ErrorCodeTimeout
	}
	if strings.Contains(errStr, "connection refused") || strings.Contains(errStr, "connection reset") || strings.Contains(errStr, "network unreachable") {
		return 0x12 // ErrorCodeNetworkError
	}
	if strings.Contains(errStr, "queue") && strings.Contains(errStr, "not configured") {
		return 0x32 // ErrorCodeInvalidConfiguration
	}

	// Default for worker-side execution failures
	if status == "failed" {
		return 0x23 // ErrorCodeJobExecutionFailed
	}
	return defaultCode
}

// MapJupyterError maps Jupyter task errors to error codes
func (m *TaskErrorMapper) MapJupyterError(t *queue.Task) ErrorCode {
	if t == nil {
		return 0x00 // ErrorCodeUnknownError
	}
	status := strings.ToLower(strings.TrimSpace(t.Status))
	errStr := strings.ToLower(strings.TrimSpace(t.Error))

	if status == "cancelled" {
		return 0x24 // ErrorCodeJobCancelled
	}
	if strings.Contains(errStr, "out of memory") || strings.Contains(errStr, "oom") {
		return 0x30 // ErrorCodeOutOfMemory
	}
	if strings.Contains(errStr, "no space left") || strings.Contains(errStr, "disk full") {
		return 0x31 // ErrorCodeDiskFull
	}
	if strings.Contains(errStr, "rate limit") || strings.Contains(errStr, "too many requests") || strings.Contains(errStr, "throttle") {
		return 0x33 // ErrorCodeServiceUnavailable
	}
	if strings.Contains(errStr, "timed out") || strings.Contains(errStr, "timeout") || strings.Contains(errStr, "deadline") {
		return 0x14 // ErrorCodeTimeout
	}
	if strings.Contains(errStr, "connection refused") || strings.Contains(errStr, "connection reset") || strings.Contains(errStr, "network unreachable") {
		return 0x12 // ErrorCodeNetworkError
	}
	if strings.Contains(errStr, "queue") && strings.Contains(errStr, "not configured") {
		return 0x32 // ErrorCodeInvalidConfiguration
	}

	// Default for worker-side execution failures
	if status == "failed" {
		return 0x23 // ErrorCodeJobExecutionFailed
	}
	return 0x00 // ErrorCodeUnknownError
}

// ResourceRequest represents resource requirements
type ResourceRequest struct {
	GPUMemory string
	CPU       int
	MemoryGB  int
	GPU       int
}

// ParseResourceRequest parses an optional resource request from bytes.
// Format: [cpu:1][memory_gb:1][gpu:1][gpu_mem_len:1][gpu_mem:var]
// If payload is empty, returns nil.
func ParseResourceRequest(payload []byte) (*ResourceRequest, error) {
	if len(payload) == 0 {
		return nil, nil
	}
	if len(payload) < 4 {
		return nil, fmt.Errorf("resource payload too short")
	}
	cpu := int(payload[0])
	mem := int(payload[1])
	gpu := int(payload[2])
	gpuMemLen := int(payload[3])
	if gpuMemLen < 0 || len(payload) < 4+gpuMemLen {
		return nil, fmt.Errorf("invalid gpu memory length")
	}
	gpuMem := ""
	if gpuMemLen > 0 {
		gpuMem = string(payload[4 : 4+gpuMemLen])
	}
	return &ResourceRequest{CPU: cpu, MemoryGB: mem, GPU: gpu, GPUMemory: gpuMem}, nil
}

// JSONResponseBuilder helps build JSON data responses
type JSONResponseBuilder struct {
	data any
}

// NewJSONResponseBuilder creates a new JSON response builder
func NewJSONResponseBuilder(data any) *JSONResponseBuilder {
	return &JSONResponseBuilder{data: data}
}

// Build marshals the data to JSON
func (b *JSONResponseBuilder) Build() ([]byte, error) {
	return json.Marshal(b.data)
}

// BuildOrEmpty marshals the data to JSON or returns empty array on error
func (b *JSONResponseBuilder) BuildOrEmpty() []byte {
	data, err := json.Marshal(b.data)
	if err != nil {
		return []byte("[]")
	}
	return data
}

// StringPtr returns a pointer to a string
func StringPtr(s string) *string {
	return &s
}

// IntPtr returns a pointer to an int
func IntPtr(i int) *int {
	return &i
}

// MarshalJSONOrEmpty marshals data to JSON or returns empty array on error
func MarshalJSONOrEmpty(data any) []byte {
	b, err := json.Marshal(data)
	if err != nil {
		return []byte("[]")
	}
	return b
}

// MarshalJSONBytes marshals data to JSON bytes with error handling
func MarshalJSONBytes(data any) ([]byte, error) {
	return json.Marshal(data)
}

// IsEmptyJSON checks if JSON data is empty or "null"
func IsEmptyJSON(data []byte) bool {
	if len(data) == 0 {
		return true
	}
	// Check for "null", "[]", "{}" or empty after trimming
	s := strings.TrimSpace(string(data))
	return s == "" || s == "null" || s == "[]" || s == "{}"
}