package formats

import (
	"bufio"
	"github.com/bluenviron/mediacommon/pkg/codecs/mpeg4audio"
	"os"
	"sync"

	"github.com/bluenviron/gortsplib/v4/pkg/format"
	"github.com/bluenviron/mediacommon/pkg/codecs/h264"
	"github.com/bluenviron/mediacommon/pkg/formats/mpegts"
)

func multiplyAndDivide(v, m, d int64) int64 {
	secs := v / d
	dec := v % d
	return (secs*m + dec*m/d)
}

// MpegtsMuxer allows to save a H264 / MPEG-4 audio stream into a MPEG-TS file.
type MpegtsMuxer struct {
	FileName         string
	H264Format       *format.H264
	Mpeg4AudioFormat *format.MPEG4Audio

	f               *os.File
	b               *bufio.Writer
	w               *mpegts.Writer
	h264Track       *mpegts.Track
	mpeg4AudioTrack *mpegts.Track
	dtsExtractor    *h264.DTSExtractor2
	mutex           sync.Mutex
}

// Initialize initializes a MpegtsMuxer.
func (e *MpegtsMuxer) Initialize() error {
	var err error
	e.f, err = os.Create(e.FileName)
	if err != nil {
		return err
	}
	e.b = bufio.NewWriter(e.f)

	e.h264Track = &mpegts.Track{
		Codec: &mpegts.CodecH264{},
	}

	e.mpeg4AudioTrack = &mpegts.Track{
		Codec: &mpegts.CodecMPEG4Audio{
			Config: mpeg4audio.Config{},
		},
	}

	e.w = mpegts.NewWriter(e.b, []*mpegts.Track{e.h264Track, e.mpeg4AudioTrack})

	return nil
}

// Close closes all the MpegtsMuxer resources.
func (e *MpegtsMuxer) Close() {
	e.b.Flush()
	e.f.Close()
}

// WriteH264 writes a H264 access unit into MPEG-TS.
func (e *MpegtsMuxer) WriteH264(pts int64, au [][]byte) error {
	e.mutex.Lock()
	defer e.mutex.Unlock()

	var filteredAU [][]byte

	nonIDRPresent := false
	idrPresent := false

	for _, nalu := range au {
		typ := h264.NALUType(nalu[0] & 0x1F)
		switch typ {
		case h264.NALUTypeSPS:
			e.H264Format.SPS = nalu
			continue

		case h264.NALUTypePPS:
			e.H264Format.PPS = nalu
			continue

		case h264.NALUTypeAccessUnitDelimiter:
			continue

		case h264.NALUTypeIDR:
			idrPresent = true

		case h264.NALUTypeNonIDR:
			nonIDRPresent = true
		}

		filteredAU = append(filteredAU, nalu)
	}

	au = filteredAU

	if au == nil || (!nonIDRPresent && !idrPresent) {
		return nil
	}

	// Add SPS and PPS before access unit that contains an IDR.
	if idrPresent {
		au = append([][]byte{e.H264Format.SPS, e.H264Format.PPS}, au...)
	}

	if e.dtsExtractor == nil {
		// Skip samples silently until we find one with an IDR.
		if !idrPresent {
			return nil
		}
		e.dtsExtractor = h264.NewDTSExtractor2()
	}

	dts, err := e.dtsExtractor.Extract(au, pts)
	if err != nil {
		return err
	}

	// Encode into MPEG-TS.
	return e.w.WriteH2642(e.h264Track, pts, dts, au)
}

// WriteAAC writes MPEG-4 audio access units into MPEG-TS.
func (e *MpegtsMuxer) WriteAAC(aus [][]byte, pts int64) error {
	e.mutex.Lock()
	defer e.mutex.Unlock()

	return e.w.WriteMPEG4Audio(e.mpeg4AudioTrack, multiplyAndDivide(pts, 90000, int64(8000)), aus)
}