package nebula

import (
	"net"
	"testing"
	"time"

	"github.com/slackhq/nebula/header"
	"github.com/slackhq/nebula/iputil"
	"github.com/slackhq/nebula/udp"
	"github.com/slackhq/nebula/util"
	"github.com/stretchr/testify/assert"
)

func Test_NewHandshakeManagerVpnIp(t *testing.T) {
	l := util.NewTestLogger()
	_, tuncidr, _ := net.ParseCIDR("172.1.1.1/24")
	_, vpncidr, _ := net.ParseCIDR("172.1.1.1/24")
	_, localrange, _ := net.ParseCIDR("10.1.1.1/24")
	ip := iputil.Ip2VpnIp(net.ParseIP("172.1.1.2"))
	preferredRanges := []*net.IPNet{localrange}
	mw := &mockEncWriter{}
	mainHM := NewHostMap(l, "test", vpncidr, preferredRanges)

	blah := NewHandshakeManager(l, tuncidr, preferredRanges, mainHM, &LightHouse{}, &udp.Conn{}, defaultHandshakeConfig)

	now := time.Now()
	blah.NextOutboundHandshakeTimerTick(now, mw)

	var initCalled bool
	initFunc := func(*HostInfo) {
		initCalled = true
	}

	i := blah.AddVpnIp(ip, initFunc)
	assert.True(t, initCalled)

	initCalled = false
	i2 := blah.AddVpnIp(ip, initFunc)
	assert.False(t, initCalled)
	assert.Same(t, i, i2)

	i.remotes = NewRemoteList()
	i.HandshakeReady = true

	// Adding something to pending should not affect the main hostmap
	assert.Len(t, mainHM.Hosts, 0)

	// Confirm they are in the pending index list
	assert.Contains(t, blah.pendingHostMap.Hosts, ip)

	// Jump ahead `HandshakeRetries` ticks, offset by one to get the sleep logic right
	for i := 1; i <= DefaultHandshakeRetries+1; i++ {
		now = now.Add(time.Duration(i) * DefaultHandshakeTryInterval)
		blah.NextOutboundHandshakeTimerTick(now, mw)
	}

	// Confirm they are still in the pending index list
	assert.Contains(t, blah.pendingHostMap.Hosts, ip)

	// Tick 1 more time, a minute will certainly flush it out
	blah.NextOutboundHandshakeTimerTick(now.Add(time.Minute), mw)

	// Confirm they have been removed
	assert.NotContains(t, blah.pendingHostMap.Hosts, ip)
}

func Test_NewHandshakeManagerTrigger(t *testing.T) {
	l := util.NewTestLogger()
	_, tuncidr, _ := net.ParseCIDR("172.1.1.1/24")
	_, vpncidr, _ := net.ParseCIDR("172.1.1.1/24")
	_, localrange, _ := net.ParseCIDR("10.1.1.1/24")
	ip := iputil.Ip2VpnIp(net.ParseIP("172.1.1.2"))
	preferredRanges := []*net.IPNet{localrange}
	mw := &mockEncWriter{}
	mainHM := NewHostMap(l, "test", vpncidr, preferredRanges)
	lh := &LightHouse{addrMap: make(map[iputil.VpnIp]*RemoteList), l: l}

	blah := NewHandshakeManager(l, tuncidr, preferredRanges, mainHM, lh, &udp.Conn{}, defaultHandshakeConfig)

	now := time.Now()
	blah.NextOutboundHandshakeTimerTick(now, mw)

	assert.Equal(t, 0, testCountTimerWheelEntries(blah.OutboundHandshakeTimer))

	hi := blah.AddVpnIp(ip, nil)
	hi.HandshakeReady = true
	assert.Equal(t, 1, testCountTimerWheelEntries(blah.OutboundHandshakeTimer))
	assert.Equal(t, 0, hi.HandshakeCounter, "Should not have attempted a handshake yet")

	// Trigger the same method the channel will but, this should set our remotes pointer
	blah.handleOutbound(ip, mw, true)
	assert.Equal(t, 1, hi.HandshakeCounter, "Trigger should have done a handshake attempt")
	assert.NotNil(t, hi.remotes, "Manager should have set my remotes pointer")

	// Make sure the trigger doesn't double schedule the timer entry
	assert.Equal(t, 1, testCountTimerWheelEntries(blah.OutboundHandshakeTimer))

	uaddr := udp.NewAddrFromString("10.1.1.1:4242")
	hi.remotes.unlockedPrependV4(ip, NewIp4AndPort(uaddr.IP, uint32(uaddr.Port)))

	// We now have remotes but only the first trigger should have pushed things forward
	blah.handleOutbound(ip, mw, true)
	assert.Equal(t, 1, hi.HandshakeCounter, "Trigger should have not done a handshake attempt")
	assert.Equal(t, 1, testCountTimerWheelEntries(blah.OutboundHandshakeTimer))
}

func testCountTimerWheelEntries(tw *SystemTimerWheel) (c int) {
	for _, i := range tw.wheel {
		n := i.Head
		for n != nil {
			c++
			n = n.Next
		}
	}
	return c
}

type mockEncWriter struct {
}

func (mw *mockEncWriter) SendMessageToVpnIp(t header.MessageType, st header.MessageSubType, vpnIp iputil.VpnIp, p, nb, out []byte) {
	return
}