goffmpeg/pkg/ffmpeg/ffmpeg.go

package ffmpeg

/*
#include <libavformat/avformat.h>
#include <libavformat/avio.h>
*/
import "C"
import "unsafe"

// FormatContext represents the input/output format context
type FormatContext struct {
	ptr *C.AVFormatContext
}

// AllocFormatContext allocates a format context
func AllocFormatContext() *FormatContext {
	return &FormatContext{
		ptr: C.avformat_alloc_context(),
	}
}

// Free frees the format context
func (fc *FormatContext) Free() {
	if fc.ptr != nil {
		C.avformat_close_input(&fc.ptr)
		fc.ptr = nil
	}
}

// FormatContextFromC converts a C pointer to FormatContext
func FormatContextFromC(ptr *C.AVFormatContext) *FormatContext {
	return &FormatContext{ptr: ptr}
}

// CPtr returns the underlying C pointer
func (fc *FormatContext) CPtr() *C.AVFormatContext {
	return fc.ptr
}

// OpenInput opens an input file
func (fc *FormatContext) OpenInput(url string) error {
	if fc.ptr == nil {
		fc.ptr = C.avformat_alloc_context()
		if fc.ptr == nil {
			return ErrInvalidInput
		}
	}

	cURL := C.CString(url)
	defer C.free(unsafe.Pointer(cURL))

	ret := C.avformat_open_input(&fc.ptr, cURL, nil, nil)
	if ret < 0 {
		return &FFmpegError{
			Code:    int(ret),
			Message: "failed to open input",
			Op:      "OpenInput",
		}
	}
	return nil
}

// Close closes the input
func (fc *FormatContext) Close() {
	if fc.ptr != nil {
		C.avformat_close_input(&fc.ptr)
	}
}

// FindStreamInfo finds stream info
func (fc *FormatContext) FindStreamInfo() error {
	if fc.ptr == nil {
		return ErrInvalidInput
	}

	ret := C.avformat_find_stream_info(fc.ptr, nil)
	if ret < 0 {
		return &FFmpegError{
			Code:    int(ret),
			Message: "failed to find stream info",
			Op:      "FindStreamInfo",
		}
	}
	return nil
}

// NbStreams returns the number of streams
func (fc *FormatContext) NbStreams() int {
	if fc.ptr == nil {
		return 0
	}
	return int(fc.ptr.nb_streams)
}

// Streams returns the streams
func (fc *FormatContext) Streams() []*Stream {
	if fc.ptr == nil {
		return nil
	}

	streams := make([]*Stream, fc.NbStreams())
	ptrSize := unsafe.Sizeof(fc.ptr.streams)
	for i := 0; i < fc.NbStreams(); i++ {
		streamPtr := (**C.AVStream)(unsafe.Pointer(uintptr(unsafe.Pointer(fc.ptr.streams)) + uintptr(i)*ptrSize))
		streams[i] = StreamFromC(*streamPtr)
	}
	return streams
}

// VideoStreams returns only video streams
func (fc *FormatContext) VideoStreams() []*Stream {
	streams := fc.Streams()
	videoStreams := make([]*Stream, 0)
	for _, s := range streams {
		if s.Type() == CodecTypeVideo {
			videoStreams = append(videoStreams, s)
		}
	}
	return videoStreams
}

// AudioStreams returns only audio streams
func (fc *FormatContext) AudioStreams() []*Stream {
	streams := fc.Streams()
	audioStreams := make([]*Stream, 0)
	for _, s := range streams {
		if s.Type() == CodecTypeAudio {
			audioStreams = append(audioStreams, s)
		}
	}
	return audioStreams
}

// ReadPacket reads a packet
func (fc *FormatContext) ReadPacket(pkt *Packet) error {
	if fc.ptr == nil {
		return ErrInvalidInput
	}

	ret := C.av_read_frame(fc.ptr, pkt.CPtr())
	if ret < 0 {
		return &FFmpegError{
			Code:    int(ret),
			Message: "failed to read packet",
			Op:      "ReadPacket",
		}
	}
	return nil
}

// WritePacket writes a packet
func (fc *FormatContext) WritePacket(pkt *Packet) error {
	if fc.ptr == nil {
		return ErrInvalidOutput
	}

	ret := C.av_interleaved_write_frame(fc.ptr, pkt.CPtr())
	if ret < 0 {
		return &FFmpegError{
			Code:    int(ret),
			Message: "failed to write packet",
			Op:      "WritePacket",
		}
	}
	return nil
}

// DumpFormat dumps format info
func (fc *FormatContext) DumpFormat(idx int, url string, isOutput bool) {
	if fc.ptr == nil {
		return
	}
	cURL := C.CString(url)
	defer C.free(unsafe.Pointer(cURL))
	C.av_dump_format(fc.ptr, C.int(idx), cURL, C.int(boolToInt(isOutput)))
}

func boolToInt(b bool) int {
	if b {
		return 1
	}
	return 0
}

// Stream represents a media stream
type Stream struct {
	ptr *C.AVStream
}

// StreamFromC converts a C pointer to Stream
func StreamFromC(ptr *C.AVStream) *Stream {
	return &Stream{ptr: ptr}
}

// CPtr returns the underlying C pointer
func (s *Stream) CPtr() *C.AVStream {
	return s.ptr
}

// Index returns the stream index
func (s *Stream) Index() int {
	if s.ptr == nil {
		return -1
	}
	return int(s.ptr.index)
}

// Type returns the codec type
func (s *Stream) Type() CodecType {
	if s.ptr == nil {
		return CodecType(0)
	}
	return CodecType(s.ptr.codecpar.codec_type)
}

// CodecParameters returns the codec parameters
func (s *Stream) CodecParameters() *CodecParameters {
	if s.ptr == nil {
		return nil
	}
	return CodecParametersFromC(s.ptr.codecpar)
}

// SetCodecParameters sets the codec parameters
func (s *Stream) SetCodecParameters(cp *CodecParameters) error {
	if s.ptr == nil || cp == nil || cp.ptr == nil {
		return ErrInvalidCodec
	}

	ret := C.avcodec_parameters_copy(s.ptr.codecpar, cp.ptr)
	if ret < 0 {
		return &FFmpegError{
			Code:    int(ret),
			Message: "failed to copy codec parameters",
			Op:      "SetCodecParameters",
		}
	}
	return nil
}

// SetCodecContextParameters copies parameters from a codec context to this stream
func (s *Stream) SetCodecContextParameters(cc *Context) error {
	if s.ptr == nil || cc == nil || cc.ptr == nil {
		return ErrInvalidCodec
	}

	// Manually copy fields from codec context to codec parameters
	s.ptr.codecpar.codec_type = cc.ptr.codec_type
	s.ptr.codecpar.codec_id = cc.ptr.codec_id
	s.ptr.codecpar.bit_rate = cc.ptr.bit_rate
	s.ptr.codecpar.width = cc.ptr.width
	s.ptr.codecpar.height = cc.ptr.height

	// Copy format based on codec type
	if cc.ptr.codec_type == C.AVMEDIA_TYPE_VIDEO {
		s.ptr.codecpar.format = C.int(cc.ptr.pix_fmt)
	} else if cc.ptr.codec_type == C.AVMEDIA_TYPE_AUDIO {
		s.ptr.codecpar.format = C.int(cc.ptr.sample_fmt)
	}

	s.ptr.codecpar.sample_rate = cc.ptr.sample_rate
	s.ptr.codecpar.channels = cc.ptr.channels
	s.ptr.codecpar.channel_layout = cc.ptr.channel_layout

	return nil
}

// Codec returns the codec context (deprecated: use CodecParameters instead)
func (s *Stream) Codec() *Context {
	// In FFmpeg 4.0+, codec field was removed from AVStream
	// Use CodecParameters() and allocate a new context if needed
	return nil
}

// TimeBase returns the time base
func (s *Stream) TimeBase() Rational {
	if s.ptr == nil {
		return Rational{}
	}
	return Rational{
		num: int(s.ptr.time_base.num),
		den: int(s.ptr.time_base.den),
	}
}

// SetTimeBase sets the time base
func (s *Stream) SetTimeBase(r Rational) {
	if s.ptr != nil {
		s.ptr.time_base.num = C.int(r.num)
		s.ptr.time_base.den = C.int(r.den)
	}
}

// Rational represents a rational number
type Rational struct {
	num int
	den int
}

// NewRational creates a new rational
func NewRational(num, den int) Rational {
	return Rational{num: num, den: den}
}

// CodecParameters represents codec parameters
type CodecParameters struct {
	ptr *C.AVCodecParameters
}

// CodecParametersFromC converts a C pointer to CodecParameters
func CodecParametersFromC(ptr *C.AVCodecParameters) *CodecParameters {
	return &CodecParameters{ptr: ptr}
}

// CPtr returns the underlying C pointer
func (cp *CodecParameters) CPtr() *C.AVCodecParameters {
	return cp.ptr
}

// CodecType returns the codec type
func (cp *CodecParameters) CodecType() CodecType {
	if cp.ptr == nil {
		return CodecType(0)
	}
	return CodecType(cp.ptr.codec_type)
}

// CodecID returns the codec ID
func (cp *CodecParameters) CodecID() int {
	if cp.ptr == nil {
		return 0
	}
	return int(cp.ptr.codec_id)
}

// Width returns the width
func (cp *CodecParameters) Width() int {
	if cp.ptr == nil {
		return 0
	}
	return int(cp.ptr.width)
}

// Height returns the height
func (cp *CodecParameters) Height() int {
	if cp.ptr == nil {
		return 0
	}
	return int(cp.ptr.height)
}

// Format returns the format
func (cp *CodecParameters) Format() int {
	if cp.ptr == nil {
		return -1
	}
	return int(cp.ptr.format)
}

// SampleRate returns the sample rate
func (cp *CodecParameters) SampleRate() int {
	if cp.ptr == nil {
		return 0
	}
	return int(cp.ptr.sample_rate)
}

// Channels returns the number of channels
func (cp *CodecParameters) Channels() int {
	if cp.ptr == nil {
		return 0
	}
	return int(cp.ptr.channels)
}

// FrameSize returns the frame size
func (cp *CodecParameters) FrameSize() int {
	if cp.ptr == nil {
		return 0
	}
	return int(cp.ptr.frame_size)
}

// BitRate returns the bit rate
func (cp *CodecParameters) BitRate() int64 {
	if cp.ptr == nil {
		return 0
	}
	return int64(cp.ptr.bit_rate)
}

// OutputFormatContext represents an output format context
type OutputFormatContext struct {
	FormatContext
}

// AllocOutputContext allocates an output format context
func AllocOutputContext(url string, fmt *OutputFormat) (*OutputFormatContext, error) {
	ofc := &OutputFormatContext{
		FormatContext: FormatContext{
			ptr: C.avformat_alloc_context(),
		},
	}

	if ofc.ptr == nil {
		return nil, ErrInvalidOutput
	}

	cURL := C.CString(url)
	defer C.free(unsafe.Pointer(cURL))

	var ret C.int
	if fmt != nil && fmt.ptr != nil {
		ret = C.avformat_alloc_output_context2(&ofc.ptr, fmt.ptr, nil, cURL)
	} else {
		ret = C.avformat_alloc_output_context2(&ofc.ptr, nil, nil, cURL)
	}

	if ret < 0 {
		C.avformat_free_context(ofc.ptr)
		return nil, &FFmpegError{
			Code:    int(ret),
			Message: "failed to allocate output context",
			Op:      "AllocOutputContext",
		}
	}

	return ofc, nil
}

// AddStream adds a new stream with optional codec
func (ofc *OutputFormatContext) AddStream(codec *Codec) (*Stream, error) {
	if ofc.ptr == nil {
		return nil, ErrInvalidOutput
	}

	var stream *C.AVStream
	if codec != nil && codec.ptr != nil {
		stream = C.avformat_new_stream(ofc.ptr, codec.ptr)
	} else {
		stream = C.avformat_new_stream(ofc.ptr, nil)
	}
	if stream == nil {
		return nil, ErrInvalidCodec
	}

	return StreamFromC(stream), nil
}

// SetOformat sets the output format
func (ofc *OutputFormatContext) SetOformat(fmt *OutputFormat) {
	if ofc.ptr != nil && fmt != nil && fmt.ptr != nil {
		ofc.ptr.oformat = fmt.ptr
	}
}

// OpenOutput opens the output file
func (ofc *OutputFormatContext) OpenOutput(url string) error {
	if ofc.ptr == nil {
		return ErrInvalidOutput
	}

	cURL := C.CString(url)
	defer C.free(unsafe.Pointer(cURL))

	ret := C.avio_open(&ofc.ptr.pb, cURL, C.AVIO_FLAG_WRITE)
	if ret < 0 {
		return &FFmpegError{
			Code:    int(ret),
			Message: "failed to open output",
			Op:      "OpenOutput",
		}
	}
	return nil
}

// CloseOutput closes the output file
func (ofc *OutputFormatContext) CloseOutput() {
	if ofc.ptr != nil && ofc.ptr.pb != nil {
		C.avio_close(ofc.ptr.pb)
	}
}

// DumpFormat dumps format info
func (ofc *OutputFormatContext) DumpFormat(idx int, url string, isOutput bool) {
	ofc.FormatContext.DumpFormat(idx, url, isOutput)
}

// WriteHeader writes the header
func (ofc *OutputFormatContext) WriteHeader() error {
	if ofc.ptr == nil {
		return ErrInvalidOutput
	}

	ret := C.avformat_write_header(ofc.ptr, nil)
	if ret < 0 {
		return &FFmpegError{
			Code:    int(ret),
			Message: "failed to write header",
			Op:      "WriteHeader",
		}
	}
	return nil
}

// WriteTrailer writes the trailer
func (ofc *OutputFormatContext) WriteTrailer() error {
	if ofc.ptr == nil {
		return ErrInvalidOutput
	}

	ret := C.av_write_trailer(ofc.ptr)
	if ret < 0 {
		return &FFmpegError{
			Code:    int(ret),
			Message: "failed to write trailer",
			Op:      "WriteTrailer",
		}
	}
	return nil
}

// OutputFormat represents an output format
type OutputFormat struct {
	ptr *C.AVOutputFormat
}

// OutputFormatFromC converts a C pointer to OutputFormat
func OutputFormatFromC(ptr *C.AVOutputFormat) *OutputFormat {
	return &OutputFormat{ptr: ptr}
}

// CPtr returns the underlying C pointer
func (of *OutputFormat) CPtr() *C.AVOutputFormat {
	return of.ptr
}

// GuessFormat guesses the output format
func GuessFormat(shortName, filename string) *OutputFormat {
	cShortName := C.CString(shortName)
	cFilename := C.CString(filename)
	defer C.free(unsafe.Pointer(cShortName))
	defer C.free(unsafe.Pointer(cFilename))

	ptr := C.av_guess_format(cShortName, cFilename, nil)
	if ptr == nil {
		return nil
	}
	return OutputFormatFromC(ptr)
}