From 4b857a6eedf41a43e3997a8dcc6a307f89b95e13 Mon Sep 17 00:00:00 2001
From: nareix <hi.nareix@gmail.com>
Date: Sun, 6 Mar 2016 16:10:40 +0800
Subject: [PATCH] add Muxer

---
 muxer.go | 166 +++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 166 insertions(+)
 create mode 100644 muxer.go

diff --git a/muxer.go b/muxer.go
new file mode 100644
index 0000000..aaa7017
--- /dev/null
+++ b/muxer.go
@@ -0,0 +1,166 @@
+
+package ts
+
+import (
+	"bytes"
+	"io"
+)
+
+type Track struct {
+	timeScale int64
+	writeSPS bool
+	SPS []byte
+	PPS []byte
+	spsHasWritten bool
+	tsw *TSWriter
+	PTS int64
+	PCR int64
+	pesBuf *bytes.Buffer
+}
+
+func (self *Track) WriteH264NALU(sync bool, duration int, nalu []byte) (err error) {
+	nalus := [][]byte{}
+
+	if !self.spsHasWritten {
+		nalus = append(nalus, self.SPS)
+		nalus = append(nalus, self.PPS)
+		self.spsHasWritten = true
+	}
+	nalus = append(nalus, nalu)
+
+	pes := PESHeader{
+		StreamId: StreamIdH264,
+		PTS: uint64(self.PTS)*PTS_HZ/uint64(self.timeScale),
+	}
+	if err = WritePESHeader(self.pesBuf, pes); err != nil {
+		return
+	}
+
+	data := &iovec{}
+	data.Append(self.pesBuf.Bytes())
+	for i, nalu := range nalus {
+		var startCode []byte
+		if i == 0 {
+			startCode = []byte{0,0,0,1,0x9,0xf0,0,0,0,1} // AUD
+		} else {
+			startCode = []byte{0,0,1}
+		}
+		data.Append(startCode)
+		data.Append(nalu)
+	}
+
+	self.tsw.RandomAccessIndicator = sync
+	self.tsw.PCR = uint64(self.PCR)*PCR_HZ/uint64(self.timeScale)
+
+	if err = self.tsw.WriteIovec(data); err != nil {
+		return
+	}
+
+	self.PTS += int64(duration)
+	self.PCR += int64(duration)
+	self.pesBuf.Reset()
+
+	return
+}
+
+func (self *Track) WriteADTSAACFrame(duration int, frame []byte) (err error) {
+	pes := PESHeader{
+		StreamId: StreamIdAAC,
+		PTS: uint64(self.PTS)*PTS_HZ/uint64(self.timeScale),
+	}
+	if err = WritePESHeader(self.pesBuf, pes); err != nil {
+		return
+	}
+
+	data := &iovec{}
+	data.Append(self.pesBuf.Bytes())
+	data.Append(frame)
+
+	self.tsw.RandomAccessIndicator = true
+	self.tsw.PCR = uint64(self.PCR)*PCR_HZ/uint64(self.timeScale)
+	if err = self.tsw.WriteIovec(data); err != nil {
+		return
+	}
+
+	self.PTS += int64(duration)
+	self.PCR += int64(duration)
+	self.pesBuf.Reset()
+
+	return
+}
+
+func newTrack(w io.Writer, pid uint, timeScale int64) (track *Track) {
+	track = &Track{
+		tsw: &TSWriter{
+			W: w,
+			PID: pid,
+			DiscontinuityIndicator: true,
+		},
+		timeScale: timeScale,
+		pesBuf: &bytes.Buffer{},
+	}
+	track.tsw.EnableVecWriter()
+	track.PTS = timeScale
+	track.PCR = timeScale
+	return
+}
+
+type Muxer struct {
+	W io.Writer
+	TimeScale int64
+	tswPAT *TSWriter
+	tswPMT *TSWriter
+	elemStreams []ElementaryStreamInfo
+}
+
+func (self *Muxer) AddAACTrack() (track *Track) {
+	self.elemStreams = append(
+		self.elemStreams,
+		ElementaryStreamInfo{StreamType: ElementaryStreamTypeAdtsAAC, ElementaryPID: 0x101},
+	)
+	return newTrack(self.W, 0x101, self.TimeScale)
+}
+
+func (self *Muxer) AddH264Track() (track *Track) {
+	self.elemStreams = append(
+		self.elemStreams,
+		ElementaryStreamInfo{StreamType: ElementaryStreamTypeH264, ElementaryPID: 0x100},
+	)
+	return newTrack(self.W, 0x100, self.TimeScale)
+}
+
+func (self *Muxer) WriteHeader() (err error) {
+	bufPAT := &bytes.Buffer{}
+	bufPMT := &bytes.Buffer{}
+
+	pat := PAT{
+		Entries: []PATEntry{
+			{ProgramNumber: 1, ProgramMapPID: 0x1000},
+		},
+	}
+	WritePAT(bufPAT, pat)
+	pmt := PMT{
+		PCRPID: 0x100,
+		ElementaryStreamInfos: self.elemStreams,
+	}
+	WritePMT(bufPMT, pmt)
+
+	tswPMT := &TSWriter{
+		W: self.W,
+		PID: 0x1000,
+		DiscontinuityIndicator: true,
+	}
+	tswPAT := &TSWriter{
+		W: self.W,
+		PID: 0,
+		DiscontinuityIndicator: true,
+	}
+	if err = tswPAT.Write(bufPAT.Bytes()); err != nil {
+		return
+	}
+	if err = tswPMT.Write(bufPMT.Bytes()); err != nil {
+		return
+	}
+	return
+}
+