package aoiprxkit

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

const (
	streamMagic = "ARX1"

	StreamCodecPCM_S32LE uint8 = 1
	StreamCodecOpus      uint8 = 2 // reserved for later phases
)

type StreamHeader struct {
	Codec        uint8
	Channels     uint8
	SampleRateHz uint32
	FrameSamples uint32
	Sequence     uint32
	Timestamp    uint64
	PayloadBytes uint32
}

func ReadStreamHeader(r io.Reader) (StreamHeader, error) {
	var raw [26]byte
	if _, err := io.ReadFull(r, raw[:]); err != nil {
		return StreamHeader{}, err
	}
	if string(raw[0:4]) != streamMagic {
		return StreamHeader{}, fmt.Errorf("invalid stream magic %q", string(raw[0:4]))
	}
	h := StreamHeader{
		Codec:        raw[4],
		Channels:     raw[5],
		SampleRateHz: binary.BigEndian.Uint32(raw[6:10]),
		FrameSamples: binary.BigEndian.Uint32(raw[10:14]),
		Sequence:     binary.BigEndian.Uint32(raw[14:18]),
		Timestamp:    binary.BigEndian.Uint64(raw[18:26]),
	}
	var payloadLenRaw [4]byte
	if _, err := io.ReadFull(r, payloadLenRaw[:]); err != nil {
		return StreamHeader{}, err
	}
	h.PayloadBytes = binary.BigEndian.Uint32(payloadLenRaw[:])
	return h, nil
}

func WritePCM32Packet(w io.Writer, channels int, sampleRateHz int, frameSamples int, sequence uint32, timestamp uint64, samples []int32) error {
	if channels < 1 || channels > 2 {
		return fmt.Errorf("channels must be 1 or 2")
	}
	if frameSamples < 0 {
		return fmt.Errorf("frameSamples must be >= 0")
	}
	if len(samples) != frameSamples*channels {
		return fmt.Errorf("sample length mismatch: got=%d want=%d", len(samples), frameSamples*channels)
	}

	payloadBytes := len(samples) * 4
	var hdr [30]byte
	copy(hdr[0:4], []byte(streamMagic))
	hdr[4] = StreamCodecPCM_S32LE
	hdr[5] = byte(channels)
	binary.BigEndian.PutUint32(hdr[6:10], uint32(sampleRateHz))
	binary.BigEndian.PutUint32(hdr[10:14], uint32(frameSamples))
	binary.BigEndian.PutUint32(hdr[14:18], sequence)
	binary.BigEndian.PutUint64(hdr[18:26], timestamp)
	binary.BigEndian.PutUint32(hdr[26:30], uint32(payloadBytes))
	if _, err := w.Write(hdr[:]); err != nil {
		return err
	}

	payload := make([]byte, payloadBytes)
	for i, s := range samples {
		binary.LittleEndian.PutUint32(payload[i*4:i*4+4], uint32(s))
	}
	_, err := w.Write(payload)
	return err
}
