From 128f96f8fd2446ca8e53c9f95aa6486ced541862 Mon Sep 17 00:00:00 2001
From: Quentin Renard <contact-github@asticode.com>
Date: Thu, 14 Nov 2024 18:00:40 +0100
Subject: [PATCH] Added custom io muxing example

---
 examples/custom_io_muxing/main.go | 185 ++++++++++++++++++++++++++++++
 1 file changed, 185 insertions(+)
 create mode 100644 examples/custom_io_muxing/main.go

diff --git a/examples/custom_io_muxing/main.go b/examples/custom_io_muxing/main.go
new file mode 100644
index 0000000..4f2dce3
--- /dev/null
+++ b/examples/custom_io_muxing/main.go
@@ -0,0 +1,185 @@
+package main
+
+import (
+	"errors"
+	"flag"
+	"fmt"
+	"log"
+	"os"
+	"strings"
+
+	"github.com/asticode/go-astiav"
+)
+
+var (
+	input  = flag.String("i", "", "the input path")
+	output = flag.String("o", "", "the output path")
+)
+
+func main() {
+	// Handle ffmpeg logs
+	astiav.SetLogLevel(astiav.LogLevelDebug)
+	astiav.SetLogCallback(func(c astiav.Classer, l astiav.LogLevel, fmt, msg string) {
+		var cs string
+		if c != nil {
+			if cl := c.Class(); cl != nil {
+				cs = " - class: " + cl.String()
+			}
+		}
+		log.Printf("ffmpeg log: %s%s - level: %d\n", strings.TrimSpace(msg), cs, l)
+	})
+
+	// Parse flags
+	flag.Parse()
+
+	// Usage
+	if *input == "" || *output == "" {
+		log.Println("Usage: <binary path> -i <input path> -o <output path>")
+		return
+	}
+
+	// Allocate packet
+	pkt := astiav.AllocPacket()
+	defer pkt.Free()
+
+	// Allocate input format context
+	inputFormatContext := astiav.AllocFormatContext()
+	if inputFormatContext == nil {
+		log.Fatal(errors.New("main: input format context is nil"))
+	}
+	defer inputFormatContext.Free()
+
+	// Open input
+	if err := inputFormatContext.OpenInput(*input, nil, nil); err != nil {
+		log.Fatal(fmt.Errorf("main: opening input failed: %w", err))
+	}
+	defer inputFormatContext.CloseInput()
+
+	// Find stream info
+	if err := inputFormatContext.FindStreamInfo(nil); err != nil {
+		log.Fatal(fmt.Errorf("main: finding stream info failed: %w", err))
+	}
+
+	// Allocate output format context
+	outputFormatContext, err := astiav.AllocOutputFormatContext(nil, "mp4", "")
+	if err != nil {
+		log.Fatal(fmt.Errorf("main: allocating output format context failed: %w", err))
+	}
+	if outputFormatContext == nil {
+		log.Fatal(errors.New("main: output format context is nil"))
+	}
+	defer outputFormatContext.Free()
+
+	// Open file
+	f, err := os.Create(*output)
+	if err != nil {
+		log.Fatal(fmt.Errorf("main: creating %s failed: %w", *output, err))
+	}
+	defer f.Close()
+
+	// Allocate io context
+	ioContext, err := astiav.AllocIOContext(
+		4096,
+		true,
+		nil,
+		func(offset int64, whence int) (n int64, err error) {
+			return f.Seek(offset, whence)
+		},
+		func(b []byte) (n int, err error) {
+			return f.Write(b)
+		},
+	)
+	if err != nil {
+		log.Fatal(fmt.Errorf("main: allocating io context failed: %w", err))
+	}
+	defer ioContext.Free()
+
+	// Store io context
+	outputFormatContext.SetPb(ioContext)
+
+	// Loop through streams
+	inputStreams := make(map[int]*astiav.Stream)  // Indexed by input stream index
+	outputStreams := make(map[int]*astiav.Stream) // Indexed by input stream index
+	for _, is := range inputFormatContext.Streams() {
+		// Only process audio or video
+		if is.CodecParameters().MediaType() != astiav.MediaTypeAudio &&
+			is.CodecParameters().MediaType() != astiav.MediaTypeVideo {
+			continue
+		}
+
+		// Add input stream
+		inputStreams[is.Index()] = is
+
+		// Add stream to output format context
+		os := outputFormatContext.NewStream(nil)
+		if os == nil {
+			log.Fatal(errors.New("main: output stream is nil"))
+		}
+
+		// Copy codec parameters
+		if err = is.CodecParameters().Copy(os.CodecParameters()); err != nil {
+			log.Fatal(fmt.Errorf("main: copying codec parameters failed: %w", err))
+		}
+
+		// Reset codec tag
+		os.CodecParameters().SetCodecTag(0)
+
+		// Add output stream
+		outputStreams[is.Index()] = os
+	}
+
+	// Write header
+	if err = outputFormatContext.WriteHeader(nil); err != nil {
+		log.Fatal(fmt.Errorf("main: writing header failed: %w", err))
+	}
+
+	// Loop through packets
+	for {
+		// We use a closure to ease unreferencing packet
+		if stop := func() bool {
+			// Read frame
+			if err = inputFormatContext.ReadFrame(pkt); err != nil {
+				if errors.Is(err, astiav.ErrEof) {
+					return true
+				}
+				log.Fatal(fmt.Errorf("main: reading frame failed: %w", err))
+			}
+
+			// Make sure to unreference packet
+			defer pkt.Unref()
+
+			// Get input stream
+			inputStream, ok := inputStreams[pkt.StreamIndex()]
+			if !ok {
+				return false
+			}
+
+			// Get output stream
+			outputStream, ok := outputStreams[pkt.StreamIndex()]
+			if !ok {
+				return false
+			}
+
+			// Update packet
+			pkt.SetStreamIndex(outputStream.Index())
+			pkt.RescaleTs(inputStream.TimeBase(), outputStream.TimeBase())
+			pkt.SetPos(-1)
+
+			// Write frame
+			if err = outputFormatContext.WriteInterleavedFrame(pkt); err != nil {
+				log.Fatal(fmt.Errorf("main: writing interleaved frame failed: %w", err))
+			}
+			return false
+		}(); stop {
+			break
+		}
+	}
+
+	// Write trailer
+	if err = outputFormatContext.WriteTrailer(); err != nil {
+		log.Fatal(fmt.Errorf("main: writing trailer failed: %w", err))
+	}
+
+	// Success
+	log.Println("success")
+}