// Package protocol implements the obs-teleport wire format.
//
// Every packet on the wire starts with a Header followed by the payload.
// Video packets are "JPEG" (Header + ImageHeader + JPEG bytes), audio
// packets are "WAVE" (Header + WaveHeader + raw PCM).
//
// All integers/float32s are little-endian, matching the reference
// implementation https://github.com/fzwoch/obs-teleport (GPL-2.0).
package protocol

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

// Header is the fixed-size prefix of every packet.
//
//	Type [4]byte | Timestamp uint64 | Size int32  (all little-endian)
type Header struct {
	Type      [4]byte
	Timestamp uint64
	Size      int32
}

// ImageHeader is written after Header for video packets. It carries the
// color parameters OBS derives from its rendering pipeline (BT.709/FULL in
// teleportfling, but the receiver happily forwards whatever we send).
//
//	ColorMatrix [16]float32 | ColorRangeMin [3]float32 | ColorRangeMax [3]float32
type ImageHeader struct {
	ColorMatrix   [16]float32
	ColorRangeMin [3]float32
	ColorRangeMax [3]float32
}

// WaveHeader is written after Header for audio packets. Format uses the
// OBS AUDIO_FORMAT_* enum values (see audioFormat_* consts below); the
// receiver feeds these straight into obs_source_output_audio.
//
//	Format int32 | SampleRate int32 | Speakers int32 | Frames int32
type WaveHeader struct {
	Format     int32
	SampleRate int32
	Speakers   int32
	Frames     int32
}

// AnnouncePayload is the JSON document broadcast on the multicast discovery
// group. It tells OBS receivers where to connect and what the pipe carries.
type AnnouncePayload struct {
	Name          string
	Port          int
	AudioAndVideo bool
	Version       string
	Address       string `json:",omitempty"`
}

// Packet type identifiers used in Header.Type.
var (
	VideoType = [4]byte{'J', 'P', 'E', 'G'}
	AudioType = [4]byte{'W', 'A', 'V', 'E'}
)

// OBS audio output format enum values written into WaveHeader.Format. Only
// the interleaved (non-planar) forms appear on the wire; the reference
// collapses the planar forms to these when packetizing. See obs-audio.h.
const (
	AudioFormatU8  int32 = 1 // unsigned 8-bit
	AudioFormatS16 int32 = 2 // signed 16-bit little-endian
	AudioFormatS32 int32 = 3 // signed 32-bit little-endian
	AudioFormatF32 int32 = 4 // IEEE-754 float little-endian
)

// headerSize / imageHeaderSize / waveHeaderSize are the fixed wire sizes.
const (
	headerSize      = 16               // Type[4] + Timestamp[8] + Size[4]
	imageHeaderSize = 16*4 + 3*4 + 3*4 // 16 float32s + 6 float32s
	waveHeaderSize  = 4 * 4            // 4 int32s
)

// DefaultBT709Full returns the ImageHeader describing a BT.709, full-range
// YCbCr stream, matching obs-teleport's fallback (video_format_get_parameters
// for VIDEO_CS_709 + VIDEO_RANGE_FULL, 8-bit). The ColorMatrix is the YUV→RGB
// conversion matrix OBS applies when rendering the frame.
func DefaultBT709Full() ImageHeader {
	var m [16]float32
	copy(m[:], bt709FullMatrix[:])
	return ImageHeader{
		ColorMatrix:   m,
		ColorRangeMin: [3]float32{0, 0, 0},
		ColorRangeMax: [3]float32{1, 1, 1},
	}
}

// bt709FullMatrix is the YUV→RGB matrix produced by OBS's
// video_format_get_parameters(VIDEO_CS_709, VIDEO_RANGE_FULL) for 8-bit
// video. The trailing column is the chroma-offset term that centres
// Cb/Cr at 0.5.
var bt709FullMatrix = [16]float32{
	1, 0, 1.5748, -0.790488,
	1, -0.187324, -0.468124, 0.329009,
	1, 1.8556, 0, -0.931439,
	0, 0, 0, 1,
}

// WritePacket serializes the full packet: Header, optional ImageHeader or
// WaveHeader, then the payload bytes. It returns the wire slice.
//
// image/wave selects which sub-header is emitted; passing both is an error,
// passing neither (with payload) produces a header-only packet. A nil header
// and the empty type is used by tests to size check framing.
func WritePacket(h Header, img *ImageHeader, wave *WaveHeader, payload []byte) ([]byte, error) {
	if img != nil && wave != nil {
		return nil, errors.New("protocol: both image and wave headers set")
	}

	out := make([]byte, 0, headerSize+len(payload)+imageHeaderSize)
	var scratch [headerSize]byte
	binary.LittleEndian.PutUint32(scratch[0:4], encodeType(h.Type))
	binary.LittleEndian.PutUint64(scratch[4:12], h.Timestamp)
	binary.LittleEndian.PutUint32(scratch[12:16], uint32(h.Size))
	out = append(out, scratch[:]...)

	if img != nil {
		out = appendImageHeader(out, img)
	}
	if wave != nil {
		out = appendWaveHeader(out, wave)
	}
	out = append(out, payload...)
	return out, nil
}

// ReadPacket reads one complete packet from r (Header + sub-header + payload)
// and returns the payload bytes plus the parsed headers.
func ReadPacket(r io.Reader) (Header, *ImageHeader, *WaveHeader, []byte, error) {
	var hdr [headerSize]byte
	if _, err := io.ReadFull(r, hdr[:]); err != nil {
		return Header{}, nil, nil, nil, err
	}

	h := Header{
		Type:      [4]byte{hdr[0], hdr[1], hdr[2], hdr[3]},
		Timestamp: binary.LittleEndian.Uint64(hdr[4:12]),
		Size:      int32(binary.LittleEndian.Uint32(hdr[12:16])),
	}

	if h.Size < 0 || int64(h.Size)+imageHeaderSize > 1<<30 {
		return Header{}, nil, nil, nil, errors.New("protocol: invalid packet size")
	}

	var img *ImageHeader
	var wave *WaveHeader
	switch h.Type {
	case VideoType:
		var ih imageHeaderBytes
		if _, err := io.ReadFull(r, ih[:]); err != nil {
			return Header{}, nil, nil, nil, err
		}
		img = &ImageHeader{}
		decodeImageHeader(ih, img)
	case AudioType:
		var wh waveHeaderBytes
		if _, err := io.ReadFull(r, wh[:]); err != nil {
			return Header{}, nil, nil, nil, err
		}
		wave = &WaveHeader{}
		decodeWaveHeader(wh, wave)
	default:
		return Header{}, nil, nil, nil, errors.New("protocol: unknown packet type")
	}

	payload := make([]byte, h.Size)
	if _, err := io.ReadFull(r, payload); err != nil {
		return Header{}, nil, nil, nil, err
	}
	return h, img, wave, payload, nil
}

func encodeType(t [4]byte) uint32 {
	return uint32(t[0]) | uint32(t[1])<<8 | uint32(t[2])<<16 | uint32(t[3])<<24
}

type imageHeaderBytes [imageHeaderSize]byte
type waveHeaderBytes [waveHeaderSize]byte

func appendImageHeader(dst []byte, img *ImageHeader) []byte {
	var b imageHeaderBytes
	for i, f := range img.ColorMatrix {
		binary.LittleEndian.PutUint32(b[i*4:], math.Float32bits(f))
	}
	off := 16 * 4
	for i, f := range img.ColorRangeMin {
		binary.LittleEndian.PutUint32(b[off+i*4:], math.Float32bits(f))
	}
	off += 3 * 4
	for i, f := range img.ColorRangeMax {
		binary.LittleEndian.PutUint32(b[off+i*4:], math.Float32bits(f))
	}
	return append(dst, b[:]...)
}

func decodeImageHeader(b imageHeaderBytes, img *ImageHeader) {
	for i := range img.ColorMatrix {
		img.ColorMatrix[i] = math.Float32frombits(binary.LittleEndian.Uint32(b[i*4:]))
	}
	off := 16 * 4
	for i := range img.ColorRangeMin {
		img.ColorRangeMin[i] = math.Float32frombits(binary.LittleEndian.Uint32(b[off+i*4:]))
	}
	off += 3 * 4
	for i := range img.ColorRangeMax {
		img.ColorRangeMax[i] = math.Float32frombits(binary.LittleEndian.Uint32(b[off+i*4:]))
	}
}

func appendWaveHeader(dst []byte, w *WaveHeader) []byte {
	var b waveHeaderBytes
	binary.LittleEndian.PutUint32(b[0:4], uint32(w.Format))
	binary.LittleEndian.PutUint32(b[4:8], uint32(w.SampleRate))
	binary.LittleEndian.PutUint32(b[8:12], uint32(w.Speakers))
	binary.LittleEndian.PutUint32(b[12:16], uint32(w.Frames))
	return append(dst, b[:]...)
}

func decodeWaveHeader(b waveHeaderBytes, w *WaveHeader) {
	w.Format = int32(binary.LittleEndian.Uint32(b[0:4]))
	w.SampleRate = int32(binary.LittleEndian.Uint32(b[4:8]))
	w.Speakers = int32(binary.LittleEndian.Uint32(b[8:12]))
	w.Frames = int32(binary.LittleEndian.Uint32(b[12:16]))
}