package protocol

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

// encode is a helper function that encodes any message with the given message type
func encode(w io.Writer, msgType MessageType, payload interface{}) error {
	var payloadBytes []byte
	var err error

	// Only marshal if payload is not nil
	if payload != nil {
		payloadBytes, err = json.Marshal(payload)
		if err != nil {
			return fmt.Errorf("failed to encode payload: %w", err)
		}
	}

	// Write message type (1 byte)
	if err := binary.Write(w, binary.BigEndian, msgType); err != nil {
		return fmt.Errorf("failed to write message type: %w", err)
	}

	// Write payload length (4 bytes, big-endian)
	if err := binary.Write(w, binary.BigEndian, uint32(len(payloadBytes))); err != nil {
		return fmt.Errorf("failed to write payload length: %w", err)
	}

	// Write JSON payload if we have one
	if len(payloadBytes) > 0 {
		if _, err := w.Write(payloadBytes); err != nil {
			return fmt.Errorf("failed to write payload: %w", err)
		}
	}

	return nil
}