TrafficSpeed/internal/project/video_iterator.go

package project

import (
	"errors"
	"fmt"
	"image"
	"log"
	"time"

	"github.com/asticode/go-astiav"
)

func init() {
	astiav.SetLogLevel(astiav.LogLevelError)
}

type Iterator struct {
	err            error
	formatContext  *astiav.FormatContext
	codecContext   *astiav.CodecContext
	videoStreamIdx int
	rect           image.Rectangle
	packet         *astiav.Packet
	frame          int
	decoded        bool
	currentFrame   *astiav.Frame
	img            image.Image
}

func (p *Iterator) Close() {
	if p.formatContext != nil {
		p.formatContext.Free()
		p.formatContext = nil
	}
	if p.codecContext != nil {
		p.codecContext.Free()
		p.codecContext = nil
	}
	if p.currentFrame != nil {
		p.currentFrame.Free()
		p.currentFrame = nil
	}
	if p.packet != nil {
		p.packet.Free()
		p.packet = nil
	}
}

func NewIterator(filename string) (iter *Iterator, err error) {
	if filename == "" {
		panic("missing filename")
	}

	iter = &Iterator{frame: -1}

	// Open input file
	formatContext := astiav.AllocFormatContext()
	if formatContext == nil {
		return nil, fmt.Errorf("allocating format context failed")
	}
	iter.formatContext = formatContext

	// Get stream info
	if err := formatContext.OpenInput(filename, nil, nil); err != nil {
		iter.Close()
		return nil, fmt.Errorf("opening input: %w", err)
	}

	// Find the first video stream
	iter.videoStreamIdx = -1
	for i, stream := range formatContext.Streams() {
		if stream.CodecParameters().MediaType() != astiav.MediaTypeVideo {
			continue
		}
		iter.videoStreamIdx = i
		width := stream.CodecParameters().Width()
		height := stream.CodecParameters().Height()
		iter.rect = image.Rect(0, 0, width, height)
		fmt.Printf("stream[%d] = video (%dx%d)\n", i, width, height)

		iter.currentFrame = astiav.AllocFrame()
		// iter.currentFrame.SetWidth(width)
		// iter.currentFrame.SetHeight(height)
		// iter.currentFrame.SetPixelFormat(astiav.PixelFormatYuv420P)
		break
	}

	if iter.videoStreamIdx == -1 {
		iter.Close()
		return nil, fmt.Errorf("no video stream found")
	}

	// Find decoder
	stream := formatContext.Streams()[iter.videoStreamIdx]
	codec := astiav.FindDecoder(stream.CodecParameters().CodecID())
	if codec == nil {
		iter.Close()
		return nil, fmt.Errorf("no codec found for stream %d", iter.videoStreamIdx)
	}
	log.Printf("codec found: %s", codec.Name())

	// Create codec context
	iter.codecContext = astiav.AllocCodecContext(codec)
	if iter.codecContext == nil {
		iter.Close()
		return nil, fmt.Errorf("failed to create codec context")
	}
	log.Printf("codec parameters for stream %d: %#v", iter.videoStreamIdx, stream.CodecParameters())

	// Copy codec parameters
	if err := iter.codecContext.FromCodecParameters(stream.CodecParameters()); err != nil {
		iter.Close()
		return nil, fmt.Errorf("copying codec parameters: %w", err)
	}
	// iter.codecContext.SetPixelFormat(astiav.PixelFormatYuv420P)

	// Open codec
	if err := iter.codecContext.Open(codec, nil); err != nil {
		iter.Close()
		return nil, fmt.Errorf("opening codec: %w", err)
	}

	// Allocate packet
	iter.packet = astiav.AllocPacket()

	log.Printf("codecContext: %s pixelFormat=%#v", iter.codecContext.String(), iter.codecContext.PixelFormat())

	return iter, nil
}

func (i *Iterator) Seek(d time.Duration) error {
	// Convert time.Duration to AV timestamp
	ts := int64(d / time.Microsecond)
	timeBase := i.formatContext.Streams()[i.videoStreamIdx].TimeBase()
	timestamp := astiav.RescaleQ(ts, astiav.NewRational(1, 1000000), timeBase)

	log.Printf("should seek %s (timestamp: %d)", d, timestamp)

	// Seek to timestamp
	seekFlags := astiav.NewSeekFlags(astiav.SeekFlagBackward, astiav.SeekFlagAny)
	if err := i.formatContext.SeekFrame(i.videoStreamIdx, timestamp, seekFlags); err != nil {
		return fmt.Errorf("seeking: %w", err)
	}

	// Reset frame counter
	i.frame = -1

	// Send empty packet to flush the decoder
	if err := i.codecContext.SendPacket(nil); err != nil {
		return fmt.Errorf("flushing decoder: %w", err)
	}

	return nil
}

func (i *Iterator) VideoResolution() string {
	return fmt.Sprintf("%dx%d", i.rect.Dx(), i.rect.Dy())
}

// func (i *Iterator) NextWithImage() bool {
// 	for i.Next() {
// 		i.err = i.DecodeFrame()
// 		if i.err != nil {
// 			return false
// 		}
// 		if i.currentFrame == nil {
// 			continue
// 		}
// 		return true
// 	}
// 	return false
// }

func (i *Iterator) Next() bool {
	i.decoded = false

	for {
		// Read frame
		err := i.formatContext.ReadFrame(i.packet)
		if err != nil {
			if errors.Is(err, astiav.ErrEof) {
				return false
			}
			i.err = fmt.Errorf("reading frame: %w", err)
			return false
		}

		// Skip if not a video packet
		if i.packet.StreamIndex() != i.videoStreamIdx {
			log.Printf("Skipping packet from stream %d (not video)", i.packet.StreamIndex())
			i.packet.Unref()
			continue
		}
		i.frame++

		// Send packet to decoder
		if i.err = i.codecContext.SendPacket(i.packet); i.err != nil {
			log.Printf("SendPacket() error: %v", i.err)
			i.packet.Unref()
			return false
		}

		// Get frame from decoder
		i.err = i.codecContext.ReceiveFrame(i.currentFrame)
		if errors.Is(i.err, astiav.ErrEagain) {
			i.packet.Unref()
			log.Printf("ReceiveFrame() Eagain for frame %d, waiting for more data", i.frame)
			i.err = nil // No frame available yet, continue to read more packets
			continue
		}
		defer i.packet.Unref()
		if i.err != nil {
			// Check if we need more data or reached EOF
			log.Printf("ReceiveFrame() error: %v", i.err)
			return false
		}

		defer i.currentFrame.Unref()

		img, _ := i.currentFrame.Data().GuessImageFormat()
		i.err = i.currentFrame.Data().ToImage(img)
		if i.err != nil {
			log.Printf("ToImage() error: %v", i.err)
			return false
		}
		i.img = img
		return true
	}
}

func (i *Iterator) Frame() int {
	return i.frame
}

func (i *Iterator) Image() image.Image {
	return i.img
}

func (i *Iterator) Error() error {
	return i.err
}

func (i *Iterator) Duration() time.Duration {
	if i.packet == nil || i.packet.Pts() == astiav.NoPtsValue {
		return 0
	}

	timeBase := i.formatContext.Streams()[i.videoStreamIdx].TimeBase()
	durationMicros := astiav.RescaleQ(i.packet.Pts(), timeBase, astiav.NewRational(1, 1000000))
	return time.Duration(durationMicros) * time.Microsecond
}

func (i *Iterator) DurationMs() time.Duration {
	return i.Duration().Round(time.Millisecond)
}

func (i *Iterator) IsKeyFrame() bool {
	return i.currentFrame.KeyFrame()
}