fetch_ml/internal/api/protocol.go

package api

import (
	"encoding/binary"
	"encoding/json"
	"fmt"
	"time"
)

// Response packet types
const (
	PacketTypeSuccess  = 0x00
	PacketTypeError    = 0x01
	PacketTypeProgress = 0x02
	PacketTypeStatus   = 0x03
	PacketTypeData     = 0x04
	PacketTypeLog      = 0x05
)

// Error codes
const (
	ErrorCodeUnknownError          = 0x00
	ErrorCodeInvalidRequest        = 0x01
	ErrorCodeAuthenticationFailed  = 0x02
	ErrorCodePermissionDenied      = 0x03
	ErrorCodeResourceNotFound      = 0x04
	ErrorCodeResourceAlreadyExists = 0x05

	ErrorCodeServerOverloaded = 0x10
	ErrorCodeDatabaseError    = 0x11
	ErrorCodeNetworkError     = 0x12
	ErrorCodeStorageError     = 0x13
	ErrorCodeTimeout          = 0x14

	ErrorCodeJobNotFound        = 0x20
	ErrorCodeJobAlreadyRunning  = 0x21
	ErrorCodeJobFailedToStart   = 0x22
	ErrorCodeJobExecutionFailed = 0x23
	ErrorCodeJobCancelled       = 0x24

	ErrorCodeOutOfMemory          = 0x30
	ErrorCodeDiskFull             = 0x31
	ErrorCodeInvalidConfiguration = 0x32
	ErrorCodeServiceUnavailable   = 0x33
)

// Progress types
const (
	ProgressTypePercentage       = 0x00
	ProgressTypeStage            = 0x01
	ProgressTypeMessage          = 0x02
	ProgressTypeBytesTransferred = 0x03
)

// Log levels
const (
	LogLevelDebug = 0x00
	LogLevelInfo  = 0x01
	LogLevelWarn  = 0x02
	LogLevelError = 0x03
)

// ResponsePacket represents a structured response packet
type ResponsePacket struct {
	PacketType byte
	Timestamp  uint64

	// Success fields
	SuccessMessage string

	// Error fields
	ErrorCode    byte
	ErrorMessage string
	ErrorDetails string

	// Progress fields
	ProgressType    byte
	ProgressValue   uint32
	ProgressTotal   uint32
	ProgressMessage string

	// Status fields
	StatusData string

	// Data fields
	DataType    string
	DataPayload []byte

	// Log fields
	LogLevel   byte
	LogMessage string
}

// NewSuccessPacket creates a success response packet
func NewSuccessPacket(message string) *ResponsePacket {
	return &ResponsePacket{
		PacketType:     PacketTypeSuccess,
		Timestamp:      uint64(time.Now().Unix()),
		SuccessMessage: message,
	}
}

// NewSuccessPacketWithPayload creates a success response packet with JSON payload
func NewSuccessPacketWithPayload(message string, payload interface{}) *ResponsePacket {
	// Convert payload to JSON for the DataPayload field
	payloadBytes, _ := json.Marshal(payload)

	return &ResponsePacket{
		PacketType:     PacketTypeData,
		Timestamp:      uint64(time.Now().Unix()),
		SuccessMessage: message,
		DataType:       "status",
		DataPayload:    payloadBytes,
	}
}

// NewErrorPacket creates an error response packet
func NewErrorPacket(errorCode byte, message string, details string) *ResponsePacket {
	return &ResponsePacket{
		PacketType:   PacketTypeError,
		Timestamp:    uint64(time.Now().Unix()),
		ErrorCode:    errorCode,
		ErrorMessage: message,
		ErrorDetails: details,
	}
}

// NewProgressPacket creates a progress response packet
func NewProgressPacket(progressType byte, value uint32, total uint32, message string) *ResponsePacket {
	return &ResponsePacket{
		PacketType:      PacketTypeProgress,
		Timestamp:       uint64(time.Now().Unix()),
		ProgressType:    progressType,
		ProgressValue:   value,
		ProgressTotal:   total,
		ProgressMessage: message,
	}
}

// NewStatusPacket creates a status response packet
func NewStatusPacket(data string) *ResponsePacket {
	return &ResponsePacket{
		PacketType: PacketTypeStatus,
		Timestamp:  uint64(time.Now().Unix()),
		StatusData: data,
	}
}

// NewDataPacket creates a data response packet
func NewDataPacket(dataType string, payload []byte) *ResponsePacket {
	return &ResponsePacket{
		PacketType:  PacketTypeData,
		Timestamp:   uint64(time.Now().Unix()),
		DataType:    dataType,
		DataPayload: payload,
	}
}

// NewLogPacket creates a log response packet
func NewLogPacket(level byte, message string) *ResponsePacket {
	return &ResponsePacket{
		PacketType: PacketTypeLog,
		Timestamp:  uint64(time.Now().Unix()),
		LogLevel:   level,
		LogMessage: message,
	}
}

// Serialize converts the packet to binary format
func (p *ResponsePacket) Serialize() ([]byte, error) {
	var buf []byte

	// Packet type
	buf = append(buf, p.PacketType)

	// Timestamp (8 bytes, big-endian)
	timestampBytes := make([]byte, 8)
	binary.BigEndian.PutUint64(timestampBytes, p.Timestamp)
	buf = append(buf, timestampBytes...)

	// Packet-specific data
	switch p.PacketType {
	case PacketTypeSuccess:
		buf = append(buf, serializeString(p.SuccessMessage)...)

	case PacketTypeError:
		buf = append(buf, p.ErrorCode)
		buf = append(buf, serializeString(p.ErrorMessage)...)
		buf = append(buf, serializeString(p.ErrorDetails)...)

	case PacketTypeProgress:
		buf = append(buf, p.ProgressType)
		valueBytes := make([]byte, 4)
		binary.BigEndian.PutUint32(valueBytes, p.ProgressValue)
		buf = append(buf, valueBytes...)

		totalBytes := make([]byte, 4)
		binary.BigEndian.PutUint32(totalBytes, p.ProgressTotal)
		buf = append(buf, totalBytes...)

		buf = append(buf, serializeString(p.ProgressMessage)...)

	case PacketTypeStatus:
		buf = append(buf, serializeString(p.StatusData)...)

	case PacketTypeData:
		buf = append(buf, serializeString(p.DataType)...)
		buf = append(buf, serializeBytes(p.DataPayload)...)

	case PacketTypeLog:
		buf = append(buf, p.LogLevel)
		buf = append(buf, serializeString(p.LogMessage)...)

	default:
		return nil, fmt.Errorf("unknown packet type: %d", p.PacketType)
	}

	return buf, nil
}

// serializeString writes a string with 2-byte length prefix
func serializeString(s string) []byte {
	length := uint16(len(s))
	buf := make([]byte, 2+len(s))
	binary.BigEndian.PutUint16(buf[:2], length)
	copy(buf[2:], s)
	return buf
}

// serializeBytes writes bytes with 4-byte length prefix
func serializeBytes(b []byte) []byte {
	length := uint32(len(b))
	buf := make([]byte, 4+len(b))
	binary.BigEndian.PutUint32(buf[:4], length)
	copy(buf[4:], b)
	return buf
}

// GetErrorMessage returns a human-readable error message for an error code
func GetErrorMessage(code byte) string {
	switch code {
	case ErrorCodeUnknownError:
		return "Unknown error occurred"
	case ErrorCodeInvalidRequest:
		return "Invalid request format"
	case ErrorCodeAuthenticationFailed:
		return "Authentication failed"
	case ErrorCodePermissionDenied:
		return "Permission denied"
	case ErrorCodeResourceNotFound:
		return "Resource not found"
	case ErrorCodeResourceAlreadyExists:
		return "Resource already exists"

	case ErrorCodeServerOverloaded:
		return "Server is overloaded"
	case ErrorCodeDatabaseError:
		return "Database error occurred"
	case ErrorCodeNetworkError:
		return "Network error occurred"
	case ErrorCodeStorageError:
		return "Storage error occurred"
	case ErrorCodeTimeout:
		return "Operation timed out"

	case ErrorCodeJobNotFound:
		return "Job not found"
	case ErrorCodeJobAlreadyRunning:
		return "Job is already running"
	case ErrorCodeJobFailedToStart:
		return "Job failed to start"
	case ErrorCodeJobExecutionFailed:
		return "Job execution failed"
	case ErrorCodeJobCancelled:
		return "Job was cancelled"

	case ErrorCodeOutOfMemory:
		return "Server out of memory"
	case ErrorCodeDiskFull:
		return "Server disk full"
	case ErrorCodeInvalidConfiguration:
		return "Invalid server configuration"
	case ErrorCodeServiceUnavailable:
		return "Service temporarily unavailable"

	default:
		return "Unknown error code"
	}
}

// GetLogLevelName returns the name for a log level
func GetLogLevelName(level byte) string {
	switch level {
	case LogLevelDebug:
		return "DEBUG"
	case LogLevelInfo:
		return "INFO"
	case LogLevelWarn:
		return "WARN"
	case LogLevelError:
		return "ERROR"
	default:
		return "UNKNOWN"
	}
}