rtsp_reader-writer/writer/internal/storage/segmenter.go

package storage

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"os"
	"time"

	"github.com/golang/snappy"
)

// writeString записывает строку: сначала длину (int32), затем байты строки.
func writeString(w io.Writer, s string) error {
	if err := binary.Write(w, binary.LittleEndian, int32(len(s))); err != nil {
		return err
	}
	_, err := w.Write([]byte(s))
	return err
}

// WritePacket записывает один interleaved пакет в writer.
func WritePacket(w io.Writer, pkt InterleavedPacket) error {
	// Записываем тип пакета (1 байт).
	if err := binary.Write(w, binary.LittleEndian, pkt.Type); err != nil {
		return err
	}
	// Записываем pts.
	if err := binary.Write(w, binary.LittleEndian, pkt.Pts); err != nil {
		return err
	}
	// В зависимости от типа пакета записываем данные access unit.
	if pkt.Type == PacketTypeH264 {
		// Для H264 AU — [][]byte.
		numAUs := int32(len(pkt.H264AUs))
		if err := binary.Write(w, binary.LittleEndian, numAUs); err != nil {
			return err
		}
		for _, au := range pkt.H264AUs {
			// Сначала длина AU.
			if err := binary.Write(w, binary.LittleEndian, int32(len(au))); err != nil {
				return err
			}
			// Затем сами данные.
			if _, err := w.Write(au); err != nil {
				return err
			}
		}
	} else if pkt.Type == PacketTypeLPCM {
		// Для G711 AU — []byte.
		if err := binary.Write(w, binary.LittleEndian, int32(len(pkt.LPCMSamples))); err != nil {
			return err
		}
		if _, err := w.Write(pkt.LPCMSamples); err != nil {
			return err
		}
	} else {
		return fmt.Errorf("неизвестный тип пакета: %d", pkt.Type)
	}
	return nil
}

// WriteHeader writes header to a file.
func WriteHeader(w io.Writer, seg Segment) error {
	var buf bytes.Buffer

	// Записываем строки Start и Duration.
	if err := writeString(&buf, seg.Start); err != nil {
		return err
	}
	if err := writeString(&buf, seg.Duration); err != nil {
		return err
	}

	// Сначала записываем длину сегмента, затем данные сегмента.
	segData := buf.Bytes()
	if err := binary.Write(w, binary.LittleEndian, int32(len(segData))); err != nil {
		return err
	}
	_, err := w.Write(segData)
	return err
}

// WriteInterleavedPacket записывает один сегмент в writer. Сначала собирается содержимое сегмента в буфер,
// затем записывается его длина (int32) и данные.
func WriteInterleavedPacket(w io.Writer, seg Segment) error {
	var buf bytes.Buffer

	// Записываем количество пакетов.
	if err := binary.Write(&buf, binary.LittleEndian, int32(1)); err != nil {
		return err
	}

	// Для каждого пакета записываем данные.
	if err := WritePacket(&buf, seg.Packets); err != nil {
		return err
	}

	// Сначала записываем длину сегмента, затем данные сегмента.
	segData := buf.Bytes()
	if err := binary.Write(w, binary.LittleEndian, int32(len(segData))); err != nil {
		return err
	}
	_, err := w.Write(segData)
	return err
}

func main() {
	now := time.Now()
	// Пример сегмента с interleaved пакетами.
	seg := Segment{
		Start:    now.Format(time.RFC3339),
		Duration: "1m", // длительность сегмента в виде строки
		Packets: InterleavedPacket{

			Type: PacketTypeH264,
			Pts:  now.UnixNano(),
			H264AUs: [][]byte{
				[]byte{0x00, 0x01, 0x02},
				[]byte{0x03, 0x04},
			},
		},
	}

	// Открываем файл для записи.
	f, err := os.Create("stream_interleaved.bin")
	if err != nil {
		fmt.Println("Ошибка создания файла:", err)
		return
	}
	defer f.Close()

	// Оборачиваем writer через snappy для быстрой компрессии (если требуется).
	writer := snappy.NewBufferedWriter(f)
	defer writer.Close()

	// Записываем заголовок файла (например, streamID).
	streamID := "example_stream"
	if err := binary.Write(writer, binary.LittleEndian, int32(len(streamID))); err != nil {
		fmt.Println("Ошибка записи заголовка:", err)
		return
	}
	if _, err := writer.Write([]byte(streamID)); err != nil {
		fmt.Println("Ошибка записи streamID:", err)
		return
	}

	// Записываем сегмент с interleaved пакетами.
	if err := WriteInterleavedPacket(writer, seg); err != nil {
		fmt.Println("Ошибка записи сегмента:", err)
		return
	}

	// Обязательно делаем Flush, чтобы данные точно записались.
	if err := writer.Flush(); err != nil {
		fmt.Println("Ошибка при сбросе данных:", err)
		return
	}

	fmt.Println("Сегмент с interleaved пакетами успешно записан.")
}