package packetcache

import (
	"bytes"
	"math/rand"
	"sync"
	"testing"

	"github.com/pion/rtcp"
)

func randomBuf() []byte {
	length := rand.Int31n(BufSize-1) + 1
	buf := make([]byte, length)
	rand.Read(buf)
	return buf
}

func TestCache(t *testing.T) {
	buf1 := randomBuf()
	buf2 := randomBuf()
	cache := New(16)
	_, i1 := cache.Store(13, 0, false, buf1)
	_, i2 := cache.Store(17, 0, false, buf2)

	buf := make([]byte, BufSize)

	l := cache.Get(13, buf)
	if !bytes.Equal(buf[:l], buf1) {
		t.Errorf("Couldn't get 13")
	}
	l = cache.GetAt(13, i1, buf)
	if !bytes.Equal(buf[:l], buf1) {
		t.Errorf("Couldn't get 13 at %v", i1)
	}
	l = cache.Get(17, buf)
	if !bytes.Equal(buf[:l], buf2) {
		t.Errorf("Couldn't get 17")
	}
	l = cache.GetAt(17, i2, buf)
	if !bytes.Equal(buf[:l], buf2) {
		t.Errorf("Couldn't get 17 at %v", i2)
	}

	l = cache.Get(42, buf)
	if l != 0 {
		t.Errorf("Creation ex nihilo")
	}

	l = cache.GetAt(17, i1, buf)
	if l != 0 {
		t.Errorf("Got 17 at %v", i1)
	}

	l = cache.GetAt(42, i2, buf)
	if l != 0 {
		t.Errorf("Got 42 at %v", i2)
	}
}

func TestCacheOverflow(t *testing.T) {
	cache := New(16)

	for i := 0; i < 32; i++ {
		cache.Store(uint16(i), 0, false, []byte{uint8(i)})
	}

	for i := 0; i < 32; i++ {
		buf := make([]byte, BufSize)
		l := cache.Get(uint16(i), buf)
		if i < 16 {
			if l > 0 {
				t.Errorf("Creation ex nihilo: %v", i)
			}
		} else {
			if l != 1 || buf[0] != uint8(i) {
				t.Errorf("Expected [%v], got %v", i, buf[:l])
			}
		}
	}
}

func TestCacheGrow(t *testing.T) {
	cache := New(16)

	for i := 0; i < 24; i++ {
		cache.Store(uint16(i), 0, false, []byte{uint8(i)})
	}

	cache.Resize(32)
	for i := 0; i < 32; i++ {
		expected := uint16(0)
		if i < 8 {
			expected = uint16(i + 16)
		}
		if i >= 24 {
			expected = uint16(i - 16)
		}
		if cache.entries[i].seqno != expected {
			t.Errorf("At %v, got %v, expected %v",
				i, cache.entries[i].seqno, expected)
		}
	}
}

func TestCacheShrink(t *testing.T) {
	cache := New(16)

	for i := 0; i < 24; i++ {
		cache.Store(uint16(i), 0, false, []byte{uint8(i)})
	}

	cache.Resize(12)
	for i := 0; i < 12; i++ {
		expected := uint16(i + 16)
		if i >= 8 {
			expected = uint16(i + 4)
		}
		if cache.entries[i].seqno != expected {
			t.Errorf("At %v, got %v, expected %v",
				i, cache.entries[i].seqno, expected)
		}
	}
}

func TestCacheGrowCond(t *testing.T) {
	cache := New(16)
	if len(cache.entries) != 16 {
		t.Errorf("Expected 16, got %v", len(cache.entries))
	}

	done := cache.ResizeCond(17)
	if done || len(cache.entries) != 16 {
		t.Errorf("Grew cache by 1")
	}

	done = cache.ResizeCond(15)
	if done || len(cache.entries) != 16 {
		t.Errorf("Shrunk cache by 1")
	}

	done = cache.ResizeCond(32)
	if !done || len(cache.entries) != 32 {
		t.Errorf("Didn't grow cache")
	}

	done = cache.ResizeCond(16)
	if !done || len(cache.entries) != 16 {
		t.Errorf("Didn't shrink cache")
	}
}

func TestKeyframe(t *testing.T) {
	cache := New(16)
	packet := make([]byte, 1)
	buf := make([]byte, BufSize)

	cache.Store(7, 57, true, packet)
	cache.Store(8, 57, true, packet)

	ts, kf := cache.Keyframe()
	if ts != 57 || len(kf) != 2 {
		t.Errorf("Got %v %v, expected %v %v", ts, len(kf), 57, 2)
	}
	for _, i := range kf {
		l := cache.Get(i, buf)
		if int(l) != len(packet) {
			t.Errorf("Couldn't get %v", i)
		}
	}

	for i := 0; i < 32; i++ {
		cache.Store(uint16(9 + i), uint32(58 + i), false, packet)
	}

	ts, kf = cache.Keyframe()
	if ts != 57 || len(kf) != 2 {
		t.Errorf("Got %v %v, expected %v %v", ts, len(kf), 57, 2)
	}
	for _, i := range kf {
		l := cache.Get(i, buf)
		if int(l) != len(packet) {
			t.Errorf("Couldn't get %v", i)
		}
	}
}

func TestKeyframeUnsorted(t *testing.T) {
	cache := New(16)
	packet := make([]byte, 1)

	cache.Store(7, 57, true, packet)
	cache.Store(9, 57, true, packet)
	cache.Store(8, 57, true, packet)
	cache.Store(10, 57, true, packet)
	cache.Store(6, 57, true, packet)
	cache.Store(8, 57, true, packet)

	_, kf := cache.Keyframe()
	if len(kf) != 5 {
		t.Errorf("Got length %v, expected 5", len(kf))
	}
	for i, v := range kf {
		if v != uint16(i + 6) {
			t.Errorf("Position %v, expected %v, got %v\n",
				i, i + 6, v)
		}
	}
}


func TestBitmap(t *testing.T) {
	value := uint64(0xcdd58f1e035379c0)
	packet := make([]byte, 1)

	cache := New(16)

	var first uint16
	for i := 0; i < 64; i++ {
		if (value & (1 << i)) != 0 {
			first, _ = cache.Store(uint16(42+i), 0, false, packet)
		}
	}

	value >>= uint16(first - 42)
	if uint32(value) != cache.bitmap {
		t.Errorf("Got %b, expected %b", cache.bitmap, value)
	}
}

func TestBitmapWrap(t *testing.T) {
	value := uint64(0xcdd58f1e035379c0)
	packet := make([]byte, 1)

	cache := New(16)

	cache.Store(0x7000, 0, false, packet)
	cache.Store(0xA000, 0, false, packet)

	var first uint16
	for i := 0; i < 64; i++ {
		if (value & (1 << i)) != 0 {
			first, _ = cache.Store(uint16(42+i), 0, false, packet)
		}
	}

	value >>= uint16(first - 42)
	if uint32(value) != cache.bitmap {
		t.Errorf("Got %b, expected %b", cache.bitmap, value)
	}
}

func TestBitmapGet(t *testing.T) {
	value := uint64(0xcdd58f1e035379c0)
	packet := make([]byte, 1)

	cache := New(16)

	for i := 0; i < 64; i++ {
		if (value & (1 << i)) != 0 {
			cache.Store(uint16(42+i), 0, false, packet)
		}
	}

	pos := uint16(42)
	for cache.bitmap != 0 {
		found, first, bitmap := cache.BitmapGet()
		if first < pos || first >= pos+64 {
			t.Errorf("First is %v, pos is %v", first, pos)
		}
		if !found {
			t.Fatalf("Didn't find any 0 bits")
		}
		value >>= (first - pos)
		pos = first
		if (value & 1) != 0 {
			t.Errorf("Value is odd")
		}
		value >>= 1
		pos += 1
		for bitmap != 0 {
			if uint8(bitmap&1) == uint8(value&1) {
				t.Errorf("Bitmap mismatch")
			}
			bitmap >>= 1
			value >>= 1
			pos += 1
		}
	}
	if value != 0 {
		t.Errorf("Value is %v", value)
	}
}

func TestBitmapPacket(t *testing.T) {
	value := uint64(0xcdd58f1e035379c0)
	packet := make([]byte, 1)

	cache := New(16)

	for i := 0; i < 64; i++ {
		if (value & (1 << i)) != 0 {
			cache.Store(uint16(42+i), 0, false, packet)
		}
	}

	found, first, bitmap := cache.BitmapGet()

	if !found {
		t.Fatalf("Didn't find any 0 bits")
	}

	p := rtcp.NackPair{first, rtcp.PacketBitmap(bitmap)}
	pl := p.PacketList()

	for _, s := range pl {
		if s < 42 || s >= 42+64 {
			if (value & (1 << (s - 42))) != 0 {
				t.Errorf("Bit %v unexpectedly set", s-42)
			}
		}
	}
}

func BenchmarkCachePutGet(b *testing.B) {
	n := 10
	chans := make([]chan uint16, n)
	for i := range chans {
		chans[i] = make(chan uint16, 8)
	}

	cache := New(96)

	var wg sync.WaitGroup
	wg.Add(len(chans))

	for i := range chans {
		go func(ch <-chan uint16) {
			defer wg.Done()
			buf := make([]byte, BufSize)
			for {
				seqno, ok := <-ch
				if !ok {
					return
				}
				l := cache.Get(seqno, buf)
				if l == 0 {
					b.Errorf("Couldn't get %v", seqno)
				}
			}
		}(chans[i])
	}

	buf := make([]byte, 1200)

	b.SetBytes(1200)
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		seqno := uint16(i)
		cache.Store(seqno, 0, false, buf)
		for _, ch := range chans {
			ch <- seqno
		}
	}
	for _, ch := range chans {
		close(ch)
	}
	wg.Wait()
}

type is struct {
	index, seqno uint16
}

func BenchmarkCachePutGetAt(b *testing.B) {
	n := 10
	chans := make([]chan is, n)
	for i := range chans {
		chans[i] = make(chan is, 8)
	}

	cache := New(96)

	var wg sync.WaitGroup
	wg.Add(len(chans))

	for i := range chans {
		go func(ch <-chan is) {
			defer wg.Done()
			buf := make([]byte, BufSize)
			for {
				is, ok := <-ch
				if !ok {
					return
				}
				l := cache.GetAt(is.seqno, is.index, buf)
				if l == 0 {
					b.Errorf("Couldn't get %v", is)
				}
			}
		}(chans[i])
	}

	buf := make([]byte, 1200)

	b.SetBytes(1200)
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		seqno := uint16(i)
		_, index := cache.Store(seqno, 0, false, buf)
		for _, ch := range chans {
			ch <- is{index, seqno}
		}
	}
	for _, ch := range chans {
		close(ch)
	}
	wg.Wait()
}