package audio

import (
	"errors"
	"sync"

	"github.com/pion/interceptor/pkg/jitterbuffer"
	"github.com/pion/rtp"
)

// ceiling on buffered packets before Push force-trims back to the target.
// 20 packets is ~400ms at 20ms framing - past this we're just adding delay.
const DefaultMaxDepth = 20

// pion's jitterbuffer + late-drop on Push + skip-ahead on persistent gap;
// PopOrSkip reports gap counts so the caller can run PLC over them.
//
// TargetDepth/MaxDepth bound playout latency: Push hard-trims to TargetDepth
// once depth exceeds MaxDepth; ShedOne walks depth back down one packet per
// pop. Dropped packets skip the audio forward briefly.
type Playout struct {
	mu      sync.Mutex
	jb      *jitterbuffer.JitterBuffer
	started bool   // first Pop succeeded -> headSeq is meaningful
	headSeq uint16 // expected next seq
	// caps the scan in PopOrSkip's slow path
	MaxSkip uint16
	// trim target; also the floor ShedOne won't drop below. Defaults to prime.
	TargetDepth int
	// depth ceiling; Push trims to TargetDepth once exceeded
	MaxDepth int
	// buffered packets (pushed minus popped); jitterbuffer exposes no length.
	// Playout latency in packets: depth*OpusFrameMs ms of audio waiting.
	depth int
	// lifetime counters for telemetry (see Stats)
	cPushed, cLateDrop, cPopped, cTrimDrop, cTrimEvents, cShedDrop int64
	maxDepth                                                       int // high-water depth
}

// PlayoutStats snapshots a Playout's lifetime counters and current/peak depth.
type PlayoutStats struct {
	Depth, MaxDepthSeen                                      int
	Pushed, LateDrop, Popped, TrimDrop, TrimEvents, ShedDrop int64
}

var ErrEmpty = errors.New("playout: empty")

// primePackets=2 is ~40ms at 20ms framing
func NewPlayout(primePackets uint16) *Playout {
	if primePackets == 0 {
		primePackets = 2
	}
	return &Playout{
		jb:          jitterbuffer.New(jitterbuffer.WithMinimumPacketCount(primePackets)),
		MaxSkip:     64,
		TargetDepth: int(primePackets),
		MaxDepth:    DefaultMaxDepth,
	}
}

// late packets (seq behind headSeq) are dropped
func (p *Playout) Push(pkt *rtp.Packet) {
	if pkt == nil {
		return
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	// int16 cast -> signed seq distance
	if p.started && int16(pkt.SequenceNumber-p.headSeq) < 0 {
		p.cLateDrop++
		return
	}
	p.jb.Push(pkt)
	p.depth++
	p.cPushed++
	if p.depth > p.maxDepth {
		p.maxDepth = p.depth
	}
	// hard cap: keep bounding latency even while the consumer is stalled
	// (Push runs on the reader goroutine, independent of the ticker)
	if p.MaxDepth > 0 && p.depth > p.MaxDepth {
		p.cTrimEvents++
		p.trimLocked(p.TargetDepth)
	}
}

// Stats snapshots the lifetime counters and current/peak depth.
func (p *Playout) Stats() PlayoutStats {
	p.mu.Lock()
	defer p.mu.Unlock()
	return PlayoutStats{
		Depth:        p.depth,
		MaxDepthSeen: p.maxDepth,
		Pushed:       p.cPushed,
		LateDrop:     p.cLateDrop,
		Popped:       p.cPopped,
		TrimDrop:     p.cTrimDrop,
		TrimEvents:   p.cTrimEvents,
		ShedDrop:     p.cShedDrop,
	}
}

// trimLocked drops the oldest buffered packets until at most target remain.
// Caller holds p.mu.
func (p *Playout) trimLocked(target int) {
	// budget guards against spinning across a long run of missing seqs
	budget := p.depth + int(p.MaxSkip)
	for p.depth > target && budget > 0 {
		budget--
		if pkt, err := p.jb.Pop(); err == nil {
			p.started = true
			p.headSeq = pkt.SequenceNumber + 1
			p.depth--
			p.cTrimDrop++
		} else {
			// hole at the head; step over it toward the next present packet
			p.jb.SetPlayoutHead(p.jb.PlayoutHead() + 1)
		}
	}
}

// ShedOne drops the oldest buffered packet if depth exceeds target, returning
// whether it dropped. Called once per pop so post-stall latency drains back to
// target gradually rather than surfacing as call-long lag.
func (p *Playout) ShedOne(target int) bool {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.depth <= target {
		return false
	}
	if pkt, err := p.jb.Pop(); err == nil {
		p.started = true
		p.headSeq = pkt.SequenceNumber + 1
		p.depth--
		p.cShedDrop++
		return true
	}
	// hole at the head; step past it and let the next pop try again
	p.jb.SetPlayoutHead(p.jb.PlayoutHead() + 1)
	return false
}

// Depth reports the number of buffered packets not yet popped, i.e. the
// current playout latency in packets (depth*OpusFrameMs ms of audio).
func (p *Playout) Depth() int {
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.depth
}

// returns (next playable packet, gap count to PLC over, err);
// ErrEmpty means priming or genuinely empty
func (p *Playout) PopOrSkip() (*rtp.Packet, int, error) {
	p.mu.Lock()
	defer p.mu.Unlock()

	if pkt, err := p.jb.Pop(); err == nil {
		p.started = true
		p.headSeq = pkt.SequenceNumber + 1
		p.depth--
		p.cPopped++
		return pkt, 0, nil
	} else if errors.Is(err, jitterbuffer.ErrPopWhileBuffering) {
		return nil, 0, ErrEmpty
	}

	// scan forward for the next available packet
	head := p.jb.PlayoutHead()
	for i := uint16(1); i <= p.MaxSkip; i++ {
		pkt, err := p.jb.PopAtSequence(head + i)
		if err == nil {
			// PopAtSequence advances by 1; head must skip past the gap
			p.jb.SetPlayoutHead(pkt.SequenceNumber + 1)
			p.started = true
			p.headSeq = pkt.SequenceNumber + 1
			p.depth--
			p.cPopped++
			return pkt, int(i), nil
		}
	}
	return nil, 0, ErrEmpty
}