package node

import (
	"fmt"
	"log"
	"math/rand"
	"net/netip"
	"sync/atomic"
	"time"
	"vppn/m"
)

type peerState interface {
	Name() string
	OnPeerUpdate(*m.Peer) peerState
	OnPing(netip.AddrPort, pingPacket) peerState
	OnPong(netip.AddrPort, pongPacket) peerState
	OnPingTimer() peerState
	OnTimeoutTimer() peerState
}

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

type stateBase struct {
	// The purpose of this state machine is to manage this published data.
	published *atomic.Pointer[peerRoutingData]

	// The other remote peers.
	peers *remotePeers

	// Immutable data.
	localIP  byte
	localPub bool
	remoteIP byte
	privKey  []byte
	conn     *connWriter

	// For sending to peer.
	counter *uint64

	// Mutable peer data.
	peer        *m.Peer
	remotePub   bool
	routingData peerRoutingData // Local copy of shared data. See publish().

	// Timers
	pingTimer    *time.Timer
	timeoutTimer *time.Timer

	buf    []byte
	encBuf []byte
}

func (sb *stateBase) Name() string { return "idle" }

func (s *stateBase) OnPeerUpdate(peer *m.Peer) peerState {
	// Both nil: no change.
	if peer == nil && s.peer == nil {
		return nil
	}

	// No change.
	if peer != nil && s.peer != nil && s.peer.Version == peer.Version {
		return nil
	}

	return s.selectStateFromPeer(peer)
}

func (s *stateBase) selectStateFromPeer(peer *m.Peer) peerState {
	s.peer = peer
	s.routingData = peerRoutingData{}

	if peer == nil {
		return newStateNoPeer(s)
	}

	s.routingData.controlCipher = newControlCipher(s.privKey, peer.EncPubKey)

	ip, isValid := netip.AddrFromSlice(peer.PublicIP)
	if isValid {
		s.remotePub = true
		s.routingData.remoteAddr = netip.AddrPortFrom(ip, peer.Port)
		s.routingData.relay = peer.Mediator

		if s.localPub && s.localIP < s.remoteIP {
			return newStateServer(s)
		}
		return newStateClient(s)
	}

	if s.localPub {
		return newStateServer(s)
	}

	return newStateSelectRelay(s)
}

func (s *stateBase) OnPing(rAddr netip.AddrPort, p pingPacket) peerState { return nil }
func (s *stateBase) OnPong(rAddr netip.AddrPort, p pongPacket) peerState { return nil }
func (s *stateBase) OnPingTimer() peerState                              { return nil }
func (s *stateBase) OnTimeoutTimer() peerState                           { return nil }

// Helpers.

func (s *stateBase) resetPingTimer()    { s.pingTimer.Reset(pingInterval) }
func (s *stateBase) resetTimeoutTimer() { s.timeoutTimer.Reset(timeoutInterval) }
func (s *stateBase) stopPingTimer()     { s.pingTimer.Stop() }
func (s *stateBase) stopTimeoutTimer()  { s.timeoutTimer.Stop() }

func (s *stateBase) logf(msg string, args ...any) {
	log.Printf(fmt.Sprintf("[%03d] ", s.remoteIP)+msg, args...)
}

func (s *stateBase) publish() {
	data := s.routingData
	s.published.Store(&data)
}

func (s *stateBase) selectRelay() byte {
	possible := make([]byte, 0, 8)
	for i, peer := range s.peers {
		if peer.CanRelay() {
			possible = append(possible, byte(i))
		}
	}

	if len(possible) == 0 {
		return 0
	}
	return possible[rand.Intn(len(possible))]
}

func (s *stateBase) sendPing(sharedKey [32]byte) {
	s.sendControlPacket(newPingPacket(sharedKey))
}

func (s *stateBase) sendPong(ping pingPacket) {
	s.sendControlPacket(newPongPacket(ping.SentAt))
}

func (s *stateBase) sendControlPacket(pkt interface{ Marshal([]byte) []byte }) {
	buf := pkt.Marshal(s.buf)
	h := header{
		StreamID: controlStreamID,
		Counter:  atomic.AddUint64(s.counter, 1),
		SourceIP: s.localIP,
		DestIP:   s.remoteIP,
	}

	buf = s.routingData.controlCipher.Encrypt(h, buf, s.encBuf)
	if s.routingData.relayIP != 0 {
		s.peers[s.routingData.relayIP].RelayFor(s.remoteIP, buf)
	} else {
		s.conn.WriteTo(buf, s.routingData.remoteAddr)
	}
}

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

type stateNoPeer struct{ *stateBase }

func newStateNoPeer(b *stateBase) *stateNoPeer {
	s := &stateNoPeer{b}
	s.publish()
	return s
}

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

type stateClient struct {
	sharedKey [32]byte
	*stateBase
}

func newStateClient(b *stateBase) peerState {
	s := &stateClient{stateBase: b}
	s.publish()

	s.routingData.dataCipher = newDataCipher()
	s.sharedKey = s.routingData.dataCipher.Key()

	s.sendPing(s.sharedKey)
	s.resetPingTimer()
	s.resetTimeoutTimer()
	return s
}

func (s *stateClient) Name() string { return "client" }

func (s *stateClient) OnPong(addr netip.AddrPort, p pongPacket) peerState {
	if !s.routingData.up {
		s.routingData.up = true
		s.publish()
	}
	s.resetTimeoutTimer()
	return nil
}

func (s *stateClient) OnPingTimer() peerState {
	s.sendPing(s.sharedKey)
	s.resetPingTimer()
	return nil
}

func (s *stateClient) OnTimeoutTimer() peerState {
	s.routingData.up = false
	s.publish()
	return nil
}

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

type stateServer struct {
	*stateBase
}

func newStateServer(b *stateBase) peerState {
	s := &stateServer{b}
	s.publish()
	s.stopPingTimer()
	s.stopTimeoutTimer()
	return s
}

func (s *stateServer) Name() string { return "server" }

func (s *stateServer) OnPing(addr netip.AddrPort, p pingPacket) peerState {
	if addr != s.routingData.remoteAddr {
		s.logf("Got new peer address: %v", addr)
		s.routingData.remoteAddr = addr
		s.routingData.up = true
		s.publish()
	}

	if s.routingData.dataCipher == nil || p.SharedKey != s.routingData.dataCipher.Key() {
		s.logf("Got new shared key.")
		s.routingData.dataCipher = newDataCipherFromKey(p.SharedKey)
		s.routingData.up = true
		s.publish()
	}

	s.sendPong(p)
	return nil
}

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

type stateSelectRelay struct {
	*stateBase
}

func newStateSelectRelay(b *stateBase) peerState {
	s := &stateSelectRelay{stateBase: b}
	s.routingData.dataCipher = nil
	s.routingData.up = false
	s.publish()

	if relay := s.selectRelay(); relay != 0 {
		s.routingData.up = false
		s.routingData.relayIP = relay
		return s.selectRole()
	}

	s.resetPingTimer()
	s.stopTimeoutTimer()
	return s
}

func (s *stateSelectRelay) selectRole() peerState {
	if s.localIP < s.remoteIP {
		return newStateServerRelayed(s.stateBase)
	}
	return newStateClientRelayed(s.stateBase)
}

func (s *stateSelectRelay) Name() string { return "select-relay" }

func (s *stateSelectRelay) OnPingTimer() peerState {
	if relay := s.selectRelay(); relay != 0 {
		s.routingData.relayIP = relay
		return s.selectRole()
	}
	s.resetPingTimer()
	return nil
}

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

type stateClientRelayed struct {
	sharedKey [32]byte
	*stateBase
}

func newStateClientRelayed(b *stateBase) peerState {
	s := &stateClientRelayed{stateBase: b}

	s.routingData.dataCipher = newDataCipher()
	s.sharedKey = s.routingData.dataCipher.Key()
	s.publish()

	s.sendPing(s.sharedKey)
	s.resetPingTimer()
	s.resetTimeoutTimer()
	return s
}

func (s *stateClientRelayed) Name() string { return "client-relayed" }

func (s *stateClientRelayed) OnPong(addr netip.AddrPort, p pongPacket) peerState {
	if !s.routingData.up {
		s.routingData.up = true
		s.publish()
	}
	s.resetTimeoutTimer()
	return nil
}

func (s *stateClientRelayed) OnPingTimer() peerState {
	s.sendPing(s.sharedKey)
	s.resetPingTimer()
	return nil
}

func (s *stateClientRelayed) OnTimeoutTimer() peerState {
	return newStateSelectRelay(s.stateBase)
}

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

type stateServerRelayed struct {
	*stateBase
}

func newStateServerRelayed(b *stateBase) peerState {
	s := &stateServerRelayed{b}
	s.stopPingTimer()
	s.resetTimeoutTimer()
	return s
}

func (s *stateServerRelayed) Name() string { return "server-relayed" }

func (s *stateServerRelayed) OnPing(addr netip.AddrPort, p pingPacket) peerState {
	if s.routingData.dataCipher == nil || p.SharedKey != s.routingData.dataCipher.Key() {
		s.logf("Got new shared key.")
		s.routingData.up = true
		s.routingData.dataCipher = newDataCipherFromKey(p.SharedKey)
		s.publish()
	}

	s.sendPong(p)
	s.resetTimeoutTimer()
	return nil
}

func (s *stateServerRelayed) OnTimeoutTimer() peerState {
	return newStateSelectRelay(s.stateBase)
}