package mux

import (
	"bytes"
	"context"
	crand "crypto/rand"
	"fmt"
	"io"
	mrand "math/rand"
	"net"
	"reflect"
	"runtime"
	"runtime/debug"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/libp2p/go-libp2p-testing/ci"
	"github.com/libp2p/go-libp2p/core/network"
	"github.com/libp2p/go-libp2p/core/peer"

	"github.com/stretchr/testify/require"
)

var randomness []byte
var Subtests map[string]TransportTest

func init() {
	// read 1MB of randomness
	randomness = make([]byte, 1<<20)
	if _, err := crand.Read(randomness); err != nil {
		panic(err)
	}

	Subtests = make(map[string]TransportTest)
	for _, f := range subtests {
		Subtests[getFunctionName(f)] = f
	}
}

func getFunctionName(i interface{}) string {
	return runtime.FuncForPC(reflect.ValueOf(i).Pointer()).Name()
}

type peerScope struct {
	mx     sync.Mutex
	memory int
}

func (p *peerScope) ReserveMemory(size int, _ uint8) error {
	p.mx.Lock()
	p.memory += size
	p.mx.Unlock()
	return nil
}

func (p *peerScope) ReleaseMemory(size int) {
	p.mx.Lock()
	defer p.mx.Unlock()
	if p.memory < size {
		panic(fmt.Sprintf("tried to release too much memory: %d (current: %d)", size, p.memory))
	}
	p.memory -= size
}

// Check checks that we don't have any more reserved memory.
func (p *peerScope) Check(t *testing.T) {
	p.mx.Lock()
	defer p.mx.Unlock()
	require.Zero(t, p.memory, "expected all reserved memory to have been released")
}

type peerScopeSpan struct {
	peerScope
}

func (p *peerScopeSpan) Done() {
	p.mx.Lock()
	defer p.mx.Unlock()
	p.memory = 0
}

func (p *peerScope) Stat() network.ScopeStat                       { return network.ScopeStat{} }
func (p *peerScope) BeginSpan() (network.ResourceScopeSpan, error) { return &peerScopeSpan{}, nil }
func (p *peerScope) Peer() peer.ID                                 { panic("implement me") }

var _ network.PeerScope = &peerScope{}

type Options struct {
	tr        network.Multiplexer
	connNum   int
	streamNum int
	msgNum    int
	msgMin    int
	msgMax    int
}

func randBuf(size int) []byte {
	n := len(randomness) - size
	if size < 1 {
		panic(fmt.Errorf("requested too large buffer (%d). max is %d", size, len(randomness)))
	}

	start := mrand.Intn(n)
	return randomness[start : start+size]
}

func checkErr(t *testing.T, err error) {
	if err != nil {
		debug.PrintStack()
		t.Fatal(err)
	}
}

func echoStream(s network.MuxedStream) {
	defer s.Close()
	io.Copy(s, s) // echo everything
}

func GoServe(t *testing.T, tr network.Multiplexer, l net.Listener) (done func()) {
	closed := make(chan struct{}, 1)

	go func() {
		for {
			c1, err := l.Accept()
			if err != nil {
				select {
				case <-closed:
					return // closed naturally.
				default:
					checkErr(t, err)
				}
			}

			sc1, err := tr.NewConn(c1, true, nil)
			checkErr(t, err)
			go func() {
				for {
					str, err := sc1.AcceptStream()
					if err != nil {
						break
					}
					go echoStream(str)
				}
			}()
		}
	}()

	return func() {
		closed <- struct{}{}
	}
}

func SubtestSimpleWrite(t *testing.T, tr network.Multiplexer) {
	l, err := net.Listen("tcp", "localhost:0")
	checkErr(t, err)
	done := GoServe(t, tr, l)
	defer done()

	nc1, err := net.Dial("tcp", l.Addr().String())
	checkErr(t, err)
	defer nc1.Close()

	scope := &peerScope{}
	c1, err := tr.NewConn(nc1, false, scope)
	checkErr(t, err)
	defer func() {
		c1.Close()
		scope.Check(t)
	}()

	// serve the outgoing conn, because some muxers assume
	// that we _always_ call serve. (this is an error?)
	go c1.AcceptStream()

	s1, err := c1.OpenStream(context.Background())
	checkErr(t, err)
	defer s1.Close()

	buf1 := randBuf(4096)
	_, err = s1.Write(buf1)
	checkErr(t, err)

	buf2 := make([]byte, len(buf1))
	_, err = io.ReadFull(s1, buf2)
	checkErr(t, err)

	require.Equal(t, buf1, buf2)
}

func SubtestStress(t *testing.T, opt Options) {
	msgsize := 1 << 11
	errs := make(chan error) // dont block anything.

	rateLimitN := 5000 // max of 5k funcs, because -race has 8k max.
	rateLimitChan := make(chan struct{}, rateLimitN)
	for i := 0; i < rateLimitN; i++ {
		rateLimitChan <- struct{}{}
	}

	rateLimit := func(f func()) {
		<-rateLimitChan
		f()
		rateLimitChan <- struct{}{}
	}

	writeStream := func(s network.MuxedStream, bufs chan<- []byte) {
		for i := 0; i < opt.msgNum; i++ {
			buf := randBuf(msgsize)
			bufs <- buf
			if _, err := s.Write(buf); err != nil {
				errs <- fmt.Errorf("s.Write(buf): %s", err)
				continue
			}
		}
	}

	readStream := func(s network.MuxedStream, bufs <-chan []byte) {
		buf2 := make([]byte, msgsize)
		for buf1 := range bufs {
			if _, err := io.ReadFull(s, buf2); err != nil {
				errs <- fmt.Errorf("io.ReadFull(s, buf2): %s", err)
				continue
			}
			if !bytes.Equal(buf1, buf2) {
				errs <- fmt.Errorf("buffers not equal (%x != %x)", buf1[:3], buf2[:3])
			}
		}
	}

	openStreamAndRW := func(c network.MuxedConn) {
		s, err := c.OpenStream(context.Background())
		if err != nil {
			errs <- fmt.Errorf("failed to create NewStream: %s", err)
			return
		}

		bufs := make(chan []byte, opt.msgNum)
		go func() {
			writeStream(s, bufs)
			close(bufs)
		}()

		readStream(s, bufs)
		s.Close()
	}

	openConnAndRW := func() {
		l, err := net.Listen("tcp", "localhost:0")
		checkErr(t, err)
		done := GoServe(t, opt.tr, l)
		defer done()

		nla := l.Addr()
		nc, err := net.Dial(nla.Network(), nla.String())
		checkErr(t, err)
		if err != nil {
			t.Fatal(fmt.Errorf("net.Dial(%s, %s): %s", nla.Network(), nla.String(), err))
			return
		}

		scope := &peerScope{}
		c, err := opt.tr.NewConn(nc, false, scope)
		if err != nil {
			t.Fatal(fmt.Errorf("a.AddConn(%s <--> %s): %s", nc.LocalAddr(), nc.RemoteAddr(), err))
			return
		}

		// serve the outgoing conn, because some muxers assume
		// that we _always_ call serve. (this is an error?)
		go func() {
			for {
				str, err := c.AcceptStream()
				if err != nil {
					break
				}
				go echoStream(str)
			}
		}()

		var wg sync.WaitGroup
		for i := 0; i < opt.streamNum; i++ {
			wg.Add(1)
			go rateLimit(func() {
				defer wg.Done()
				openStreamAndRW(c)
			})
		}
		wg.Wait()
		c.Close()
		scope.Check(t)
	}

	openConnsAndRW := func() {
		var wg sync.WaitGroup
		for i := 0; i < opt.connNum; i++ {
			wg.Add(1)
			go rateLimit(func() {
				defer wg.Done()
				openConnAndRW()
			})
		}
		wg.Wait()
	}

	go func() {
		openConnsAndRW()
		close(errs) // done
	}()

	for err := range errs {
		t.Error(err)
	}

}

func tcpPipe(t *testing.T) (net.Conn, net.Conn) {
	list, err := net.Listen("tcp", "0.0.0.0:0")
	if err != nil {
		t.Fatal(err)
	}

	con1, err := net.Dial("tcp", list.Addr().String())
	if err != nil {
		t.Fatal(err)
	}

	con2, err := list.Accept()
	if err != nil {
		t.Fatal(err)
	}

	return con1, con2
}

func SubtestStreamOpenStress(t *testing.T, tr network.Multiplexer) {
	wg := new(sync.WaitGroup)

	a, b := tcpPipe(t)
	defer a.Close()
	defer b.Close()

	defer wg.Wait()

	wg.Add(1)
	count := 10000
	workers := 5
	go func() {
		defer wg.Done()
		muxa, err := tr.NewConn(a, true, nil)
		if err != nil {
			t.Error(err)
			return
		}
		stress := func() {
			defer wg.Done()
			for i := 0; i < count; i++ {
				s, err := muxa.OpenStream(context.Background())
				if err != nil {
					t.Error(err)
					return
				}
				err = s.CloseWrite()
				if err != nil {
					t.Error(err)
				}
				n, err := s.Read([]byte{0})
				if n != 0 {
					t.Error("expected to read no bytes")
				}
				if err != io.EOF {
					t.Errorf("expected an EOF, got %s", err)
				}
			}
		}

		for i := 0; i < workers; i++ {
			wg.Add(1)
			go stress()
		}
	}()

	scope := &peerScope{}
	muxb, err := tr.NewConn(b, false, scope)
	if err != nil {
		t.Fatal(err)
	}
	defer func() {
		muxb.Close()
		scope.Check(t)
	}()

	time.Sleep(time.Millisecond * 50)

	wg.Add(1)
	recv := make(chan struct{}, count*workers)
	go func() {
		defer wg.Done()
		for i := 0; i < count*workers; i++ {
			str, err := muxb.AcceptStream()
			if err != nil {
				break
			}
			wg.Add(1)
			go func() {
				defer wg.Done()
				str.Close()
				select {
				case recv <- struct{}{}:
				default:
					t.Error("too many stream")
				}
			}()
		}
	}()

	timeout := time.Second * 10
	if ci.IsRunning() {
		timeout *= 10
	}

	limit := time.After(timeout)
	for i := 0; i < count*workers; i++ {
		select {
		case <-recv:
		case <-limit:
			t.Fatal("timed out receiving streams")
		}
	}

	wg.Wait()
}

func SubtestStreamReset(t *testing.T, tr network.Multiplexer) {
	wg := new(sync.WaitGroup)
	defer wg.Wait()

	a, b := tcpPipe(t)
	defer a.Close()
	defer b.Close()

	wg.Add(1)
	scopea := &peerScope{}
	muxa, err := tr.NewConn(a, true, scopea)
	if err != nil {
		t.Error(err)
		return
	}
	defer func() {
		muxa.Close()
		scopea.Check(t)
	}()

	go func() {
		defer wg.Done()
		s, err := muxa.OpenStream(context.Background())
		if err != nil {
			t.Error(err)
			return
		}
		time.Sleep(time.Millisecond * 50)

		_, err = s.Write([]byte("foo"))
		if err != network.ErrReset {
			t.Error("should have been stream reset")
		}
		s.Close()
	}()

	scopeb := &peerScope{}
	muxb, err := tr.NewConn(b, false, scopeb)
	if err != nil {
		t.Fatal(err)
	}
	defer func() {
		muxb.Close()
		scopeb.Check(t)
	}()

	str, err := muxb.AcceptStream()
	checkErr(t, err)
	str.Reset()

	wg.Wait()
}

// check that Close also closes the underlying net.Conn
func SubtestWriteAfterClose(t *testing.T, tr network.Multiplexer) {
	a, b := tcpPipe(t)

	scopea := &peerScope{}
	muxa, err := tr.NewConn(a, true, scopea)
	checkErr(t, err)

	scopeb := &peerScope{}
	muxb, err := tr.NewConn(b, false, scopeb)
	checkErr(t, err)

	checkErr(t, muxa.Close())
	scopea.Check(t)
	checkErr(t, muxb.Close())
	scopeb.Check(t)

	// make sure the underlying net.Conn was closed
	if _, err := a.Write([]byte("foobar")); err == nil || !strings.Contains(err.Error(), "use of closed network connection") {
		t.Fatal("write should have failed")
	}
	if _, err := b.Write([]byte("foobar")); err == nil || !strings.Contains(err.Error(), "use of closed network connection") {
		t.Fatal("write should have failed")
	}
}

func SubtestStreamLeftOpen(t *testing.T, tr network.Multiplexer) {
	a, b := tcpPipe(t)

	const numStreams = 10
	const dataLen = 50 * 1024

	scopea := &peerScope{}
	muxa, err := tr.NewConn(a, true, scopea)
	checkErr(t, err)

	scopeb := &peerScope{}
	muxb, err := tr.NewConn(b, false, scopeb)
	checkErr(t, err)

	var wg sync.WaitGroup
	wg.Add(1 + numStreams)
	go func() {
		defer wg.Done()
		for i := 0; i < numStreams; i++ {
			stra, err := muxa.OpenStream(context.Background())
			checkErr(t, err)
			go func() {
				defer wg.Done()
				_, err = stra.Write(randBuf(dataLen))
				checkErr(t, err)
				// do NOT close or reset the stream
			}()
		}
	}()

	wg.Add(1 + numStreams)
	go func() {
		defer wg.Done()
		for i := 0; i < numStreams; i++ {
			str, err := muxb.AcceptStream()
			checkErr(t, err)
			go func() {
				defer wg.Done()
				_, err = io.ReadFull(str, make([]byte, dataLen))
				checkErr(t, err)
			}()
		}
	}()

	// Now we have a bunch of open streams.
	// Make sure that their memory is returned when we close the connection.
	wg.Wait()

	muxa.Close()
	scopea.Check(t)
	muxb.Close()
	scopeb.Check(t)
}

func SubtestStress1Conn1Stream1Msg(t *testing.T, tr network.Multiplexer) {
	SubtestStress(t, Options{
		tr:        tr,
		connNum:   1,
		streamNum: 1,
		msgNum:    1,
		msgMax:    100,
		msgMin:    100,
	})
}

func SubtestStress1Conn1Stream100Msg(t *testing.T, tr network.Multiplexer) {
	SubtestStress(t, Options{
		tr:        tr,
		connNum:   1,
		streamNum: 1,
		msgNum:    100,
		msgMax:    100,
		msgMin:    100,
	})
}

func SubtestStress1Conn100Stream100Msg(t *testing.T, tr network.Multiplexer) {
	SubtestStress(t, Options{
		tr:        tr,
		connNum:   1,
		streamNum: 100,
		msgNum:    100,
		msgMax:    100,
		msgMin:    100,
	})
}

func SubtestStress10Conn10Stream50Msg(t *testing.T, tr network.Multiplexer) {
	SubtestStress(t, Options{
		tr:        tr,
		connNum:   10,
		streamNum: 10,
		msgNum:    50,
		msgMax:    100,
		msgMin:    100,
	})
}

func SubtestStress1Conn1000Stream10Msg(t *testing.T, tr network.Multiplexer) {
	SubtestStress(t, Options{
		tr:        tr,
		connNum:   1,
		streamNum: 1000,
		msgNum:    10,
		msgMax:    100,
		msgMin:    100,
	})
}

func SubtestStress1Conn100Stream100Msg10MB(t *testing.T, tr network.Multiplexer) {
	SubtestStress(t, Options{
		tr:        tr,
		connNum:   1,
		streamNum: 100,
		msgNum:    100,
		msgMax:    10000,
		msgMin:    1000,
	})
}

// Subtests are all the subtests run by SubtestAll
var subtests = []TransportTest{
	SubtestSimpleWrite,
	SubtestWriteAfterClose,
	SubtestStress1Conn1Stream1Msg,
	SubtestStress1Conn1Stream100Msg,
	SubtestStress1Conn100Stream100Msg,
	SubtestStress10Conn10Stream50Msg,
	SubtestStress1Conn1000Stream10Msg,
	SubtestStress1Conn100Stream100Msg10MB,
	SubtestStreamOpenStress,
	SubtestStreamReset,
	SubtestStreamLeftOpen,
}

// SubtestAll runs all the stream multiplexer tests against the target
// transport.
func SubtestAll(t *testing.T, tr network.Multiplexer) {
	for name, f := range Subtests {
		t.Run(name, func(t *testing.T) {
			f(t, tr)
		})
	}
}

// TransportTest is a stream multiplex transport test case
type TransportTest func(t *testing.T, tr network.Multiplexer)