package playback

import (
	"context"
	"sync"
	"time"
)

type VideoMediaStats struct {
	Available   bool
	Label       string
	Index       uint64
	Width       uint32
	Height      uint32
	Stride      uint32
	PayloadSize uint32
	DeclaredFPS float64
	ReceivedFPS float64
	FrameDT     time.Duration
	Invalid     uint64
}

type AudioMediaStats struct {
	Available     bool
	Label         string
	Index         uint64
	SampleRateHz  float64
	Channels      uint64
	SampleCount   uint64
	BatchDuration time.Duration
}

type MediaStatsSnapshot struct {
	Video VideoMediaStats
	Audio AudioMediaStats
}

type MediaStatsStore struct {
	mu       sync.RWMutex
	snapshot MediaStatsSnapshot
	now      func() time.Time

	lastVideoAt      time.Time
	videoWindowAt    time.Time
	videoWindowCount uint64
}

func NewMediaStatsStore() *MediaStatsStore {
	return &MediaStatsStore{now: time.Now}
}

func (s *MediaStatsStore) ObserveVideo(frame VideoFrame) {
	now := s.now()
	s.mu.Lock()
	defer s.mu.Unlock()

	stats := &s.snapshot.Video
	stats.Available = true
	stats.Label = frame.Label
	stats.Index = frame.Index
	stats.Width = frame.Width
	stats.Height = frame.Height
	stats.Stride = frame.Stride
	stats.PayloadSize = frame.Size
	if frame.FrameRateDenominator > 0 {
		stats.DeclaredFPS = float64(frame.FrameRateNumerator) /
			float64(frame.FrameRateDenominator)
	}
	if !s.lastVideoAt.IsZero() {
		stats.FrameDT = now.Sub(s.lastVideoAt)
	}
	s.lastVideoAt = now
	if frame.Invalid {
		stats.Invalid++
	}
	if s.videoWindowAt.IsZero() {
		s.videoWindowAt = now
	}
	s.videoWindowCount++
	if elapsed := now.Sub(s.videoWindowAt); elapsed >= time.Second {
		stats.ReceivedFPS = float64(s.videoWindowCount) / elapsed.Seconds()
		s.videoWindowAt = now
		s.videoWindowCount = 0
	}
}

func (s *MediaStatsStore) ObserveAudio(frame AudioFrame) {
	s.mu.Lock()
	defer s.mu.Unlock()

	stats := &s.snapshot.Audio
	stats.Available = true
	stats.Label = frame.Label
	stats.Index = frame.Index
	stats.Channels = frame.Channels
	stats.SampleCount = frame.SampleCount
	if frame.SampleRateDenominator > 0 {
		stats.SampleRateHz = float64(frame.SampleRateNumerator) /
			float64(frame.SampleRateDenominator)
	}
	if stats.SampleRateHz > 0 {
		stats.BatchDuration = time.Duration(
			float64(time.Second) * float64(frame.SampleCount) / stats.SampleRateHz,
		)
	}
}

func (s *MediaStatsStore) Snapshot() MediaStatsSnapshot {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.snapshot
}

type VideoStatsSink struct {
	Stats *MediaStatsStore
	Sink  VideoSink
}

func (s VideoStatsSink) ConsumeVideo(ctx context.Context, frame VideoFrame) error {
	if s.Stats != nil {
		s.Stats.ObserveVideo(frame)
	}
	return s.Sink.ConsumeVideo(ctx, frame)
}

type AudioStatsSink struct {
	Stats *MediaStatsStore
	Sink  AudioSink
}

func (s AudioStatsSink) ConsumeAudio(ctx context.Context, frame AudioFrame) error {
	if s.Stats != nil {
		s.Stats.ObserveAudio(frame)
	}
	return s.Sink.ConsumeAudio(ctx, frame)
}