package encryption

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

// MaxFrameSize is the maximum AES-GCM ciphertext length accepted on a framed
// TCP tunnel (uint32 length prefix + ciphertext).
const MaxFrameSize = 1024 * 1024 // 1 MiB

// WriteEncryptedFrame encrypts plaintext with AES-GCM (EncryptData) and writes
// a uint32 big-endian length prefix followed by the ciphertext.
func WriteEncryptedFrame(w io.Writer, plaintext, key []byte) error {
	ciphertext, err := EncryptData(plaintext, key)
	if err != nil {
		return err
	}
	if len(ciphertext) == 0 {
		return fmt.Errorf("encrypted data cannot be empty")
	}
	if len(ciphertext) > MaxFrameSize {
		return fmt.Errorf("frame too large: %d bytes (max %d)", len(ciphertext), MaxFrameSize)
	}

	var length [4]byte
	binary.BigEndian.PutUint32(length[:], uint32(len(ciphertext)))
	if err := writeFull(w, length[:]); err != nil {
		return err
	}
	return writeFull(w, ciphertext)
}

// ReadEncryptedFrame reads a uint32 big-endian length, caps it, reads the
// ciphertext, and decrypts it with AES-GCM (DecryptData).
func ReadEncryptedFrame(r io.Reader, key []byte) ([]byte, error) {
	var length uint32
	if err := binary.Read(r, binary.BigEndian, &length); err != nil {
		return nil, err
	}
	if length == 0 {
		return nil, fmt.Errorf("frame length cannot be zero")
	}
	if length > MaxFrameSize {
		return nil, fmt.Errorf("frame too large: %d bytes (max %d)", length, MaxFrameSize)
	}

	ciphertext := make([]byte, length)
	if _, err := io.ReadFull(r, ciphertext); err != nil {
		return nil, fmt.Errorf("failed to read frame data: %v", err)
	}
	return DecryptData(ciphertext, key)
}

func writeFull(w io.Writer, data []byte) error {
	for len(data) > 0 {
		n, err := w.Write(data)
		if err != nil {
			return err
		}
		if n == 0 {
			return io.ErrShortWrite
		}
		data = data[n:]
	}
	return nil
}