package peer

import (
	"log"
	"math/rand"
	"net/netip"
	"sync"
	"sync/atomic"
	"time"

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

type supervisor struct {
	writeToUDPAddrPort func([]byte, netip.AddrPort) (int, error)
	staged             routingTable
	shared             *atomic.Pointer[routingTable]
	peers              [256]*peerSuper
	lock               sync.Mutex

	buf1 []byte
	buf2 []byte
}

func newSupervisor(
	writeToUDPAddrPort func([]byte, netip.AddrPort) (int, error),
	rt *atomic.Pointer[routingTable],
	privKey []byte,
) *supervisor {

	routes := rt.Load()

	s := &supervisor{
		writeToUDPAddrPort: writeToUDPAddrPort,
		staged:             *routes,
		shared:             rt,
		buf1:               newBuf(),
		buf2:               newBuf(),
	}

	pubAddrs := newPubAddrStore(routes.LocalAddr)

	for i := range s.peers {
		state := &pState{
			publish:           s.publish,
			sendControlPacket: s.send,
			localIP:           routes.LocalIP,
			remoteIP:          byte(i),
			privKey:           privKey,
			localAddr:         routes.LocalAddr,
			pubAddrs:          pubAddrs,
			staged:            routes.Peers[i],
			limiter: ratelimiter.New(ratelimiter.Config{
				FillPeriod:   20 * time.Millisecond,
				MaxWaitCount: 1,
			}),
		}
		s.peers[i] = newPeerSuper(state)
	}

	return s
}

func (s *supervisor) Start() {
	for i := range s.peers {
		go s.peers[i].Run()
	}
}

func (s *supervisor) HandleControlMsg(destIP byte, msg any) {
	s.peers[destIP].HandleControlMsg(msg)
}

func (s *supervisor) send(peer remotePeer, pkt marshaller) {
	s.lock.Lock()
	defer s.lock.Unlock()

	enc := peer.EncryptControlPacket(pkt, s.buf1, s.buf2)
	if peer.Direct {
		s.writeToUDPAddrPort(enc, peer.DirectAddr)
		return
	}

	relay, ok := s.staged.GetRelay()
	if !ok {
		return
	}

	enc = relay.EncryptDataPacket(peer.IP, enc, s.buf1)
	s.writeToUDPAddrPort(enc, relay.DirectAddr)
}

func (s *supervisor) publish(rp remotePeer) {
	s.lock.Lock()
	defer s.lock.Unlock()

	s.staged.Peers[rp.IP] = rp
	s.ensureRelay()
	copy := s.staged
	s.shared.Store(&copy)
}

func (s *supervisor) ensureRelay() {
	if _, ok := s.staged.GetRelay(); ok {
		return
	}

	// TODO: Random selection?
	for _, peer := range s.staged.Peers {
		if peer.Up && peer.Direct && peer.Relay {
			s.staged.RelayIP = peer.IP
			return
		}
	}
}

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

type peerSuper struct {
	messages chan any
	state    peerState
}

func newPeerSuper(state *pState) *peerSuper {
	return &peerSuper{
		messages: make(chan any, 8),
		state:    state.OnPeerUpdate(nil),
	}
}

func (s *peerSuper) HandleControlMsg(msg any) {
	select {
	case s.messages <- msg:
	default:
	}
}

func (s *peerSuper) Run() {
	go func() {
		// Randomize ping timers.
		time.Sleep(time.Duration(rand.Intn(4000)) * time.Millisecond)
		for range time.Tick(pingInterval) {
			s.messages <- pingTimerMsg{}
		}
	}()

	for rawMsg := range s.messages {
		switch msg := rawMsg.(type) {

		case peerUpdateMsg:
			s.state = s.state.OnPeerUpdate(msg.Peer)

		case controlMsg[packetSyn]:
			s.state = s.state.OnSyn(msg)

		case controlMsg[packetAck]:
			s.state.OnAck(msg)

		case controlMsg[packetProbe]:
			s.state = s.state.OnProbe(msg)

		case controlMsg[packetLocalDiscovery]:
			s.state.OnLocalDiscovery(msg)

		case pingTimerMsg:
			s.state = s.state.OnPingTimer()

		default:
			log.Printf("WARNING: unknown message type: %+v", msg)
		}
	}
}