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 := &peerData{
		publish: func(rp remotePeer) {
			h.Published = rp
		},
		sendControlPacket: func(rp remotePeer, pkt marshaller) {
			h.Sent = append(h.Sent, PeerStateControlMsg{rp, pkt})
		},
		pingTimer: time.NewTicker(pingInterval),
		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) {
	h.State = h.State.OnMsg(peerUpdateMsg{p})
}

func (h *PeerStateTestHarness) OnInit(msg controlMsg[packetInit]) {
	h.State = h.State.OnMsg(msg)
}

func (h *PeerStateTestHarness) OnSyn(msg controlMsg[packetSyn]) {
	h.State = h.State.OnMsg(msg)
}

func (h *PeerStateTestHarness) OnProbe(msg controlMsg[packetProbe]) {
	h.State = h.State.OnMsg(msg)
}

func (h *PeerStateTestHarness) OnPingTimer() {
	h.State = h.State.OnMsg(pingTimerMsg{})
}

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) ConfigClientInit(t *testing.T) *stateClientInit {
	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[*stateClientInit](t, h.State)
}

func (h *PeerStateTestHarness) ConfigClientDirect(t *testing.T) *stateClient {
	h.ConfigClientInit(t)
	init := assertType[packetInit](t, h.Sent[0].Packet)
	h.OnInit(controlMsg[packetInit]{
		Packet: init,
	})

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

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

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

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

	// TODO: Fix me.

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