package peer

import (
	"net/netip"
	"testing"
	"time"
	"vppn/m"

	"git.crumpington.com/lib/go/ratelimiter"
)

// ----------------------------------------------------------------------------

type PeerStateControlMsg struct {
	Peer   remotePeer
	Packet any
}

type PeerStateTestHarness struct {
	State     peerState
	Published remotePeer
	Sent      []PeerStateControlMsg
}

func NewPeerStateTestHarness() *PeerStateTestHarness {
	h := &PeerStateTestHarness{}

	keys := generateKeys()

	state := &pState{
		publish: func(rp remotePeer) {
			h.Published = rp
		},
		sendControlPacket: func(rp remotePeer, pkt marshaller) {
			h.Sent = append(h.Sent, PeerStateControlMsg{rp, pkt})
		},
		localIP:  2,
		remoteIP: 3,
		privKey:  keys.PrivKey,
		pubAddrs: newPubAddrStore(netip.AddrPort{}),
		limiter: ratelimiter.New(ratelimiter.Config{
			FillPeriod:   20 * time.Millisecond,
			MaxWaitCount: 1,
		}),
	}

	h.State = enterStateDisconnected(state)
	return h
}

func (h *PeerStateTestHarness) PeerUpdate(p *m.Peer) {
	if s := h.State.OnPeerUpdate(p); s != nil {
		h.State = s
	}
}

func (h *PeerStateTestHarness) OnSyn(msg controlMsg[packetSyn]) {
	if s := h.State.OnSyn(msg); s != nil {
		h.State = s
	}
}

func (h *PeerStateTestHarness) OnProbe(msg controlMsg[packetProbe]) {
	if s := h.State.OnProbe(msg); s != nil {
		h.State = s
	}
}

func (h *PeerStateTestHarness) OnPingTimer() {
	if s := h.State.OnPingTimer(); s != nil {
		h.State = s
	}
}

func (h *PeerStateTestHarness) ConfigServer_Public(t *testing.T) *stateServer {
	keys := generateKeys()

	state := h.State.(*stateDisconnected)
	state.localAddr = addrPort4(1, 1, 1, 2, 200)

	peer := &m.Peer{
		PeerIP:     3,
		PublicIP:   []byte{1, 1, 1, 3},
		Port:       456,
		PubKey:     keys.PubKey,
		PubSignKey: keys.PubSignKey,
	}

	h.PeerUpdate(peer)
	assertEqual(t, h.Published.Up, false)
	return assertType[*stateServer](t, h.State)
}

func (h *PeerStateTestHarness) ConfigServer_Relayed(t *testing.T) *stateServer {
	keys := generateKeys()
	peer := &m.Peer{
		PeerIP:     3,
		Port:       456,
		PubKey:     keys.PubKey,
		PubSignKey: keys.PubSignKey,
	}

	h.PeerUpdate(peer)
	assertEqual(t, h.Published.Up, false)
	return assertType[*stateServer](t, h.State)
}

func (h *PeerStateTestHarness) ConfigClientDirect(t *testing.T) *stateClientDirect {
	keys := generateKeys()
	peer := &m.Peer{
		PeerIP:     3,
		PublicIP:   []byte{1, 2, 3, 4},
		Port:       456,
		PubKey:     keys.PubKey,
		PubSignKey: keys.PubSignKey,
	}

	h.PeerUpdate(peer)
	assertEqual(t, h.Published.Up, false)
	return assertType[*stateClientDirect](t, h.State)
}

func (h *PeerStateTestHarness) ConfigClientRelayed(t *testing.T) *stateClientRelayed {
	keys := generateKeys()

	state := h.State.(*stateDisconnected)
	state.remoteIP = 1

	peer := &m.Peer{
		PeerIP:     3,
		Port:       456,
		PubKey:     keys.PubKey,
		PubSignKey: keys.PubSignKey,
	}

	h.PeerUpdate(peer)
	assertEqual(t, h.Published.Up, false)
	return assertType[*stateClientRelayed](t, h.State)
}

// ----------------------------------------------------------------------------

func TestPeerState_OnPeerUpdate_nilPeer(t *testing.T) {
	h := NewPeerStateTestHarness()
	h.PeerUpdate(nil)
	assertType[*stateDisconnected](t, h.State)
}

func TestPeerState_OnPeerUpdate_publicLocalIsServer(t *testing.T) {
	keys := generateKeys()
	h := NewPeerStateTestHarness()

	state := h.State.(*stateDisconnected)
	state.localAddr = addrPort4(1, 1, 1, 2, 200)

	peer := &m.Peer{
		PeerIP:     3,
		Port:       456,
		PubKey:     keys.PubKey,
		PubSignKey: keys.PubSignKey,
	}

	h.PeerUpdate(peer)
	assertEqual(t, h.Published.Up, false)
	assertType[*stateServer](t, h.State)
}

func TestPeerState_OnPeerUpdate_serverDirect(t *testing.T) {
	h := NewPeerStateTestHarness()
	h.ConfigServer_Public(t)
}

func TestPeerState_OnPeerUpdate_serverRelayed(t *testing.T) {
	h := NewPeerStateTestHarness()
	h.ConfigServer_Relayed(t)
}

func TestPeerState_OnPeerUpdate_clientDirect(t *testing.T) {
	h := NewPeerStateTestHarness()
	h.ConfigClientDirect(t)
}

func TestPeerState_OnPeerUpdate_clientRelayed(t *testing.T) {
	h := NewPeerStateTestHarness()
	h.ConfigClientRelayed(t)
}

func TestStateServer_directSyn(t *testing.T) {
	h := NewPeerStateTestHarness()
	h.ConfigServer_Relayed(t)

	assertEqual(t, h.Published.Up, false)

	synMsg := controlMsg[packetSyn]{
		SrcIP:   3,
		SrcAddr: addrPort4(1, 1, 1, 3, 300),
		Packet: packetSyn{
			TraceID: newTraceID(),
			//SentAt:        time.Now().UnixMilli(),
			//SharedKeyType: 1,
			Direct: true,
		},
	}

	h.State.OnSyn(synMsg)

	assertEqual(t, len(h.Sent), 1)
	ack := assertType[packetAck](t, h.Sent[0].Packet)
	assertEqual(t, ack.TraceID, synMsg.Packet.TraceID)
	assertEqual(t, h.Sent[0].Peer.IP, 3)
	assertEqual(t, ack.PossibleAddrs[0].IsValid(), false)
	assertEqual(t, h.Published.Up, true)
}

func TestStateServer_relayedSyn(t *testing.T) {
	h := NewPeerStateTestHarness()
	state := h.ConfigServer_Relayed(t)

	state.pubAddrs.Store(addrPort4(4, 5, 6, 7, 1234))

	assertEqual(t, h.Published.Up, false)

	synMsg := controlMsg[packetSyn]{
		SrcIP:   3,
		SrcAddr: addrPort4(1, 1, 1, 3, 300),
		Packet: packetSyn{
			TraceID: newTraceID(),
			//SentAt:        time.Now().UnixMilli(),
			//SharedKeyType: 1,
			Direct: false,
		},
	}
	synMsg.Packet.PossibleAddrs[0] = addrPort4(1, 1, 1, 3, 300)
	synMsg.Packet.PossibleAddrs[1] = addrPort4(2, 2, 2, 3, 300)

	h.State.OnSyn(synMsg)

	assertEqual(t, len(h.Sent), 3)

	ack := assertType[packetAck](t, h.Sent[0].Packet)
	assertEqual(t, ack.TraceID, synMsg.Packet.TraceID)
	assertEqual(t, h.Sent[0].Peer.IP, 3)
	assertEqual(t, ack.PossibleAddrs[0], addrPort4(4, 5, 6, 7, 1234))
	assertEqual(t, ack.PossibleAddrs[1].IsValid(), false)
	assertEqual(t, h.Published.Up, true)

	assertType[packetProbe](t, h.Sent[1].Packet)
	assertType[packetProbe](t, h.Sent[2].Packet)
	assertEqual(t, h.Sent[1].Peer.DirectAddr, addrPort4(1, 1, 1, 3, 300))
	assertEqual(t, h.Sent[2].Peer.DirectAddr, addrPort4(2, 2, 2, 3, 300))
}

func TestStateServer_onProbe(t *testing.T) {
	h := NewPeerStateTestHarness()
	h.ConfigServer_Relayed(t)
	assertEqual(t, h.Published.Up, false)

	probeMsg := controlMsg[packetProbe]{
		SrcIP:   3,
		SrcAddr: addrPort4(1, 1, 1, 3, 300),
		Packet:  packetProbe{TraceID: newTraceID()},
	}

	h.State.OnProbe(probeMsg)

	assertEqual(t, len(h.Sent), 1)

	probe := assertType[packetProbe](t, h.Sent[0].Packet)
	assertEqual(t, probe.TraceID, probeMsg.Packet.TraceID)
	assertEqual(t, h.Sent[0].Peer.DirectAddr, addrPort4(1, 1, 1, 3, 300))
}

func TestStateServer_OnPingTimer_timeout(t *testing.T) {
	h := NewPeerStateTestHarness()
	h.ConfigServer_Relayed(t)

	synMsg := controlMsg[packetSyn]{
		SrcIP:   3,
		SrcAddr: addrPort4(1, 1, 1, 3, 300),
		Packet: packetSyn{
			TraceID: newTraceID(),
			//SentAt:        time.Now().UnixMilli(),
			//SharedKeyType: 1,
			Direct: true,
		},
	}

	h.State.OnSyn(synMsg)
	assertEqual(t, len(h.Sent), 1)
	assertEqual(t, h.Published.Up, true)

	// Ping shouldn't timeout.
	h.OnPingTimer()
	assertEqual(t, h.Published.Up, true)

	// Advance the time, then ping.
	state := assertType[*stateServer](t, h.State)
	state.lastSeen = time.Now().Add(-timeoutInterval - time.Second)

	h.OnPingTimer()
	assertEqual(t, h.Published.Up, false)
}

func TestStateClientDirect_OnAck(t *testing.T) {
	h := NewPeerStateTestHarness()
	h.ConfigClientDirect(t)

	assertEqual(t, h.Published.Up, false)

	// On entering the state, a SYN should have been sent.
	assertEqual(t, len(h.Sent), 1)
	syn := assertType[packetSyn](t, h.Sent[0].Packet)

	ack := controlMsg[packetAck]{
		Packet: packetAck{TraceID: syn.TraceID},
	}
	h.State.OnAck(ack)
	assertEqual(t, h.Published.Up, true)
}

func TestStateClientDirect_OnAck_incorrectTraceID(t *testing.T) {
	h := NewPeerStateTestHarness()
	h.ConfigClientDirect(t)

	assertEqual(t, h.Published.Up, false)

	// On entering the state, a SYN should have been sent.
	assertEqual(t, len(h.Sent), 1)
	syn := assertType[packetSyn](t, h.Sent[0].Packet)

	ack := controlMsg[packetAck]{
		Packet: packetAck{TraceID: syn.TraceID + 1},
	}
	h.State.OnAck(ack)
	assertEqual(t, h.Published.Up, false)
}

func TestStateClientDirect_OnPingTimer(t *testing.T) {
	h := NewPeerStateTestHarness()
	h.ConfigClientDirect(t)

	// On entering the state, a SYN should have been sent.
	assertEqual(t, len(h.Sent), 1)
	assertType[packetSyn](t, h.Sent[0].Packet)

	h.OnPingTimer()

	// On ping timer, another syn should be sent. Additionally, we should remain
	// in the same state.
	assertEqual(t, len(h.Sent), 2)
	assertType[packetSyn](t, h.Sent[1].Packet)
	assertType[*stateClientDirect](t, h.State)
	assertEqual(t, h.Published.Up, false)
}

func TestStateClientDirect_OnPingTimer_timeout(t *testing.T) {
	h := NewPeerStateTestHarness()
	h.ConfigClientDirect(t)

	assertEqual(t, h.Published.Up, false)

	// On entering the state, a SYN should have been sent.
	assertEqual(t, len(h.Sent), 1)
	syn := assertType[packetSyn](t, h.Sent[0].Packet)

	ack := controlMsg[packetAck]{
		Packet: packetAck{TraceID: syn.TraceID},
	}
	h.State.OnAck(ack)
	assertEqual(t, h.Published.Up, true)

	state := assertType[*stateClientDirect](t, h.State)
	state.lastSeen = time.Now().Add(-(timeoutInterval + time.Second))

	h.OnPingTimer()

	// On ping timer, we should timeout, causing the client to reset. Another SYN
	// will be sent when re-entering the state, but the connection should be down.
	assertEqual(t, len(h.Sent), 2)
	assertType[packetSyn](t, h.Sent[1].Packet)
	assertType[*stateClientDirect](t, h.State)
	assertEqual(t, h.Published.Up, false)
}

func TestStateClientRelayed_OnAck(t *testing.T) {
	h := NewPeerStateTestHarness()
	h.ConfigClientRelayed(t)

	assertEqual(t, h.Published.Up, false)

	// On entering the state, a SYN should have been sent.
	assertEqual(t, len(h.Sent), 1)
	syn := assertType[packetSyn](t, h.Sent[0].Packet)

	ack := controlMsg[packetAck]{
		Packet: packetAck{TraceID: syn.TraceID},
	}
	h.State.OnAck(ack)
	assertEqual(t, h.Published.Up, true)
}

func TestStateClientRelayed_OnPingTimer_noAddrs(t *testing.T) {
	h := NewPeerStateTestHarness()
	h.ConfigClientRelayed(t)

	assertEqual(t, h.Published.Up, false)

	// On entering the state, a SYN should have been sent.
	assertEqual(t, len(h.Sent), 1)

	// If we haven't had an ack yet, we won't have addresses to probe. Therefore
	// we'll have just one more syn packet sent.
	h.OnPingTimer()
	assertEqual(t, len(h.Sent), 2)
}

func TestStateClientRelayed_OnPingTimer_withAddrs(t *testing.T) {
	h := NewPeerStateTestHarness()
	h.ConfigClientRelayed(t)

	assertEqual(t, h.Published.Up, false)

	// On entering the state, a SYN should have been sent.
	assertEqual(t, len(h.Sent), 1)

	syn := assertType[packetSyn](t, h.Sent[0].Packet)

	ack := controlMsg[packetAck]{Packet: packetAck{TraceID: syn.TraceID}}
	ack.Packet.PossibleAddrs[0] = addrPort4(1, 1, 1, 1, 300)
	ack.Packet.PossibleAddrs[1] = addrPort4(1, 1, 1, 2, 300)

	h.State.OnAck(ack)

	// Add a local discovery address. Note that the port will be configured port
	// and no the one provided here.
	h.State.OnLocalDiscovery(controlMsg[packetLocalDiscovery]{
		SrcIP:   3,
		SrcAddr: addrPort4(2, 2, 2, 3, 300),
	})

	// We should see one SYN and three probe packets.
	h.OnPingTimer()
	assertEqual(t, len(h.Sent), 5)
	assertType[packetSyn](t, h.Sent[1].Packet)
	assertType[packetProbe](t, h.Sent[2].Packet)
	assertType[packetProbe](t, h.Sent[3].Packet)
	assertType[packetProbe](t, h.Sent[4].Packet)

	assertEqual(t, h.Sent[2].Peer.DirectAddr, addrPort4(1, 1, 1, 1, 300))
	assertEqual(t, h.Sent[3].Peer.DirectAddr, addrPort4(1, 1, 1, 2, 300))
	assertEqual(t, h.Sent[4].Peer.DirectAddr, addrPort4(2, 2, 2, 3, 456))
}

func TestStateClientRelayed_OnPingTimer_timeout(t *testing.T) {
	h := NewPeerStateTestHarness()
	h.ConfigClientRelayed(t)

	// On entering the state, a SYN should have been sent.
	assertEqual(t, len(h.Sent), 1)
	syn := assertType[packetSyn](t, h.Sent[0].Packet)

	ack := controlMsg[packetAck]{
		Packet: packetAck{TraceID: syn.TraceID},
	}
	h.State.OnAck(ack)
	assertEqual(t, h.Published.Up, true)

	state := assertType[*stateClientRelayed](t, h.State)
	state.lastSeen = time.Now().Add(-(timeoutInterval + time.Second))

	h.OnPingTimer()

	// On ping timer, we should timeout, causing the client to reset. Another SYN
	// will be sent when re-entering the state, but the connection should be down.
	assertEqual(t, len(h.Sent), 2)
	assertType[packetSyn](t, h.Sent[1].Packet)
	assertType[*stateClientRelayed](t, h.State)
	assertEqual(t, h.Published.Up, false)
}

func TestStateClientRelayed_OnProbe_unknownAddr(t *testing.T) {
	h := NewPeerStateTestHarness()
	h.ConfigClientRelayed(t)

	h.OnProbe(controlMsg[packetProbe]{
		Packet: packetProbe{TraceID: newTraceID()},
	})

	assertType[*stateClientRelayed](t, h.State)
}

func TestStateClientRelayed_OnProbe_upgradeDirect(t *testing.T) {
	h := NewPeerStateTestHarness()
	h.ConfigClientRelayed(t)

	syn := assertType[packetSyn](t, h.Sent[0].Packet)

	ack := controlMsg[packetAck]{Packet: packetAck{TraceID: syn.TraceID}}
	ack.Packet.PossibleAddrs[0] = addrPort4(1, 1, 1, 1, 300)
	ack.Packet.PossibleAddrs[1] = addrPort4(1, 1, 1, 2, 300)

	h.State.OnAck(ack)
	h.OnPingTimer()

	probe := assertType[packetProbe](t, h.Sent[2].Packet)
	h.OnProbe(controlMsg[packetProbe]{Packet: probe})

	assertType[*stateClientDirect](t, h.State)
}