package dispatcher

import (
	"crypto/rand"
	"crypto/sha256"
	"errors"
	"fmt"

	"github.com/fxamacker/cbor/v2"
)

const (
	ProtocolVersion  = 1
	MaxSphinxPayload = 2000
	FrameDataBytes   = 1500
	MaxFrames        = (MaxEnvelopeBytes + FrameDataBytes - 1) / FrameDataBytes
)

type Frame struct {
	Version uint8    `cbor:"1,keyasint"`
	ID      [16]byte `cbor:"2,keyasint"`
	Digest  [32]byte `cbor:"3,keyasint"`
	Index   uint16   `cbor:"4,keyasint"`
	Total   uint16   `cbor:"5,keyasint"`
	Data    []byte   `cbor:"6,keyasint"`
}

func SplitEnvelope(encoded []byte) ([][]byte, error) {
	if len(encoded) == 0 || len(encoded) > MaxEnvelopeBytes {
		return nil, errors.New("invalid encoded envelope size")
	}

	var id [16]byte
	if _, err := rand.Read(id[:]); err != nil {
		return nil, fmt.Errorf("generate message id: %w", err)
	}
	digest := sha256.Sum256(encoded)
	total := (len(encoded) + FrameDataBytes - 1) / FrameDataBytes
	if total > MaxFrames {
		return nil, errors.New("encoded envelope requires too many frames")
	}

	frames := make([][]byte, 0, total)
	for index, start := 0, 0; start < len(encoded); index, start = index+1, start+FrameDataBytes {
		end := min(start+FrameDataBytes, len(encoded))
		frame := Frame{
			Version: ProtocolVersion,
			ID:      id,
			Digest:  digest,
			Index:   uint16(index),
			Total:   uint16(total),
			Data:    encoded[start:end],
		}
		payload, err := cbor.Marshal(frame)
		if err != nil {
			return nil, fmt.Errorf("encode frame %d: %w", index, err)
		}
		if len(payload) > MaxSphinxPayload {
			return nil, fmt.Errorf("frame %d exceeds Sphinx payload limit", index)
		}
		frames = append(frames, payload)
	}
	return frames, nil
}

func DecodeFrame(payload []byte) (Frame, error) {
	if len(payload) == 0 || len(payload) > MaxSphinxPayload {
		return Frame{}, errors.New("invalid frame size")
	}
	var frame Frame
	if err := cbor.Unmarshal(payload, &frame); err != nil {
		return Frame{}, fmt.Errorf("decode frame: %w", err)
	}
	if frame.Version != ProtocolVersion {
		return Frame{}, errors.New("unsupported frame version")
	}
	if frame.Total == 0 || int(frame.Total) > MaxFrames || frame.Index >= frame.Total {
		return Frame{}, errors.New("invalid frame sequence")
	}
	if len(frame.Data) == 0 || len(frame.Data) > FrameDataBytes {
		return Frame{}, errors.New("invalid frame data size")
	}
	return frame, nil
}
