// Package frame implements Aura's wire framing: the 5-byte protocol header and the
// application-level Frame{Data,Ping,Pong,Close}.
//
// This is a byte-for-byte port of crates/aura-proto/src/frame.rs. The Rust unit tests in that
// file are the wire spec; matching them here keeps the Go port interoperable with the Rust
// server.
//
// Wire layout (from docs/protocol.md §6.1):
//
//	byte 0     : msg_type (u8)
//	bytes 1..4 : length   (u24, big-endian) — payload length in bytes
//	byte 4     : version  = 0x01
//	bytes 5..  : payload  (length bytes)
package frame

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

// HeaderLen is the size of the protocol header in bytes.
const HeaderLen = 5

// ProtocolVersion is the constant carried in byte 4 of every header.
const ProtocolVersion byte = 0x01

// MaxPayloadLen is the largest payload expressible by the u24 length field.
const MaxPayloadLen = 0x00FF_FFFF

// MsgType is the on-wire message-type discriminant carried in byte 0 of the header.
type MsgType byte

// Message-type bytes (must match the Rust MsgType repr in aura-proto/frame.rs).
const (
	MsgClientHello MsgType = 0x01
	MsgServerHello MsgType = 0x02
	MsgClientAuth  MsgType = 0x03
	MsgServerAuth  MsgType = 0x04
	MsgFinished    MsgType = 0x05
	MsgData        MsgType = 0x06
	MsgAlert       MsgType = 0xFF
)

// String returns the short name of the message type, for logs.
func (m MsgType) String() string {
	switch m {
	case MsgClientHello:
		return "ClientHello"
	case MsgServerHello:
		return "ServerHello"
	case MsgClientAuth:
		return "ClientAuth"
	case MsgServerAuth:
		return "ServerAuth"
	case MsgFinished:
		return "Finished"
	case MsgData:
		return "Data"
	case MsgAlert:
		return "Alert"
	default:
		return fmt.Sprintf("MsgType(0x%02X)", byte(m))
	}
}

// Errors returned by the codec. They mirror the ProtoError variants the Rust side returns so
// callers can map them onto identical wire alerts.
var (
	ErrFrameTooLarge   = errors.New("aura/frame: payload exceeds 16 MiB u24 length field")
	ErrBadVersion      = errors.New("aura/frame: header byte 4 is not protocol version 0x01")
	ErrUnknownMsgType  = errors.New("aura/frame: unknown message-type byte")
	ErrMalformedFrame  = errors.New("aura/frame: malformed application frame")
)

// EncodeHeader builds a 5-byte header for msgType carrying a payload of payloadLen bytes.
func EncodeHeader(msgType MsgType, payloadLen int) ([HeaderLen]byte, error) {
	var h [HeaderLen]byte
	if payloadLen < 0 || payloadLen > MaxPayloadLen {
		return h, fmt.Errorf("%w: len=%d", ErrFrameTooLarge, payloadLen)
	}
	h[0] = byte(msgType)
	// u24 big-endian.
	h[1] = byte((payloadLen >> 16) & 0xFF)
	h[2] = byte((payloadLen >> 8) & 0xFF)
	h[3] = byte(payloadLen & 0xFF)
	h[4] = ProtocolVersion
	return h, nil
}

// DecodeHeader parses a 5-byte header into (msgType, payloadLen).
func DecodeHeader(h [HeaderLen]byte) (MsgType, int, error) {
	if h[4] != ProtocolVersion {
		return 0, 0, fmt.Errorf("%w: got 0x%02X", ErrBadVersion, h[4])
	}
	mt := MsgType(h[0])
	switch mt {
	case MsgClientHello, MsgServerHello, MsgClientAuth, MsgServerAuth, MsgFinished, MsgData, MsgAlert:
		// recognized
	default:
		return 0, 0, fmt.Errorf("%w: got 0x%02X", ErrUnknownMsgType, h[0])
	}
	plen := int(h[1])<<16 | int(h[2])<<8 | int(h[3])
	return mt, plen, nil
}

// RawFrame is a frame as it was on the wire: type, header bytes (useful for AEAD AAD and the
// handshake transcript hash), and payload bytes.
type RawFrame struct {
	MsgType MsgType
	Header  [HeaderLen]byte
	Payload []byte
}

// WireBytes returns header || payload in a fresh slice — used to feed the transcript hash, which
// hashes the bytes exactly as transmitted.
func (rf *RawFrame) WireBytes() []byte {
	out := make([]byte, 0, HeaderLen+len(rf.Payload))
	out = append(out, rf.Header[:]...)
	out = append(out, rf.Payload...)
	return out
}

// WriteFrame serializes header || payload and writes it to w. Single Write, so on a streaming
// transport a single TCP segment is preferred.
func WriteFrame(w io.Writer, msgType MsgType, payload []byte) error {
	h, err := EncodeHeader(msgType, len(payload))
	if err != nil {
		return err
	}
	buf := make([]byte, 0, HeaderLen+len(payload))
	buf = append(buf, h[:]...)
	buf = append(buf, payload...)
	_, err = w.Write(buf)
	return err
}

// ReadFrame reads one full frame (header || payload) from r.
func ReadFrame(r io.Reader) (*RawFrame, error) {
	var h [HeaderLen]byte
	if _, err := io.ReadFull(r, h[:]); err != nil {
		return nil, err
	}
	mt, plen, err := DecodeHeader(h)
	if err != nil {
		return nil, err
	}
	payload := make([]byte, plen)
	if plen > 0 {
		if _, err := io.ReadFull(r, payload); err != nil {
			return nil, err
		}
	}
	return &RawFrame{MsgType: mt, Header: h, Payload: payload}, nil
}

// ----------------------------------------------------------------------------------------------
// Application frames (§6.3) — Data, Ping, Pong, Close, Control.
// ----------------------------------------------------------------------------------------------

// FrameKind identifies the Application-frame variant.
type FrameKind byte

// On-wire frame tags (must match crates/aura-proto/src/frame.rs frame_tag::*).
const (
	FrameData  FrameKind = 0x01
	FramePing  FrameKind = 0x02
	FramePong  FrameKind = 0x03
	FrameClose FrameKind = 0x04
)

// Frame is the post-handshake application payload carried inside an AEAD-sealed MsgData record.
// One Frame is mapped to one of the four variants by Kind.
type Frame struct {
	Kind     FrameKind
	StreamID uint32 // Data only
	Payload  []byte // Data only
	Seq      uint32 // Ping / Pong only
	Code     byte   // Close only
	Reason   string // Close only
}

// EncodeFrame serializes f into its compact byte encoding (all multi-byte ints big-endian):
//
//	Data  : 0x01 || stream_id(u32) || payload
//	Ping  : 0x02 || seq(u32)
//	Pong  : 0x03 || seq(u32)
//	Close : 0x04 || code(u8) || reason_len(u32) || reason_utf8
func EncodeFrame(f *Frame) []byte {
	switch f.Kind {
	case FrameData:
		out := make([]byte, 1+4+len(f.Payload))
		out[0] = byte(FrameData)
		binary.BigEndian.PutUint32(out[1:5], f.StreamID)
		copy(out[5:], f.Payload)
		return out
	case FramePing:
		out := make([]byte, 1+4)
		out[0] = byte(FramePing)
		binary.BigEndian.PutUint32(out[1:5], f.Seq)
		return out
	case FramePong:
		out := make([]byte, 1+4)
		out[0] = byte(FramePong)
		binary.BigEndian.PutUint32(out[1:5], f.Seq)
		return out
	case FrameClose:
		reason := []byte(f.Reason)
		out := make([]byte, 1+1+4+len(reason))
		out[0] = byte(FrameClose)
		out[1] = f.Code
		binary.BigEndian.PutUint32(out[2:6], uint32(len(reason)))
		copy(out[6:], reason)
		return out
	default:
		// Programmer error — encode nothing rather than panic so call sites can defensively
		// inspect the returned length.
		return nil
	}
}

// DecodeFrame parses one byte-encoded Frame (the inverse of EncodeFrame).
func DecodeFrame(b []byte) (*Frame, error) {
	if len(b) == 0 {
		return nil, fmt.Errorf("%w: empty frame", ErrMalformedFrame)
	}
	tag := FrameKind(b[0])
	rest := b[1:]
	switch tag {
	case FrameData:
		if len(rest) < 4 {
			return nil, fmt.Errorf("%w: Data: missing stream_id", ErrMalformedFrame)
		}
		sid := binary.BigEndian.Uint32(rest[:4])
		// Payload is everything after the 4-byte stream_id.
		payload := make([]byte, len(rest)-4)
		copy(payload, rest[4:])
		return &Frame{Kind: FrameData, StreamID: sid, Payload: payload}, nil
	case FramePing:
		if len(rest) < 4 {
			return nil, fmt.Errorf("%w: Ping: truncated seq", ErrMalformedFrame)
		}
		return &Frame{Kind: FramePing, Seq: binary.BigEndian.Uint32(rest[:4])}, nil
	case FramePong:
		if len(rest) < 4 {
			return nil, fmt.Errorf("%w: Pong: truncated seq", ErrMalformedFrame)
		}
		return &Frame{Kind: FramePong, Seq: binary.BigEndian.Uint32(rest[:4])}, nil
	case FrameClose:
		if len(rest) < 1 {
			return nil, fmt.Errorf("%w: Close: missing code", ErrMalformedFrame)
		}
		code := rest[0]
		if len(rest) < 5 {
			return nil, fmt.Errorf("%w: Close: missing reason_len", ErrMalformedFrame)
		}
		rlen := int(binary.BigEndian.Uint32(rest[1:5]))
		if len(rest) < 5+rlen {
			return nil, fmt.Errorf("%w: Close: truncated reason", ErrMalformedFrame)
		}
		// We do not enforce strict UTF-8 here (Go strings can hold any bytes); the Rust side
		// rejects non-UTF-8 in this slot, so peers that follow the spec only ever send valid
		// strings.
		return &Frame{Kind: FrameClose, Code: code, Reason: string(rest[5 : 5+rlen])}, nil
	default:
		return nil, fmt.Errorf("%w: unknown frame tag 0x%02X", ErrMalformedFrame, byte(tag))
	}
}